memory_export.go

package yaad

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sort"
	"time"
)

// currentExportVersion is stamped into every export so consumers can detect
// incompatible format changes in the future.
const currentExportVersion = "1.0"

// TierStats holds the aggregate entry count and token usage for one tier.
type TierStats struct {
	Count  int `json:"count"`
	Tokens int `json:"tokens"`
}

// TierStatsMap maps each TierType to its aggregate statistics.
type TierStatsMap map[TierType]TierStats

// MemoryExport is a serializable snapshot of a SpatialMemory's entries and
// tier statistics. It is produced by ExportMemory and consumed by
// ImportMemory to back up and restore spatial memory state.
type MemoryExport struct {
	Version    string        `json:"version"`
	ExportedAt time.Time     `json:"exported_at"`
	Entries    []MemoryEntry `json:"entries"`
	Stats      TierStatsMap  `json:"stats"`
}

// ExportConfig controls what ExportMemory includes in a MemoryExport and in
// which on-wire format the export is intended to be written.
type ExportConfig struct {
	// IncludeColdTier includes cold-tier (rarely accessed) entries when true.
	// When false, cold entries are omitted, reducing export size.
	IncludeColdTier bool

	// MaxEntries caps the number of entries in the export. Entries are sorted
	// by LastAccessed descending before truncation, so the most recently
	// accessed entries are retained. Zero or negative means unlimited.
	MaxEntries int

	// Format declares the intended serialization format: "json" (default) or
	// "gzipped". ExportMemory validates this value; the actual compression is
	// the caller's responsibility (e.g. wrapping the writer in gzip.NewWriter).
	Format string
}

// ExportMemory captures the current state of sm into a MemoryExport, shaped
// by config. The returned Stats describe only the entries that made it into
// the export (after tier filtering and MaxEntries truncation).
func ExportMemory(sm *SpatialMemory, config ExportConfig) (*MemoryExport, error) {
	switch config.Format {
	case "", "json", "gzipped":
		// accepted
	default:
		return nil, fmt.Errorf("unsupported export format: %q", config.Format)
	}

	// Snapshot entries under a read lock; copy each entry to decouple from
	// sm's live pointers before doing any sorting or truncation.
	sm.mu.RLock()
	snapshot := make([]MemoryEntry, 0, len(sm.entries))
	for _, e := range sm.entries {
		if !config.IncludeColdTier && e.Tier == TierCold {
			continue
		}
		snapshot = append(snapshot, *e)
	}
	sm.mu.RUnlock()

	// Keep the most recently accessed entries when truncating.
	sort.Slice(snapshot, func(i, j int) bool {
		return snapshot[i].LastAccessed.After(snapshot[j].LastAccessed)
	})
	if config.MaxEntries > 0 && len(snapshot) > config.MaxEntries {
		snapshot = snapshot[:config.MaxEntries]
	}

	stats := make(TierStatsMap)
	for i := range snapshot {
		s := stats[snapshot[i].Tier]
		s.Count++
		s.Tokens += snapshot[i].TokenCount
		stats[snapshot[i].Tier] = s
	}

	return &MemoryExport{
		Version:    currentExportVersion,
		ExportedAt: time.Now().UTC(),
		Entries:    snapshot,
		Stats:      stats,
	}, nil
}

// ImportMemory copies the entries from export into sm, overwriting any
// existing entries that share the same ID. Each imported entry retains its
// original tier, access count, and timestamps so the restored state matches
// the exported state. Entries with an empty ID or an unrecognized tier are
// skipped. It returns the number of entries successfully imported.
func ImportMemory(sm *SpatialMemory, export *MemoryExport) (int, error) {
	if export == nil {
		return 0, errors.New("nil export")
	}

	sm.mu.Lock()
	defer sm.mu.Unlock()

	imported := 0
	for i := range export.Entries {
		e := export.Entries[i] // copy to avoid aliasing the export's slice
		if e.ID == "" {
			continue
		}
		switch e.Tier {
		case TierHot, TierWarm, TierCold:
			// valid
		default:
			continue
		}

		ptr := sm.entries[e.ID]
		if ptr == nil {
			ptr = &MemoryEntry{}
			sm.entries[e.ID] = ptr
		}
		*ptr = e
		imported++
	}
	return imported, nil
}

// WriteTo serializes the memory export as indented JSON and writes it to w.
// It implements io.WriterTo.
func (me *MemoryExport) WriteTo(w io.Writer) (int64, error) {
	if me == nil {
		return 0, errors.New("nil memory export")
	}
	cw := &countingWriter{w: w}
	enc := json.NewEncoder(cw)
	enc.SetIndent("", "  ")
	if err := enc.Encode(me); err != nil {
		return cw.n, fmt.Errorf("writing memory export: %w", err)
	}
	return cw.n, nil
}

// countingWriter wraps an io.Writer and counts bytes written through it.
type countingWriter struct {
	w io.Writer
	n int64
}

func (cw *countingWriter) Write(p []byte) (int, error) {
	n, err := cw.w.Write(p)
	cw.n += int64(n)
	return n, err
}

// ReadMemoryExport decodes a MemoryExport from the JSON format produced by
// WriteTo.
func ReadMemoryExport(r io.Reader) (*MemoryExport, error) {
	var me MemoryExport
	if err := json.NewDecoder(r).Decode(&me); err != nil {
		return nil, fmt.Errorf("reading memory export: %w", err)
	}
	return &me, nil
}

// Summary returns a one-line human-readable description of the export,
// including its version, timestamp, and per-tier entry breakdown.
func (me *MemoryExport) Summary() string {
	if me == nil || len(me.Entries) == 0 {
		return "empty memory export"
	}
	return fmt.Sprintf(
		"memory export v%s (%s): %d entries (hot=%d warm=%d cold=%d)",
		me.Version,
		me.ExportedAt.Format(time.RFC3339),
		len(me.Entries),
		me.Stats[TierHot].Count,
		me.Stats[TierWarm].Count,
		me.Stats[TierCold].Count,
	)
}