package memorysearch

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	lifecycleSchemaVersion = 1
	lifecycleFileName      = "lifecycle.json"
)

type BindingLifecycleState string

const (
	BindingLifecycleActive       BindingLifecycleState = "active"
	BindingLifecycleFenced       BindingLifecycleState = "fenced"
	BindingLifecycleUnbound      BindingLifecycleState = "unbound"
	BindingLifecycleUnknownOwner BindingLifecycleState = "unknown_owner"
)

type GenerationLifecycleKind string

const (
	GenerationLifecycleCurrent  GenerationLifecycleKind = "current"
	GenerationLifecycleRollback GenerationLifecycleKind = "rollback"
)

type GenerationLifecycle struct {
	Kind                   GenerationLifecycleKind `json:"kind"`
	PublishedAt            time.Time               `json:"published_at"`
	SupersededAt           *time.Time              `json:"superseded_at,omitempty"`
	FirstSuccessfulQueryAt *time.Time              `json:"first_successful_query_at,omitempty"`
}

type LifecycleState struct {
	SchemaVersion      int                            `json:"schema_version"`
	BindingFingerprint string                         `json:"binding_fingerprint"`
	BindingState       BindingLifecycleState          `json:"binding_state"`
	BindingChangedAt   time.Time                      `json:"binding_changed_at"`
	Generations        map[string]GenerationLifecycle `json:"generations"`
}

// lifecycleStore serializes all sidecar read-modify-write operations for one
// reconciler. The file sits at the per-agent derivative root and therefore
// shares the same crash-safe write primitive as current.json/pending.json.
type lifecycleStore struct {
	mu             sync.Mutex
	derivativeRoot string
	expectedFP     string
	clock          func() time.Time
}

func newLifecycleStore(
	derivativeRoot string,
	binding *Binding,
	clock func() time.Time,
) *lifecycleStore {
	return &lifecycleStore{
		derivativeRoot: derivativeRoot,
		expectedFP:     binding.Fingerprint(),
		clock:          clock,
	}
}

func (s *lifecycleStore) recordPublish(
	current Current,
	previousBasename string,
) error {
	if err := validateBasename(current.GenerationBasename); err != nil {
		return err
	}
	if err := validateBasename(previousBasename); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	state, err := s.loadForPublishLocked(current.PublishedAt)
	if err != nil {
		return err
	}
	for basename, generation := range state.Generations {
		if basename == current.GenerationBasename ||
			generation.Kind != GenerationLifecycleCurrent {
			continue
		}
		generation.Kind = GenerationLifecycleRollback
		supersededAt := current.PublishedAt
		generation.SupersededAt = &supersededAt
		state.Generations[basename] = generation
	}
	if previousBasename != "" && previousBasename != current.GenerationBasename {
		previous := state.Generations[previousBasename]
		previous.Kind = GenerationLifecycleRollback
		if previous.PublishedAt.IsZero() {
			previous.PublishedAt = publishedAtForGeneration(
				s.derivativeRoot,
				previousBasename,
				current.PublishedAt,
			)
		}
		supersededAt := current.PublishedAt
		previous.SupersededAt = &supersededAt
		state.Generations[previousBasename] = previous
	}
	currentState := state.Generations[current.GenerationBasename]
	currentState.Kind = GenerationLifecycleCurrent
	currentState.PublishedAt = current.PublishedAt
	currentState.SupersededAt = nil
	state.Generations[current.GenerationBasename] = currentState
	if state.BindingState != BindingLifecycleActive {
		state.BindingState = BindingLifecycleActive
		state.BindingChangedAt = current.PublishedAt
	}
	return s.writeLocked(state)
}

func publishedAtForGeneration(
	derivativeRoot string,
	basename string,
	fallback time.Time,
) time.Time {
	info, err := os.Stat(filepath.Join(
		derivativeRoot,
		generationsDirName,
		basename,
	))
	if err != nil {
		return fallback
	}
	return info.ModTime()
}

// ensureCurrent bootstraps lifecycle metadata for a generation published by
// an older binary. A valid existing current record is preserved verbatim,
// including FirstSuccessfulQueryAt. Missing/corrupt/cross-binding state is
// repaired from the authoritative current pointer.
func (s *lifecycleStore) ensureCurrent(current Current) error {
	if err := validateBasename(current.GenerationBasename); err != nil ||
		current.GenerationBasename == "" {
		return fmt.Errorf("%w: current basename %q",
			ErrLifecycleCorrupt, current.GenerationBasename)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	state, err := s.loadLocked()
	if err == nil {
		generation, ok := state.Generations[current.GenerationBasename]
		if ok && generation.Kind == GenerationLifecycleCurrent {
			return nil
		}
		// The pointer is authoritative. A valid sidecar that disagrees is
		// stale state from an interrupted upgrade/publish and is rebuilt.
	}
	if err != nil &&
		!errors.Is(err, os.ErrNotExist) &&
		!errors.Is(err, ErrLifecycleCorrupt) &&
		!errors.Is(err, ErrLifecycleBindingMismatch) {
		return err
	}
	repaired := &LifecycleState{
		SchemaVersion:      lifecycleSchemaVersion,
		BindingFingerprint: s.expectedFP,
		BindingState:       BindingLifecycleActive,
		BindingChangedAt:   current.PublishedAt,
		Generations: map[string]GenerationLifecycle{
			current.GenerationBasename: {
				Kind:        GenerationLifecycleCurrent,
				PublishedAt: current.PublishedAt,
			},
		},
	}
	return s.writeLocked(repaired)
}

func (s *lifecycleStore) loadForPublishLocked(now time.Time) (*LifecycleState, error) {
	state, err := s.loadLocked()
	if err == nil {
		return state, nil
	}
	if !errors.Is(err, os.ErrNotExist) &&
		!errors.Is(err, ErrLifecycleCorrupt) &&
		!errors.Is(err, ErrLifecycleBindingMismatch) {
		return nil, err
	}
	// A successful publish is the authoritative repair point for a missing,
	// corrupt, or previous-binding sidecar. Reconstruct only the state the
	// current process can prove from current.json; retention treats older
	// unrecorded directories conservatively in Commit D.
	return &LifecycleState{
		SchemaVersion:      lifecycleSchemaVersion,
		BindingFingerprint: s.expectedFP,
		BindingState:       BindingLifecycleActive,
		BindingChangedAt:   now,
		Generations:        make(map[string]GenerationLifecycle),
	}, nil
}

// markQuerySuccess records the first successful query for the currently
// published generation. It is first-write-wins and rechecks current.json under
// the store mutex so a query which finishes after a publish cannot stamp the
// now-rollback generation as the new current's success.
func (s *lifecycleStore) markQuerySuccess(
	ctx context.Context,
	generationBasename string,
) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := validateBasename(generationBasename); err != nil || generationBasename == "" {
		return fmt.Errorf("%w: generation basename %q", ErrInvalidBasename, generationBasename)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	current, err := loadCurrent(s.derivativeRoot)
	if err != nil {
		return err
	}
	if current.BindingFingerprint != s.expectedFP ||
		current.GenerationBasename != generationBasename {
		return fmt.Errorf("%w: query generation %q is no longer current",
			ErrQueryGenerationSuperseded, generationBasename)
	}
	state, err := s.loadOrNewLocked(current.PublishedAt)
	if err != nil {
		return err
	}
	generation, ok := state.Generations[generationBasename]
	if !ok || generation.Kind != GenerationLifecycleCurrent {
		return fmt.Errorf("%w: lifecycle has no current generation %q",
			ErrLifecycleCorrupt, generationBasename)
	}
	if generation.FirstSuccessfulQueryAt != nil {
		return nil
	}
	now := s.clock()
	generation.FirstSuccessfulQueryAt = &now
	state.Generations[generationBasename] = generation
	return s.writeLocked(state)
}

func (s *lifecycleStore) setBindingState(next BindingLifecycleState) error {
	switch next {
	case BindingLifecycleActive, BindingLifecycleFenced,
		BindingLifecycleUnbound, BindingLifecycleUnknownOwner:
	default:
		return fmt.Errorf("%w: binding state %q", ErrLifecycleCorrupt, next)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	state, err := s.loadOrNewLocked(s.clock())
	if err != nil {
		return err
	}
	if state.BindingState == next {
		return nil
	}
	state.BindingState = next
	state.BindingChangedAt = s.clock()
	return s.writeLocked(state)
}

func (s *lifecycleStore) loadOrNewLocked(now time.Time) (*LifecycleState, error) {
	state, err := s.loadLocked()
	if err == nil {
		return state, nil
	}
	if !errors.Is(err, os.ErrNotExist) {
		return nil, err
	}
	return &LifecycleState{
		SchemaVersion:      lifecycleSchemaVersion,
		BindingFingerprint: s.expectedFP,
		BindingState:       BindingLifecycleActive,
		BindingChangedAt:   now,
		Generations:        make(map[string]GenerationLifecycle),
	}, nil
}

func (s *lifecycleStore) load() (*LifecycleState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.loadLocked()
}

func (s *lifecycleStore) loadLocked() (*LifecycleState, error) {
	raw, err := os.ReadFile(filepath.Join(s.derivativeRoot, lifecycleFileName))
	if err != nil {
		return nil, err
	}
	var state LifecycleState
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&state); err != nil {
		return nil, fmt.Errorf("%w: decode: %v", ErrLifecycleCorrupt, err)
	}
	if err := validateLifecycleState(&state, s.expectedFP); err != nil {
		return nil, err
	}
	return &state, nil
}

func (s *lifecycleStore) writeLocked(state *LifecycleState) error {
	if err := validateLifecycleState(state, s.expectedFP); err != nil {
		return err
	}
	return writeJSONAtomic(filepath.Join(s.derivativeRoot, lifecycleFileName), state)
}

func (s *lifecycleStore) replaceForTest(state *LifecycleState) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.writeLocked(state)
}

func validateLifecycleState(state *LifecycleState, expectedFP string) error {
	if state == nil {
		return fmt.Errorf("%w: nil state", ErrLifecycleCorrupt)
	}
	if state.SchemaVersion != lifecycleSchemaVersion {
		return fmt.Errorf("%w: schema %d", ErrLifecycleCorrupt, state.SchemaVersion)
	}
	if state.BindingFingerprint != expectedFP {
		return fmt.Errorf("%w: lifecycle fingerprint %q does not match %q",
			ErrLifecycleBindingMismatch, state.BindingFingerprint, expectedFP)
	}
	switch state.BindingState {
	case BindingLifecycleActive, BindingLifecycleFenced,
		BindingLifecycleUnbound, BindingLifecycleUnknownOwner:
	default:
		return fmt.Errorf("%w: binding state %q", ErrLifecycleCorrupt, state.BindingState)
	}
	if state.BindingChangedAt.IsZero() {
		return fmt.Errorf("%w: zero binding_changed_at", ErrLifecycleCorrupt)
	}
	if state.Generations == nil {
		return fmt.Errorf("%w: nil generations", ErrLifecycleCorrupt)
	}
	currentCount := 0
	for basename, generation := range state.Generations {
		if err := validateBasename(basename); err != nil || basename == "" {
			return fmt.Errorf("%w: generation basename %q", ErrLifecycleCorrupt, basename)
		}
		switch generation.Kind {
		case GenerationLifecycleCurrent:
			currentCount++
			if generation.SupersededAt != nil {
				return fmt.Errorf("%w: current %q has superseded_at",
					ErrLifecycleCorrupt, basename)
			}
		case GenerationLifecycleRollback:
			if generation.SupersededAt == nil {
				return fmt.Errorf("%w: rollback %q missing superseded_at",
					ErrLifecycleCorrupt, basename)
			}
		default:
			return fmt.Errorf("%w: generation %q kind %q",
				ErrLifecycleCorrupt, basename, generation.Kind)
		}
		if generation.PublishedAt.IsZero() {
			return fmt.Errorf("%w: generation %q zero published_at",
				ErrLifecycleCorrupt, basename)
		}
	}
	if currentCount > 1 {
		return fmt.Errorf("%w: %d current generations", ErrLifecycleCorrupt, currentCount)
	}
	return nil
}

var (
	ErrLifecycleCorrupt          = errors.New("memory search: lifecycle sidecar corrupt")
	ErrLifecycleBindingMismatch  = errors.New("memory search: lifecycle binding mismatch")
	ErrQueryGenerationSuperseded = errors.New("memory search: query generation superseded")
)