package store

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	"neuroforge/internal/core"
)

const maxSegmentRecordBytes = 128 << 20

type segmentRecord struct {
	Revision uint64       `json:"revision"`
	Op       string       `json:"op"`
	ID       string       `json:"id"`
	Memory   *core.Memory `json:"memory,omitempty"`
}

type segmentLocation struct {
	Path     string
	Offset   int64
	Length   uint32
	Revision uint64
	Deleted  bool
}

type SegmentStats struct {
	Enabled      bool   `json:"enabled"`
	Segments     int    `json:"segments"`
	Records      int    `json:"records"`
	Live         int    `json:"live"`
	Tombstones   int    `json:"tombstones"`
	Bytes        int64  `json:"bytes"`
	Active       string `json:"active,omitempty"`
	MmapSegments int    `json:"mmap_segments"`
}

type SegmentStore struct {
	mu              sync.RWMutex
	dir             string
	maxSegmentBytes int64
	mmapSealed      bool
	seq             int
	activePath      string
	activeSize      int64
	index           map[string]segmentLocation
	mmaps           map[string][]byte
	records         int
	tombstones      int
	scanMetadata    map[string]core.Memory
}

func openSegmentStore(dir string, maxBytes int64, mmapSealed bool) (*SegmentStore, error) {
	if maxBytes < 1<<20 {
		maxBytes = 128 << 20
	}
	ss := &SegmentStore{
		dir: dir, maxSegmentBytes: maxBytes, mmapSealed: mmapSealed,
		index: map[string]segmentLocation{}, mmaps: map[string][]byte{}, scanMetadata: map[string]core.Memory{},
	}
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, err
	}
	if err := ss.scan(); err != nil {
		ss.Close()
		return nil, err
	}
	return ss, nil
}

func (s *SegmentStore) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for path, b := range s.mmaps {
		_ = unmapSegmentFile(b)
		delete(s.mmaps, path)
	}
}

func segmentName(seq int) string { return fmt.Sprintf("segment-%06d.nfs", seq) }

func parseSegmentSeq(name string) (int, bool) {
	if !strings.HasPrefix(name, "segment-") || !strings.HasSuffix(name, ".nfs") {
		return 0, false
	}
	x := strings.TrimSuffix(strings.TrimPrefix(name, "segment-"), ".nfs")
	n, err := strconv.Atoi(x)
	return n, err == nil && n > 0
}

func (s *SegmentStore) scan() error {
	ents, err := os.ReadDir(s.dir)
	if err != nil {
		return err
	}
	type item struct {
		seq  int
		path string
	}
	items := []item{}
	for _, ent := range ents {
		if ent.IsDir() {
			continue
		}
		seq, ok := parseSegmentSeq(ent.Name())
		if ok {
			items = append(items, item{seq: seq, path: filepath.Join(s.dir, ent.Name())})
		}
	}
	sort.Slice(items, func(i, j int) bool { return items[i].seq < items[j].seq })
	for _, it := range items {
		if err := s.scanFile(it.path); err != nil {
			return err
		}
		s.seq = it.seq
	}
	if s.seq == 0 {
		s.seq = 1
	}
	s.activePath = filepath.Join(s.dir, segmentName(s.seq))
	if st, err := os.Stat(s.activePath); err == nil {
		s.activeSize = st.Size()
	}
	if s.activeSize >= s.maxSegmentBytes {
		s.seq++
		s.activePath = filepath.Join(s.dir, segmentName(s.seq))
		s.activeSize = 0
	}
	return nil
}

func (s *SegmentStore) scanFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	var offset int64
	for {
		var hdr [4]byte
		_, err := io.ReadFull(f, hdr[:])
		if errors.Is(err, io.EOF) {
			return nil
		}
		if errors.Is(err, io.ErrUnexpectedEOF) {
			// A crash can leave a partial tail. Ignore only that tail.
			return nil
		}
		if err != nil {
			return err
		}
		n := binary.BigEndian.Uint32(hdr[:])
		if n == 0 || n > maxSegmentRecordBytes {
			return fmt.Errorf("invalid segment record length %d at %s:%d", n, path, offset)
		}
		buf := make([]byte, n)
		if _, err := io.ReadFull(f, buf); err != nil {
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
				return nil
			}
			return err
		}
		var rec segmentRecord
		if err := json.Unmarshal(buf, &rec); err != nil {
			return fmt.Errorf("decode segment %s:%d: %w", path, offset, err)
		}
		if rec.ID == "" {
			return fmt.Errorf("empty memory id in segment %s:%d", path, offset)
		}
		loc := segmentLocation{Path: path, Offset: offset + 4, Length: n, Revision: rec.Revision, Deleted: rec.Op == "delete"}
		if old, ok := s.index[rec.ID]; !ok || rec.Revision >= old.Revision {
			s.index[rec.ID] = loc
			if rec.Op == "delete" {
				delete(s.scanMetadata, rec.ID)
			} else if rec.Memory != nil {
				m := cloneMemory(*rec.Memory)
				if m.VectorDim == 0 && len(m.Vector) > 0 {
					m.VectorDim = len(m.Vector)
				}
				m.Text = ""
				m.Vector = nil
				s.scanMetadata[rec.ID] = m
			}
		}
		s.records++
		if rec.Op == "delete" {
			s.tombstones++
		}
		offset += 4 + int64(n)
	}
}

func (s *SegmentStore) appendRecord(rec segmentRecord) error {
	return s.appendRecords([]segmentRecord{rec})
}

func (s *SegmentStore) appendRecords(records []segmentRecord) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	var f *os.File
	var openPath string
	closeFile := func() error {
		if f == nil {
			return nil
		}
		err := f.Sync()
		cerr := f.Close()
		f = nil
		openPath = ""
		if err != nil {
			return err
		}
		return cerr
	}
	defer func() { _ = closeFile() }()
	for _, rec := range records {
		if rec.ID == "" {
			return errors.New("segment record id required")
		}
		deleted := rec.Op == "delete"
		if old, ok := s.index[rec.ID]; ok && old.Revision == rec.Revision && old.Deleted == deleted {
			continue
		}
		payload, err := json.Marshal(rec)
		if err != nil {
			return err
		}
		if len(payload) > maxSegmentRecordBytes {
			return fmt.Errorf("memory segment record exceeds %d bytes", maxSegmentRecordBytes)
		}
		recordBytes := int64(4 + len(payload))
		if s.activeSize > 0 && s.activeSize+recordBytes > s.maxSegmentBytes {
			if err := closeFile(); err != nil {
				return err
			}
			if s.mmapSealed {
				_, _ = s.mapLocked(s.activePath)
			}
			s.seq++
			s.activePath = filepath.Join(s.dir, segmentName(s.seq))
			s.activeSize = 0
		}
		if f == nil || openPath != s.activePath {
			if err := closeFile(); err != nil {
				return err
			}
			f, err = os.OpenFile(s.activePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
			if err != nil {
				return err
			}
			openPath = s.activePath
		}
		var hdr [4]byte
		binary.BigEndian.PutUint32(hdr[:], uint32(len(payload)))
		start := s.activeSize
		if _, err := f.Write(hdr[:]); err != nil {
			return err
		}
		if _, err := f.Write(payload); err != nil {
			return err
		}
		s.activeSize += recordBytes
		s.index[rec.ID] = segmentLocation{Path: s.activePath, Offset: start + 4, Length: uint32(len(payload)), Revision: rec.Revision, Deleted: rec.Op == "delete"}
		s.records++
		if rec.Op == "delete" {
			s.tombstones++
		}
	}
	return closeFile()
}

func (s *SegmentStore) AppendUpsert(revision uint64, memories []core.Memory) error {
	recs := make([]segmentRecord, 0, len(memories))
	for i := range memories {
		m := cloneMemory(memories[i])
		recs = append(recs, segmentRecord{Revision: revision, Op: "upsert", ID: m.ID, Memory: &m})
	}
	return s.appendRecords(recs)
}

func (s *SegmentStore) AppendDelete(revision uint64, ids []string) error {
	recs := make([]segmentRecord, 0, len(ids))
	for _, id := range ids {
		recs = append(recs, segmentRecord{Revision: revision, Op: "delete", ID: id})
	}
	return s.appendRecords(recs)
}

func (s *SegmentStore) mapLocked(path string) ([]byte, bool) {
	if !s.mmapSealed || path == "" || path == s.activePath {
		return nil, false
	}
	if b, ok := s.mmaps[path]; ok {
		return b, true
	}
	b, ok, err := mapSegmentFile(path)
	if err != nil || !ok {
		return nil, false
	}
	s.mmaps[path] = b
	return b, true
}

func (s *SegmentStore) readLocation(loc segmentLocation) (segmentRecord, error) {
	s.mu.Lock()
	if b, ok := s.mapLocked(loc.Path); ok {
		start, end := loc.Offset, loc.Offset+int64(loc.Length)
		if start >= 0 && end <= int64(len(b)) {
			buf := append([]byte(nil), b[start:end]...)
			s.mu.Unlock()
			var rec segmentRecord
			return rec, json.Unmarshal(buf, &rec)
		}
	}
	s.mu.Unlock()
	f, err := os.Open(loc.Path)
	if err != nil {
		return segmentRecord{}, err
	}
	defer f.Close()
	buf := make([]byte, loc.Length)
	if _, err := f.ReadAt(buf, loc.Offset); err != nil {
		return segmentRecord{}, err
	}
	var rec segmentRecord
	return rec, json.Unmarshal(buf, &rec)
}

func (s *SegmentStore) Get(id string) (core.Memory, bool, bool, error) {
	s.mu.RLock()
	loc, ok := s.index[id]
	s.mu.RUnlock()
	if !ok {
		return core.Memory{}, false, false, nil
	}
	if loc.Deleted {
		return core.Memory{}, true, true, nil
	}
	rec, err := s.readLocation(loc)
	if err != nil {
		return core.Memory{}, true, false, err
	}
	if rec.Memory == nil {
		return core.Memory{}, true, false, errors.New("segment upsert has no memory body")
	}
	return cloneMemory(*rec.Memory), true, false, nil
}

type liveSegmentCursor struct {
	id  string
	loc segmentLocation
}

// iterateLivePayloadsSequential captures a point-in-time view of the latest
// live records and then walks each segment strictly forward. It never performs
// one ReadAt/seek per memory: bytes between live records are consumed through a
// buffered stream and the payload buffer is reused. Because segment files are
// append-only, captured locations remain valid while newer writes continue.
func (s *SegmentStore) iterateLivePayloadsSequential(fn func(id string, payload []byte) error) error {
	if fn == nil {
		return errors.New("segment payload iterator callback required")
	}

	s.mu.RLock()
	byPath := make(map[string][]liveSegmentCursor)
	for id, loc := range s.index {
		if loc.Deleted {
			continue
		}
		byPath[loc.Path] = append(byPath[loc.Path], liveSegmentCursor{id: id, loc: loc})
	}
	s.mu.RUnlock()

	paths := make([]string, 0, len(byPath))
	for path := range byPath {
		paths = append(paths, path)
	}
	sort.Strings(paths)

	var payload []byte
	for _, path := range paths {
		cursors := byPath[path]
		sort.Slice(cursors, func(i, j int) bool { return cursors[i].loc.Offset < cursors[j].loc.Offset })

		f, err := os.Open(path)
		if err != nil {
			return err
		}
		br := bufio.NewReaderSize(f, 1<<20)
		var pos int64
		for _, cursor := range cursors {
			loc := cursor.loc
			if loc.Offset < pos {
				_ = f.Close()
				return fmt.Errorf("segment iterator moved backwards in %s: current=%d target=%d", path, pos, loc.Offset)
			}
			if gap := loc.Offset - pos; gap > 0 {
				if _, err := io.CopyN(io.Discard, br, gap); err != nil {
					_ = f.Close()
					return fmt.Errorf("scan segment %s to offset %d: %w", path, loc.Offset, err)
				}
				pos += gap
			}

			n := int(loc.Length)
			if cap(payload) < n {
				payload = make([]byte, n)
			} else {
				payload = payload[:n]
			}
			if _, err := io.ReadFull(br, payload); err != nil {
				_ = f.Close()
				return fmt.Errorf("read segment %s:%d: %w", path, loc.Offset, err)
			}
			pos += int64(n)
			if err := fn(cursor.id, payload); err != nil {
				_ = f.Close()
				return err
			}
		}
		if err := f.Close(); err != nil {
			return err
		}
	}
	return nil
}

// IterateLiveMemories walks the latest live records using the sequential
// segment scanner. The callback receives a detached Memory and may retain it.
func (s *SegmentStore) IterateLiveMemories(fn func(core.Memory) error) error {
	if fn == nil {
		return errors.New("segment iterator callback required")
	}
	return s.iterateLivePayloadsSequential(func(expectedID string, payload []byte) error {
		var rec segmentRecord
		if err := json.Unmarshal(payload, &rec); err != nil {
			return fmt.Errorf("decode live segment memory %s: %w", expectedID, err)
		}
		if rec.ID != expectedID {
			return fmt.Errorf("segment index mismatch: expected %q, payload contains %q", expectedID, rec.ID)
		}
		if rec.Op == "delete" || rec.Memory == nil {
			return nil
		}
		return fn(cloneMemory(*rec.Memory))
	})
}

type segmentVectorMemory struct {
	Status    string    `json:"status"`
	Vector    []float32 `json:"vector"`
	VectorDim int       `json:"vector_dim,omitempty"`
}

type segmentVectorRecord struct {
	Op     string               `json:"op"`
	ID     string               `json:"id"`
	Memory *segmentVectorMemory `json:"memory,omitempty"`
}

// IterateLiveVectorsSequential is the disk-ANN migration/build fast path. It
// scans every segment at most once in ascending physical offset order and only
// decodes the fields required by the vector builder. The vector slice is owned
// by the iterator and must not be retained after the callback returns.
func (s *SegmentStore) IterateLiveVectorsSequential(dim int, fn func(id string, vector []float32) error) error {
	if dim < 1 {
		return errors.New("vector dimension must be positive")
	}
	if fn == nil {
		return errors.New("segment vector iterator callback required")
	}
	return s.iterateLivePayloadsSequential(func(expectedID string, payload []byte) error {
		var rec segmentVectorRecord
		if err := json.Unmarshal(payload, &rec); err != nil {
			return fmt.Errorf("decode live segment vector %s: %w", expectedID, err)
		}
		if rec.ID != expectedID {
			return fmt.Errorf("segment index mismatch: expected %q, payload contains %q", expectedID, rec.ID)
		}
		if rec.Op == "delete" || rec.Memory == nil {
			return nil
		}
		status := rec.Memory.Status
		if status != "" && status != core.MemoryActive && status != core.MemoryConflicted {
			return nil
		}
		actualDim := rec.Memory.VectorDim
		if actualDim == 0 {
			actualDim = len(rec.Memory.Vector)
		}
		if actualDim != dim || len(rec.Memory.Vector) != dim {
			return nil
		}
		return fn(rec.ID, rec.Memory.Vector)
	})
}

func (s *SegmentStore) Hydrate(memories map[string]*core.Memory) error {
	for id := range memories {
		m, found, deleted, err := s.Get(id)
		if err != nil {
			return err
		}
		if !found {
			// v0.3 migration: the body can still be present in state.json.
			continue
		}
		if deleted {
			delete(memories, id)
			continue
		}
		memories[id] = &m
	}
	return nil
}

func (s *SegmentStore) ConsumeMetadata() map[string]*core.Memory {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make(map[string]*core.Memory, len(s.scanMetadata))
	for id, m := range s.scanMetadata {
		cp := cloneMemory(m)
		out[id] = &cp
	}
	s.scanMetadata = nil
	return out
}

func (s *SegmentStore) HasRecords() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.records > 0
}

func (s *SegmentStore) Rebuild(memories map[string]*core.Memory, revision uint64) error {
	s.mu.Lock()
	for _, b := range s.mmaps {
		_ = unmapSegmentFile(b)
	}
	s.mmaps = map[string][]byte{}
	s.mu.Unlock()
	tmp := s.dir + ".rebuild"
	_ = os.RemoveAll(tmp)
	if err := os.MkdirAll(tmp, 0700); err != nil {
		return err
	}
	fresh, err := openSegmentStore(tmp, s.maxSegmentBytes, s.mmapSealed)
	if err != nil {
		return err
	}
	ids := make([]string, 0, len(memories))
	for id := range memories {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	for _, id := range ids {
		if err := fresh.AppendUpsert(revision, []core.Memory{cloneMemory(*memories[id])}); err != nil {
			fresh.Close()
			return err
		}
	}
	fresh.Close()
	backup := s.dir + ".old"
	_ = os.RemoveAll(backup)
	if err := os.Rename(s.dir, backup); err != nil {
		return err
	}
	if err := os.Rename(tmp, s.dir); err != nil {
		_ = os.Rename(backup, s.dir)
		return err
	}
	_ = os.RemoveAll(backup)

	reloaded, err := openSegmentStore(s.dir, s.maxSegmentBytes, s.mmapSealed)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.seq = reloaded.seq
	s.activePath = reloaded.activePath
	s.activeSize = reloaded.activeSize
	s.index = reloaded.index
	s.records = reloaded.records
	s.tombstones = reloaded.tombstones
	s.mmaps = reloaded.mmaps
	s.scanMetadata = nil
	s.mu.Unlock()
	// Ownership of mmaps moved to s.
	reloaded.mmaps = map[string][]byte{}
	return nil
}

func (s *SegmentStore) Stats() SegmentStats {
	s.mu.RLock()
	defer s.mu.RUnlock()
	live := 0
	for _, loc := range s.index {
		if !loc.Deleted {
			live++
		}
	}
	stats := SegmentStats{Enabled: true, Records: s.records, Live: live, Tombstones: s.tombstones, Active: filepath.Base(s.activePath), MmapSegments: len(s.mmaps)}
	ents, _ := os.ReadDir(s.dir)
	for _, ent := range ents {
		if ent.IsDir() {
			continue
		}
		if _, ok := parseSegmentSeq(ent.Name()); !ok {
			continue
		}
		stats.Segments++
		if info, err := ent.Info(); err == nil {
			stats.Bytes += info.Size()
		}
	}
	return stats
}

func (s *SegmentStore) TombstoneRatio() float64 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.records == 0 {
		return 0
	}
	return float64(s.tombstones) / float64(s.records)
}

func (s *SegmentStore) HasLive(id string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	loc, ok := s.index[id]
	return ok && !loc.Deleted
}