package memorysearch

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"unicode/utf8"
)

// Scan bounds. Admitting beyond these caps would let the derivative grow
// unbounded against pathological Memory clones; the scan reports the cap as a
// degraded/limited status instead of silently truncating.
const (
	MaxFileBytes  int64 = 1 << 20   // 1 MiB per file
	MaxFiles            = 20 * 1000 // per agent snapshot
	MaxTotalBytes int64 = 100 << 20 // per agent snapshot
)

// allowedExtensions maps a lowercase file extension to its stable source
// Format. The scanner records Format separately from Kind so that the
// path-domain Kind used by filters and result grouping stays independent of
// how the chunker parses the bytes.
var allowedExtensions = map[string]string{
	".md":       SourceFormatMarkdown,
	".markdown": SourceFormatMarkdown,
	".txt":      SourceFormatText,
	".json":     SourceFormatJSON,
	".yaml":     SourceFormatYAML,
	".yml":      SourceFormatYAML,
}

// retentionMarkerName is the clone-side retention marker file. It is owned by
// the daemon materialization path, not Memory content, so it is excluded.
const retentionMarkerName = ".agentteam-memory-retention.json"

// ScanStatus records the bounded outcome of a scan.
type ScanStatus struct {
	Admitted   int
	Skipped    int
	TotalBytes int64
	Limited    bool // true if MaxFiles or MaxTotalBytes forced a cutoff
	Degraded   bool // true if a scan error was encountered for a file or subtree
}

// Scan walks the resolved source root and produces a deterministic SourceSnapshot.
//
// root must be an absolute, cleaned directory path. The walk is confined to
// the opened root via os.OpenRoot: regular-file reads happen through
// root.Open(relPath), so a path that is swapped to an outside-pointing symlink
// between Lstat and Open is rejected rather than followed out of the root.
// The returned SourceSnapshot is stable for the same on-disk tree regardless
// of traversal timing.
func Scan(ctx context.Context, root string) (SourceSnapshot, ScanStatus, error) {
	return scanWithConfig(ctx, root, scanConfig{
		maxFiles:      MaxFiles,
		maxTotalBytes: MaxTotalBytes,
		maxFileBytes:  MaxFileBytes,
	})
}

type scanConfig struct {
	maxFiles      int
	maxTotalBytes int64
	maxFileBytes  int64
	// beforeHashFile fires after WalkDir has admitted a file but before the
	// confined Open+read. Tests use it to deterministically swap or grow the
	// file at the seam, exercising TOCTOU defenses. It receives the
	// source-relative slash-normalized path.
	beforeHashFile func(relPath string)
}

func scanWithConfig(ctx context.Context, root string, cfg scanConfig) (SourceSnapshot, ScanStatus, error) {
	if !filepath.IsAbs(root) {
		return SourceSnapshot{}, ScanStatus{}, ErrInvalidRoot
	}
	cleaned := filepath.Clean(root)
	if cleaned == string(filepath.Separator) || hasNUL(cleaned) {
		return SourceSnapshot{}, ScanStatus{}, ErrInvalidRoot
	}
	info, err := os.Lstat(cleaned)
	if err != nil {
		return SourceSnapshot{}, ScanStatus{}, err
	}
	if !info.IsDir() {
		return SourceSnapshot{}, ScanStatus{}, ErrInvalidRoot
	}
	rootFile, err := os.OpenRoot(cleaned)
	if err != nil {
		return SourceSnapshot{}, ScanStatus{}, err
	}
	defer rootFile.Close()
	rootStat, err := rootFile.Stat(".")
	if err != nil {
		return SourceSnapshot{}, ScanStatus{}, err
	}

	status := ScanStatus{}
	snap := SourceSnapshot{Files: map[string]SourceFile{}}

	prefix := cleaned + string(filepath.Separator)
	walkErr := filepath.WalkDir(cleaned, func(path string, d fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			status.Degraded = true
			return nil
		}
		if err := ctx.Err(); err != nil {
			return err
		}
		if d.IsDir() {
			if path != cleaned && shouldSkipDir(filepath.Base(path)) {
				return filepath.SkipDir
			}
			return nil
		}
		if !d.Type().IsRegular() {
			// Reject symlinks, devices, pipes, sockets, etc. at walk time.
			return nil
		}
		rel := strings.TrimPrefix(path, prefix)
		normalized := filepath.ToSlash(rel)
		if !isIncludedPath(normalized) {
			return nil
		}
		ext := strings.ToLower(filepath.Ext(normalized))
		format, ok := allowedExtensions[ext]
		if !ok {
			return nil
		}
		fi, err := d.Info()
		if err != nil {
			status.Degraded = true
			return nil
		}
		if fi.Size() > cfg.maxFileBytes {
			status.Skipped++
			return nil
		}
		if status.Admitted >= cfg.maxFiles {
			status.Limited = true
			return errBoundsCutoff
		}
		if status.TotalBytes+fi.Size() > cfg.maxTotalBytes {
			status.Limited = true
			return errBoundsCutoff
		}
		walkStat, err := os.Lstat(path)
		if err != nil {
			status.Degraded = true
			return nil
		}
		if cfg.beforeHashFile != nil {
			cfg.beforeHashFile(normalized)
		}
		data, hash, hashErr := readAndHash(rootFile, rel, walkStat, cfg.maxFileBytes)
		if hashErr != nil {
			if errors.Is(hashErr, ErrFileTooLarge) || errors.Is(hashErr, ErrBinaryContent) || errors.Is(hashErr, ErrInvalidUTF8) {
				status.Skipped++
			} else {
				status.Degraded = true
			}
			return nil
		}
		size := int64(len(data))
		// Re-check the total-byte cap against the actual bytes returned by
		// the bounded read. The pre-read check used fi.Size() from the
		// walk-time DirEntry, but the file may have grown between that stat
		// and the open+read seam. readAndHash only guarantees preStat ==
		// postStat (no mutation during the read itself), not preStat ==
		// fi.Size(). Use the real len(data) here so a growth-via-read race
		// cannot push TotalBytes past the cap without being flagged.
		if status.TotalBytes+size > cfg.maxTotalBytes {
			status.Limited = true
			status.Skipped++
			return nil
		}
		status.Admitted++
		status.TotalBytes += size
		snap.Files[normalized] = SourceFile{
			Path:        normalized,
			Size:        size,
			MTimeNS:     fi.ModTime().UnixNano(),
			ContentHash: hash,
			Kind:        kindFromPath(normalized),
			Format:      format,
			Archived:    isArchivedPath(normalized),
		}
		return nil
	})
	if walkErr != nil && !errors.Is(walkErr, errBoundsCutoff) {
		return SourceSnapshot{}, status, walkErr
	}
	// Defensive: the open fd must still be rooted at the same directory inode
	// we started with. A change here means the root itself was swapped.
	endStat, err := rootFile.Stat(".")
	if err == nil && !os.SameFile(rootStat, endStat) {
		status.Degraded = true
	}
	snap.Revision = computeRevision(snap.Files)
	return snap, status, nil
}

var (
	// ErrInvalidRoot signals that the supplied source root is not an
	// absolute directory path that can be scanned.
	ErrInvalidRoot = errors.New("memory search: invalid source root")
	// ErrBinaryContent signals that a candidate file contained NUL bytes and
	// was rejected from the snapshot.
	ErrBinaryContent = errors.New("memory search: binary content rejected")
	// ErrInvalidUTF8 signals that a candidate file contained invalid UTF-8
	// and was rejected from the snapshot.
	ErrInvalidUTF8 = errors.New("memory search: invalid utf-8 rejected")
	// ErrFileTooLarge signals that a file grew past MaxFileBytes between the
	// pre-open size check and the bounded read.
	ErrFileTooLarge = errors.New("memory search: file exceeds size bound")
	// ErrRootEscape signals that an open attempt would have escaped the
	// confined source root (e.g. a symlink swap to an outside target).
	ErrRootEscape = errors.New("memory search: open escaped source root")

	errBoundsCutoff = errors.New("memory search: skip remainder for bounds")
)

// SnapshotDiff describes the change set between two snapshots. Rename is
// represented as a delete plus an add because the derivative keys chunks by
// stable path-derived IDs.
type SnapshotDiff struct {
	Added    []SourceFile
	Modified []SourceFile
	Removed  []string
}

// Diff returns the change set from old to cur. Both snapshots must come from
// the same source root layout; cross-root diffs are meaningless.
func Diff(old, cur SourceSnapshot) SnapshotDiff {
	out := SnapshotDiff{}
	for path, of := range old.Files {
		cf, ok := cur.Files[path]
		if !ok {
			out.Removed = append(out.Removed, path)
			continue
		}
		if cf.ContentHash != of.ContentHash {
			out.Modified = append(out.Modified, cf)
		}
	}
	for path, cf := range cur.Files {
		if _, ok := old.Files[path]; !ok {
			out.Added = append(out.Added, cf)
		}
	}
	sort.Strings(out.Removed)
	sort.Slice(out.Added, func(i, j int) bool { return out.Added[i].Path < out.Added[j].Path })
	sort.Slice(out.Modified, func(i, j int) bool { return out.Modified[i].Path < out.Modified[j].Path })
	return out
}

func computeRevision(files map[string]SourceFile) string {
	paths := make([]string, 0, len(files))
	for p := range files {
		paths = append(paths, p)
	}
	sort.Strings(paths)
	h := sha256.New()
	for _, p := range paths {
		h.Write([]byte(p))
		h.Write([]byte{0})
		h.Write([]byte(files[p].ContentHash))
		h.Write([]byte{0})
	}
	return hex.EncodeToString(h.Sum(nil))
}

// isIncludedPath applies the source filter: exclude dot
// directories, the `.git` and `.derived` subtrees (anywhere in the path),
// the daemon-owned retention marker, and temp/swap/backup file names.
func isIncludedPath(normalized string) bool {
	if normalized == "" {
		return false
	}
	if normalized == retentionMarkerName {
		return false
	}
	segments := strings.Split(normalized, "/")
	for _, seg := range segments {
		if seg == "" {
			continue
		}
		if seg == ".git" || seg == DerivedDirName {
			return false
		}
		if len(seg) > 1 && seg[0] == '.' {
			return false
		}
		if isTempName(strings.ToLower(seg)) {
			return false
		}
	}
	return true
}

func shouldSkipDir(base string) bool {
	if base == ".git" || base == DerivedDirName {
		return true
	}
	if len(base) > 1 && base[0] == '.' {
		return true
	}
	return false
}

func isTempName(lowered string) bool {
	ext := ""
	if idx := strings.LastIndexByte(lowered, '.'); idx >= 0 {
		ext = lowered[idx:]
	}
	base := strings.TrimSuffix(lowered, ext)
	switch {
	case strings.HasSuffix(base, "~"):
		return true
	case strings.HasPrefix(base, "~$"):
		return true
	}
	switch ext {
	case ".tmp", ".temp", ".swp", ".swo", ".bak", ".old":
		return true
	}
	return false
}

func isArchivedPath(normalized string) bool {
	segments := strings.Split(normalized, "/")
	for _, seg := range segments {
		switch seg {
		case "archive", "archived":
			return true
		}
	}
	base := strings.ToLower(filepath.Base(normalized))
	stem := strings.TrimSuffix(base, filepath.Ext(base))
	if strings.HasPrefix(stem, "archived-") || strings.Contains(stem, ".archived.") {
		return true
	}
	return false
}

// kindFromPath derives the path-domain Kind for a source-relative path.
// The mapping mirrors the Memory starter template layout:
//
//   - top-level files (no slash) → root
//   - system/...                 → system
//   - profiles/..., profile/...  → profile
//   - knowledge/concepts/...     → concept
//   - knowledge/workflows/...    → workflow
//   - events/..., event/...      → event
//   - reflections/..., reflection/... → reflection
//   - anything else              → other
//
// Singular variants are tolerated but the current template uses the plural
// form. Kind is independent of SourceFile.Format.
func kindFromPath(normalized string) string {
	if normalized == "" {
		return SourceKindOther
	}
	if !strings.Contains(normalized, "/") {
		return SourceKindRoot
	}
	rest := normalized
	seg := rest
	if idx := strings.Index(rest, "/"); idx >= 0 {
		seg = rest[:idx]
		rest = rest[idx+1:]
	} else {
		rest = ""
	}
	switch strings.ToLower(seg) {
	case SourceKindSystem:
		return SourceKindSystem
	case "profiles", "profile":
		return SourceKindProfile
	case "events", "event":
		return SourceKindEvent
	case "reflections", "reflection":
		return SourceKindReflection
	case "knowledge":
		category := rest
		if idx := strings.Index(rest, "/"); idx >= 0 {
			category = rest[:idx]
		}
		switch strings.ToLower(category) {
		case "concepts", "concept":
			return SourceKindConcept
		case "workflows", "workflow":
			return SourceKindWorkflow
		}
		return SourceKindOther
	default:
		return SourceKindOther
	}
}

// readAndHash opens relPath relative to the confined root and reads it
// verbatim. It performs a bounded read of at most maxBytes+1 bytes and
// rejects NUL bytes and invalid UTF-8.
//
// Confinement and identity:
//   - Opening via root.Open(relPath) ensures the kernel resolves the path
//     relative to the opened root directory. A symlink that would escape the
//     root is rejected with an escape error rather than followed.
//   - After Open we stat the file descriptor and compare against the Lstat
//     info from WalkDir. If the inode or device changed between Lstat and
//     Open (path swap), we refuse to index the content.
//   - After the bounded read we stat again. If the size changed between
//     pre-open Lstat and post-read stat (concurrent growth), we refuse to
//     index content whose length is unstable.
//
// All three checks are required because the Memory source clone is a
// watched, mutable directory and the scanner must not be fooled by an
// attacker who can race the scan.
func readAndHash(root *os.Root, relPath string, walkStat os.FileInfo, maxBytes int64) ([]byte, string, error) {
	f, err := root.Open(relPath)
	if err != nil {
		// os.Root.Open returns *fs.PathError wrapping syscall errors that
		// include EAGENTTEAM when an escape symlink is detected. Report as
		// ErrRootEscape so callers can mark the scan degraded.
		if isEscapeError(err) {
			return nil, "", fmt.Errorf("%w: open %q: %v", ErrRootEscape, relPath, err)
		}
		return nil, "", err
	}
	defer f.Close()

	preStat, err := f.Stat()
	if err != nil {
		return nil, "", err
	}
	if !preStat.Mode().IsRegular() {
		return nil, "", ErrBinaryContent
	}
	// Identity check: the inode we just opened must match the inode WalkDir
	// Lstat'd. A mismatch means the path was swapped between Lstat and Open.
	if walkStat != nil && !os.SameFile(walkStat, preStat) {
		return nil, "", fmt.Errorf("%w: identity changed between lstat and open for %q", ErrRootEscape, relPath)
	}

	data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
	if err != nil {
		return nil, "", err
	}
	if int64(len(data)) > maxBytes {
		return nil, "", ErrFileTooLarge
	}

	postStat, err := f.Stat()
	if err != nil {
		return nil, "", err
	}
	if !os.SameFile(preStat, postStat) {
		return nil, "", fmt.Errorf("%w: identity changed during read for %q", ErrRootEscape, relPath)
	}
	if postStat.Size() != preStat.Size() {
		// File grew or shrank during the read; treat as concurrent mutation.
		return nil, "", ErrFileTooLarge
	}

	if bytes.IndexByte(data, 0) >= 0 {
		return nil, "", ErrBinaryContent
	}
	if !utf8.Valid(data) {
		return nil, "", ErrInvalidUTF8
	}

	h := sha256.New()
	if _, err := h.Write(data); err != nil {
		return nil, "", err
	}
	return data, hex.EncodeToString(h.Sum(nil)), nil
}

// isEscapeError returns true for syscall errors that indicate an attempted
// escape from the confined root (symlink to outside, absolute path injection,
// etc.). os.Root.Open uses openat with RESOLVE_BENEATH / RESOLVE_IN_ROOT
// under the hood; the resulting EXDEV or EAGENTTEAM is reported here.
func isEscapeError(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	// os.PathError stringifies as "op path: err"; we look for the sentinel
	// substrings that openat-on-root produces on escape attempts.
	return strings.Contains(msg, "symbolic link") ||
		strings.Contains(msg, "not a directory") ||
		strings.Contains(msg, "escape") ||
		errors.Is(err, fs.ErrNotExist)
}

func hasNUL(s string) bool {
	return strings.IndexByte(s, 0) >= 0
}