package memorysearch

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/blevesearch/bleve/v2"
	"github.com/blevesearch/bleve/v2/search/query"
)

// DerivedDirName is the agent work directory's conventional home for generated
// artifacts. The source scanner skips it (alongside .git) so derived
// state is never indexed as Memory, and the search derivative lives under it.
const DerivedDirName = ".derived"

// DerivativeDirName is the final path segment of the search derivative, nested
// under DerivedDirName in the agent work directory. It is never placed inside
// the Memory source clone. Exposed so callers (the daemon) can construct
// the derivative root from the same resolver they use for Binding.
const DerivativeDirName = "memory_search"

// generationsDirName holds one subdirectory per build, named by fingerprint
// and timestamp. Only one generation per fingerprint is current at a time;
// older generations may be retained for rollback or removed by GC.
const generationsDirName = "generations"

// indexDirName is the Bleve scorch root inside a generation directory.
const indexDirName = "index"

// buildCompleteMarker is the atomic marker renamed into place as the final
// step of a successful build. A generation directory without this marker is
// considered partial and must never be opened by the query path.
const buildCompleteMarker = "BUILD_COMPLETE"

// smokeMatchAll is the fuzzy ceiling for the build-time smoke query. The
// smoke exists to catch a build that wrote chunks into the manifest but
// failed to land them in Bleve; it does not need to be precise.
const smokeMatchAll = 5

// Generation is the openable handle to one completed generation directory.
// The caller obtains a Generation from Build (which performs smoke
// validation, then closes the Bleve writer before returning) or from
// OpenGeneration (which opens the index read-only after revalidating the
// manifest). A Generation returned by Build never holds an open Bleve
// writer: the build-time writer is closed as the last step before return so
// that the query path can reopen the same directory without competing for
// bbolt's exclusive flock. The index handle inside a Generation returned
// by OpenGeneration is held until Close.
//
// Commit B extends Generation to support pooled borrowing via the
// poolRegistry:
//   - release, when non-nil, is the pool-provided release callback that
//     decrements refs and (on last release after retire) closes the
//     underlying Bleve handle. Build/OpenGeneration return Generations
//     with release nil and Close falls back to direct index close.
//   - closeOnce/closeErr make Close idempotent across all paths and
//     concurrent callers. A wrapper Close racing a registry shutdown
//     both observe the same cached result.
type Generation struct {
	dir      string
	manifest *Manifest
	index    bleve.Index // nil for Build-returned generations
	release  func() error
	// closeIndex is a narrow close seam used by the pool tests. Production
	// OpenGeneration leaves it nil and Close calls index.Close directly.
	closeIndex func() error

	closeOnce sync.Once
	closeErr  error
}

// Dir returns the absolute generation directory path.
func (g *Generation) Dir() string { return g.dir }

// Manifest returns the in-memory manifest for the generation. The manifest
// has already been validated against the binding at open time.
func (g *Generation) Manifest() *Manifest { return g.manifest }

// Index returns the open Bleve index, or nil if the Generation was produced
// by Build (which closes the writer before returning). Callers that need to
// query the index should call OpenGeneration to obtain a Generation that
// owns a reader handle. Holding a Build-returned Generation does not pin
// the index's bbolt lock.
func (g *Generation) Index() bleve.Index { return g.index }

// Close releases the Bleve reader handle for generations produced by
// OpenGeneration. For generations produced by Build it is a no-op because
// Build already closed its writer. Safe to call multiple times.
//
// Commit B: Close is now concurrency-safe and idempotent across all
// generation paths via closeOnce. Pooled wrappers (release != nil)
// dispatch to the pool release callback; direct OpenGeneration handles
// (release == nil, index != nil) close the Bleve handle inline. Either
// path caches the result in closeErr so a second Close returns the same
// error the first observed.
func (g *Generation) Close() error {
	if g == nil {
		return nil
	}
	g.closeOnce.Do(func() {
		if g.release != nil {
			g.closeErr = g.release()
			return
		}
		if g.closeIndex != nil {
			g.closeErr = g.closeIndex()
			return
		}
		if g.index != nil {
			g.closeErr = g.index.Close()
			g.index = nil
		}
	})
	return g.closeErr
}

// BuildOptions is the input to Build. Caller supplies the resolved binding,
// a fresh SourceSnapshot, and the pre-computed chunks for that snapshot.
// Build does not re-chunk; it trusts that chunks correspond to snap. Both
// the source clone root and the derivative root are derived from the
// resolver plus the binding's LocalAgentID; the caller cannot supply or
// override either path. This is the trust boundary: a runtime-supplied
// path must never land in a derivative file path.
type BuildOptions struct {
	// Binding is the validated binding the derivative is being built for.
	// Build re-runs ValidateWith(Resolver) before doing any disk work, so
	// a Binding deserialized from disk without re-validation is rejected.
	Binding *Binding
	// Resolver is the agent work-dir resolver. It must be the same
	// resolver used to construct Binding. Build uses it to re-check
	// source-root containment and to derive the derivative root;
	// the response-contract query path will re-check it again on every
	// OpenGeneration.
	Resolver AgentWorkDirResolver
	// Snapshot is the source view this generation indexes.
	Snapshot SourceSnapshot
	// Chunks is the full chunk inventory for Snapshot. Build does not
	// re-derive chunks because chunking is the responsibility of the
	// reconciler and may legitimately be cached across builds.
	Chunks []SearchChunk
	// Clock returns the build timestamps recorded in the manifest. Tests
	// inject a fake clock; production leaves it nil and Build uses
	// time.Now(). The clock is only consulted twice per build.
	Clock func() time.Time
}

// Build runs a full generation build: it creates a fresh generation
// directory, indexes every chunk into a new scorch index, writes the
// manifest, and runs smoke validation before returning. On any failure the
// partial generation directory is removed.
//
// Build does NOT publish the generation: it does not move the manifest into
// the derivative root or update the current-generation pointer. The
// reconciler (Task 5) performs the publish step after Build returns.
func Build(ctx context.Context, opts BuildOptions) (*Generation, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if opts.Binding == nil {
		return nil, fmt.Errorf("%w: binding is required", ErrInvalidBuildOptions)
	}
	if opts.Resolver == nil {
		return nil, fmt.Errorf("%w: resolver is required", ErrInvalidBuildOptions)
	}
	if err := opts.Binding.ValidateWith(opts.Resolver); err != nil {
		return nil, fmt.Errorf("%w: binding failed resolver validation: %v", ErrInvalidBuildOptions, err)
	}
	// Derive both roots from the resolver plus the binding's agent ID.
	// Runtime-supplied paths are forbidden anywhere in the build pipeline.
	workDir := opts.Resolver(opts.Binding.LocalAgentID)
	if workDir == "" {
		return nil, fmt.Errorf("%w: resolver returned empty work dir for %q",
			ErrInvalidBuildOptions, opts.Binding.LocalAgentID)
	}
	expectedSourceRoot := filepath.Join(workDir, memoryCloneDir)
	expectedDerivative := DerivativeRootFromWorkDir(workDir)
	if err := VerifyDerivativeRootSibling(expectedSourceRoot, expectedDerivative); err != nil {
		// This should be unreachable because the layout is fixed, but the
		// check remains as a defensive guard against a future resolver
		// returning a workDir whose source/derivative pair violates
		// the sibling invariant.
		return nil, err
	}
	clock := opts.Clock
	if clock == nil {
		clock = time.Now
	}

	startedAt := clock()

	// Validate chunks against the snapshot BEFORE touching disk. This runs
	// ahead of indexing so a forged batch (missing path, hash drift, empty
	// or duplicate IDs) fails fast without leaving a partial generation
	// directory or burning indexing work that will be discarded.
	manifest, err := newManifestFromSnapshot(opts.Binding, opts.Snapshot, opts.Chunks, startedAt, clock())
	if err != nil {
		return nil, err
	}

	genDir, err := allocateGenerationDir(expectedDerivative, opts.Binding, startedAt)
	if err != nil {
		return nil, err
	}
	cleanup := func(reason string) {
		_ = os.RemoveAll(genDir)
	}
	indexPath := filepath.Join(genDir, indexDirName)
	mapping, err := BuildIndexMapping()
	if err != nil {
		cleanup("mapping")
		return nil, err
	}
	idx, err := bleve.New(indexPath, mapping)
	if err != nil {
		cleanup("bleve_new")
		return nil, fmt.Errorf("memory search: bleve new failed: %w", err)
	}

	if err := indexChunks(ctx, idx, opts.Chunks); err != nil {
		_ = idx.Close()
		cleanup("index_chunks")
		return nil, err
	}

	// Re-stamp the completedAt timestamp now that indexing is done; the
	// first manifest pass recorded a placeholder completedAt so it could
	// run before disk allocation.
	manifest.BuildCompletedAt = clock()
	manifestPath := filepath.Join(genDir, manifestFilename)
	if err := manifest.SaveAtomic(manifestPath); err != nil {
		_ = idx.Close()
		cleanup("manifest_save")
		return nil, err
	}

	if err := smokeValidate(ctx, idx, len(opts.Chunks)); err != nil {
		_ = idx.Close()
		cleanup("smoke")
		return nil, fmt.Errorf("memory search: smoke validation failed: %w", err)
	}

	if err := writeBuildCompleteMarker(genDir); err != nil {
		_ = idx.Close()
		cleanup("marker")
		return nil, err
	}

	// Close the build-time writer before returning. The Generation handle
	// does not own an open Bleve index once Build returns: the index's
	// bbolt flock must be free so that the publish step and the query path
	// can reopen the same directory without competing with a writer that
	// is no longer doing work. Callers that need to query the index must
	// call OpenGeneration, which opens a fresh reader handle.
	if err := idx.Close(); err != nil {
		cleanup("writer_close")
		return nil, fmt.Errorf("memory search: close build writer: %w", err)
	}

	return &Generation{dir: genDir, manifest: manifest, index: nil}, nil
}

// OpenGeneration opens an existing generation directory for read. The
// directory must carry a BUILD_COMPLETE marker; otherwise the open is
// fail-closed with ErrGenerationPartial.
//
// Trust boundary: genDir is treated as untrusted input. Before any file under
// it is read, OpenGeneration re-validates the binding against the resolver,
// derives the expected generations parent (`<workDir>/memory_search/generations`),
// and verifies genDir is contained under that parent. A caller that passes
// a path outside the resolver-derived derivative root — for example a forged
// current-generation pointer, or a directory copied verbatim from another
// agent — is rejected with ErrGenerationOutsideRoot. This stops an attacker
// who can write a manifest + index + marker anywhere on disk from talking
// the query path into opening that arbitrary directory.
//
// After containment passes, the marker is checked, then the manifest is
// loaded. The manifest load re-runs Binding.ValidateWith(resolver) and the
// bijection check, so a copied-but-tampered manifest still fails closed.
func OpenGeneration(ctx context.Context, binding *Binding, resolver AgentWorkDirResolver, genDir string) (*Generation, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if binding == nil {
		return nil, fmt.Errorf("%w: binding is required", ErrInvalidBuildOptions)
	}
	if resolver == nil {
		return nil, fmt.Errorf("%w: resolver is required", ErrInvalidBuildOptions)
	}
	if err := binding.ValidateWith(resolver); err != nil {
		return nil, fmt.Errorf("%w: binding failed resolver validation: %v",
			ErrInvalidBuildOptions, err)
	}
	if genDir == "" {
		return nil, fmt.Errorf("%w: generation directory is empty", ErrInvalidBuildOptions)
	}
	workDir := resolver(binding.LocalAgentID)
	if workDir == "" {
		return nil, fmt.Errorf("%w: resolver returned empty work dir for %q",
			ErrInvalidBuildOptions, binding.LocalAgentID)
	}
	expectedParent := filepath.Join(DerivativeRootFromWorkDir(workDir), generationsDirName)
	if err := verifyGenerationContained(expectedParent, genDir); err != nil {
		return nil, err
	}
	if _, err := os.Stat(filepath.Join(genDir, buildCompleteMarker)); err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, fmt.Errorf("%w: %s", ErrGenerationPartial, genDir)
		}
		return nil, err
	}
	manifestPath := filepath.Join(genDir, manifestFilename)
	manifest, err := LoadManifest(ctx, manifestPath, resolver)
	if err != nil {
		return nil, err
	}
	// Cross-binding gate. LoadManifest → Manifest.Validate only verifies that
	// manifest.BindingFingerprint matches manifest.Binding's own fingerprint;
	// it does NOT compare against the caller's binding. Without this gate two
	// bindings that share agent + resolver-derived root but differ on
	// MemoryID / Branch / analyzer / chunker both pass Validate, and one can
	// open the other's generation. The query path then uses manifest.Binding
	// for source reads while SearchSource.Binding() reports the caller's,
	// so serving root and data root would silently belong to different
	// bindings. codex3 round 5 final gate: close the fence before bleve.Open.
	callerFP := binding.Fingerprint()
	if callerFP != manifest.BindingFingerprint {
		return nil, fmt.Errorf("%w: caller fingerprint %q does not match manifest fingerprint %q",
			ErrManifestBindingMismatch, callerFP, manifest.BindingFingerprint)
	}
	indexPath := filepath.Join(genDir, indexDirName)
	idx, err := bleve.Open(indexPath)
	if err != nil {
		return nil, fmt.Errorf("memory search: bleve open failed: %w", err)
	}
	return &Generation{dir: genDir, manifest: manifest, index: idx}, nil
}

// verifyGenerationContained ensures genDir is a descendant of expectedParent.
//
// Both paths are EvalSymlinks-resolved first. A purely lexical filepath.Rel
// check is not enough: a symlink inside expectedParent that points elsewhere
// on disk would pass the lexical check while letting OpenGeneration read
// files outside the agent's derivative root. The generations parent is
// mutable state (current-pointer writes, GC, local development), so leaving
// symlinks out of the threat model is unsound.
//
// After resolution, genDir must be strictly below expectedParent — equality
// is rejected because the parent itself is not a generation. A residual
// TOCTOU window between EvalSymlinks and the subsequent os.Stat / bleve.Open
// remains; closing it fully requires openat2 / O_PATH which Go's stdlib does
// not expose portably. The EvalSymlinks gate stops the straightforward
// planted-link attack codex3 reproduced.
func verifyGenerationContained(expectedParent, genDir string) error {
	if expectedParent == "" || genDir == "" {
		return fmt.Errorf("%w: empty path", ErrInvalidBuildOptions)
	}
	resolvedParent, err := filepath.EvalSymlinks(expectedParent)
	if err != nil {
		// An absent parent is a setup error, not a containment violation;
		// return it as invalid build options so the caller sees the real
		// cause rather than a misleading "outside root".
		return fmt.Errorf("%w: resolve generations parent: %v", ErrInvalidBuildOptions, err)
	}
	resolvedGen, err := filepath.EvalSymlinks(genDir)
	if err != nil {
		return fmt.Errorf("%w: resolve generation dir: %v", ErrInvalidBuildOptions, err)
	}
	c := filepath.Clean(resolvedParent)
	d := filepath.Clean(resolvedGen)
	if d == c {
		return fmt.Errorf("%w: generation dir equals generations parent", ErrGenerationOutsideRoot)
	}
	rel, err := filepath.Rel(c, d)
	if err != nil {
		return fmt.Errorf("%w: rel: %v", ErrGenerationOutsideRoot, err)
	}
	if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
		return fmt.Errorf("%w: %s resolves outside %s", ErrGenerationOutsideRoot, genDir, c)
	}
	return nil
}

// indexChunks batches chunks into the Bleve index. Bleve's batch API is the
// only safe way to bulk-index without interleaving segment flushes; the
// per-document Index path forces a flush on every call and would dominate
// build time on source clones with thousands of files.
func indexChunks(ctx context.Context, idx bleve.Index, chunks []SearchChunk) error {
	if len(chunks) == 0 {
		return nil
	}
	const batchSize = 256
	batch := idx.NewBatch()
	for _, c := range chunks {
		if err := ctx.Err(); err != nil {
			return err
		}
		doc := chunkDocument(c)
		if err := batch.Index(c.ID, doc); err != nil {
			return fmt.Errorf("memory search: batch index %q: %w", c.ID, err)
		}
		if batch.Size() >= batchSize {
			if err := idx.Batch(batch); err != nil {
				return fmt.Errorf("memory search: batch flush: %w", err)
			}
			batch.Reset()
		}
	}
	if batch.Size() > 0 {
		if err := idx.Batch(batch); err != nil {
			return fmt.Errorf("memory search: final batch flush: %w", err)
		}
	}
	return nil
}

// chunkDocument is the Bleve document shape for a SearchChunk. The field
// names match the *fieldXxx* constants in mapping.go; renaming a field here
// without a matching mapping change silently drops the value because
// IndexDynamic is false.
//
// Identifiers are normalized (lowercased) before indexing so the
// exact-identifier route can find them regardless of source case
// (e.g. "@CamelCase" indexed, "@camelcase" queried). SearchChunk.Identifiers
// retains its original source case for callers that use it outside Bleve;
// only the indexed value is normalized.
func chunkDocument(c SearchChunk) map[string]interface{} {
	return map[string]interface{}{
		fieldPath:            c.Path,
		fieldTitle:           c.Title,
		fieldHeadingPath:     c.HeadingPath,
		fieldHeadingText:     c.HeadingText,
		fieldHeadingAnchor:   c.HeadingAnchor,
		fieldBody:            c.Body,
		fieldKind:            c.Kind,
		fieldFormat:          c.Format,
		fieldArchived:        c.Archived,
		fieldIdentifiers:     normalizedIdentifiers(c.Identifiers),
		fieldIdentifierParts: c.IdentifierParts,
		fieldStartLine:       c.StartLine,
		fieldEndLine:         c.EndLine,
		fieldOrdinal:         c.Ordinal,
		fieldContentHash:     c.ContentHash,
		fieldSourceRevision:  c.SourceRevision,
	}
}

// smokeValidate runs a trivial match-all query against the freshly built
// index and asserts that Bleve's reported document count matches the number
// of chunks the caller indexed. A mismatch means the indexing pipeline
// silently dropped documents or Bleve lost them on flush; either case is
// fail-closed.
//
// The query caps at smokeMatchAll so the smoke is fast on large generations:
// we only need to know the index answers queries at all, not enumerate every
// document. The DocCount check is the real guarantee that every chunk landed.
func smokeValidate(ctx context.Context, idx bleve.Index, expectChunks int) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	got, err := idx.DocCount()
	if err != nil {
		return fmt.Errorf("doc count: %w", err)
	}
	if int(got) != expectChunks {
		return fmt.Errorf("%w: indexed %d chunks but bleve reports %d docs",
			ErrSmokeDocCount, expectChunks, got)
	}
	if expectChunks == 0 {
		// An empty generation is a valid outcome of scanning an empty
		// source clone. The smoke query path is skipped because a
		// match-all over zero docs is not a useful guarantee.
		return nil
	}
	q := query.NewMatchAllQuery()
	req := bleve.NewSearchRequestOptions(q, smokeMatchAll, 0, false)
	req.Fields = []string{fieldPath}
	res, err := idx.SearchInContext(ctx, req)
	if err != nil {
		return fmt.Errorf("smoke search: %w", err)
	}
	if res == nil || len(res.Hits) == 0 {
		return fmt.Errorf("%w: match-all returned no hits over a non-empty index",
			ErrSmokeEmpty)
	}
	return nil
}

// allocateGenerationDir creates and returns a fresh generation directory
// under derivativeRoot/generations. The directory name combines the binding
// fingerprint short, the start time in nanoseconds, and an OS-chosen random
// suffix. The OS-level MkdirTemp guarantees uniqueness even when two builds
// race in the same nanosecond (e.g. a retry agentteam with an injected fake
// clock), so a same-second rebuild never collides with a prior attempt.
func allocateGenerationDir(derivativeRoot string, b *Binding, startedAt time.Time) (string, error) {
	generationsDir := filepath.Join(derivativeRoot, generationsDirName)
	if err := os.MkdirAll(generationsDir, 0o755); err != nil {
		return "", fmt.Errorf("memory search: mkdir generations: %w", err)
	}
	fp := b.Fingerprint()
	if len(fp) > 12 {
		fp = fp[:12]
	}
	prefix := fmt.Sprintf("%s-%d-", fp, startedAt.UnixNano())
	genDir, err := os.MkdirTemp(generationsDir, prefix+"*")
	if err != nil {
		return "", fmt.Errorf("memory search: allocate generation: %w", err)
	}
	if err := os.MkdirAll(filepath.Join(genDir, indexDirName), 0o755); err != nil {
		_ = os.RemoveAll(genDir)
		return "", fmt.Errorf("memory search: mkdir index: %w", err)
	}
	return genDir, nil
}

// writeBuildCompleteMarker atomically renames a temp file into
// buildCompleteMarker inside genDir. The marker is the last thing the build
// writes; if it appears, the generation is fully formed and may be opened
// by OpenGeneration.
func writeBuildCompleteMarker(genDir string) error {
	tmp, err := os.CreateTemp(genDir, ".build-complete-*.tmp")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	if _, err := tmp.WriteString(buildCompleteMarker); err != nil {
		_ = tmp.Close()
		_ = os.Remove(tmpName)
		return err
	}
	if err := tmp.Sync(); err != nil {
		_ = tmp.Close()
		_ = os.Remove(tmpName)
		return err
	}
	if err := tmp.Close(); err != nil {
		_ = os.Remove(tmpName)
		return err
	}
	final := filepath.Join(genDir, buildCompleteMarker)
	if err := os.Rename(tmpName, final); err != nil {
		_ = os.Remove(tmpName)
		return err
	}
	if err := syncDir(genDir); err != nil {
		return err
	}
	return nil
}

// DerivativeRootFromWorkDir joins an agent work directory with the
// derivative segment. It mirrors the source-root derivation in
// NewBinding so the daemon can derive both paths from a single resolver
// call.
func DerivativeRootFromWorkDir(workDir string) string {
	if workDir == "" {
		return ""
	}
	return filepath.Join(workDir, DerivedDirName, DerivativeDirName)
}

// VerifyDerivativeRootSibling tests whether derivativeRoot is a sibling (not
// a descendant) of resolvedRoot, so the build pipeline can refuse to place
// the derivative inside the source clone. The spec prohibits nesting
// because the source clone is the source of truth and must not carry
// derived state.
func VerifyDerivativeRootSibling(resolvedRoot, derivativeRoot string) error {
	c := filepath.Clean(resolvedRoot)
	d := filepath.Clean(derivativeRoot)
	if c == "" || d == "" {
		return fmt.Errorf("%w: empty root", ErrInvalidBuildOptions)
	}
	if c == d {
		return fmt.Errorf("%w: derivative equals source root", ErrInvalidBuildOptions)
	}
	rel, err := filepath.Rel(c, d)
	if err != nil {
		return fmt.Errorf("%w: rel: %v", ErrInvalidBuildOptions, err)
	}
	if !strings.HasPrefix(rel, ".."+string(filepath.Separator)) && rel != ".." {
		return fmt.Errorf("%w: derivative %q is inside source root %q",
			ErrInvalidBuildOptions, d, c)
	}
	return nil
}

var (
	// ErrInvalidBuildOptions signals that BuildOptions failed validation.
	ErrInvalidBuildOptions = errors.New("memory search: invalid build options")

	// ErrGenerationPartial signals that a generation directory is missing
	// its BUILD_COMPLETE marker. The query path must treat this as not
	// ready.
	ErrGenerationPartial = errors.New("memory search: generation is partial")

	// ErrGenerationOutsideRoot signals that OpenGeneration was asked to open
	// a generation directory that is not contained under the resolver-
	// derived derivative root for the binding's agent. A copied generation
	// placed elsewhere on disk — or a forged pointer — is rejected here.
	ErrGenerationOutsideRoot = errors.New("memory search: generation outside derivative root")

	// ErrSmokeDocCount signals that Bleve's reported document count does
	// not match the number of chunks the build indexed.
	ErrSmokeDocCount = errors.New("memory search: smoke document count mismatch")

	// ErrSmokeEmpty signals that a match-all query against a non-empty
	// index returned no hits. The index is unusable.
	ErrSmokeEmpty = errors.New("memory search: smoke returned no hits")
)