package ctxmgr

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"time"
	"unicode/utf8"
)

const (
	ShadowStatusPending = "pending"
	shadowJobMaxAttempt = 3
	shadowLeaseTTL      = 2 * time.Minute
)

// SemanticShadowJob is a durable unit of work for advisory semantic generation.
// Under the fail-closed gray rollout (2026-09-26) a validated OK report may
// replace the served Narrative for allowlisted sessions only; every failure
// path keeps the deterministic narrative.
type SemanticShadowJob struct {
	ID                   int64              `json:"id,omitempty"`
	SessionID            string             `json:"session_id"`
	SummaryRevision      int                `json:"summary_revision"`
	InputDigest          string             `json:"input_digest"`
	VisibilityDigest     string             `json:"visibility_digest"`
	WatermarkToMessageID int64              `json:"watermark_to_message_id"`
	PriorRevision        int                `json:"prior_revision"`
	Messages             []MinimizedMessage `json:"messages"`
	PriorNarrative       string             `json:"prior_narrative,omitempty"`
	PriorFacts           map[string]string  `json:"prior_facts,omitempty"`
	Status               string             `json:"status,omitempty"`
	Attempt              int                `json:"attempt,omitempty"`
}

type shadowJobPayload struct {
	Messages       []MinimizedMessage `json:"messages"`
	PriorNarrative string             `json:"prior_narrative,omitempty"`
	PriorFacts     map[string]string  `json:"prior_facts,omitempty"`
}

// EnqueueSemanticShadow schedules background semantic evaluation after the
// deterministic summary is already persisted and returned to the caller.
func (m *ContextManager) EnqueueSemanticShadow(ctx context.Context, job SemanticShadowJob) error {
	if m == nil || job.SessionID == "" || job.InputDigest == "" {
		return fmt.Errorf("invalid semantic shadow job")
	}
	if job.Status == "" {
		job.Status = "pending"
	}
	if m.pool == nil {
		m.rememberShadowJob(job)
		m.kickShadowWorker()
		return nil
	}
	payload, err := json.Marshal(shadowJobPayload{
		Messages:       job.Messages,
		PriorNarrative: job.PriorNarrative,
		PriorFacts:     job.PriorFacts,
	})
	if err != nil {
		return err
	}
	_, err = m.pool.Exec(ctx, `
		INSERT INTO chat_summary_shadow_jobs (
			session_id, summary_revision, input_digest, visibility_digest,
			watermark_to_message_id, prior_revision, payload, status, attempt, updated_at
		) VALUES ($1,$2,$3,$4,$5,$6,$7,'pending',0,now())
		ON CONFLICT (session_id, summary_revision, input_digest) DO NOTHING
	`, job.SessionID, job.SummaryRevision, job.InputDigest, job.VisibilityDigest,
		job.WatermarkToMessageID, job.PriorRevision, payload)
	if err != nil {
		slog.Warn("semantic shadow job persist failed; using in-process queue", "session_id", job.SessionID, "err", err)
		reportBackground(err)
		m.rememberShadowJob(job)
	}
	m.kickShadowWorker()
	return nil
}

func (m *ContextManager) rememberShadowJob(job SemanticShadowJob) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.shadowJobs == nil {
		m.shadowJobs = make([]SemanticShadowJob, 0, 4)
	}
	for _, existing := range m.shadowJobs {
		if existing.SessionID == job.SessionID && existing.SummaryRevision == job.SummaryRevision && existing.InputDigest == job.InputDigest {
			return
		}
	}
	m.shadowJobs = append(m.shadowJobs, job)
}

func (m *ContextManager) kickShadowWorker() {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.shadowWorkerRunning {
		return
	}
	m.shadowWorkerRunning = true
	go m.shadowWorkerLoop()
}

// StartSemanticShadowWorker starts a long-lived reclaim loop for durable jobs.
func (m *ContextManager) StartSemanticShadowWorker(stop <-chan struct{}) {
	if m == nil || m.pool == nil {
		return
	}
	m.mu.Lock()
	if m.shadowReclaimStarted {
		m.mu.Unlock()
		return
	}
	m.shadowReclaimStarted = true
	m.mu.Unlock()
	go func() {
		ticker := time.NewTicker(5 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-stop:
				return
			case <-ticker.C:
				m.kickShadowWorker()
			}
		}
	}()
	m.kickShadowWorker()
}

// DrainSemanticShadowTasks processes queued shadow jobs synchronously (tests / recovery).
func (m *ContextManager) DrainSemanticShadowTasks(ctx context.Context) int {
	if m == nil {
		return 0
	}
	n := 0
	for {
		job, ok := m.claimNextShadowJob(ctx)
		if !ok {
			return n
		}
		m.runShadowJob(ctx, job)
		n++
	}
}

func (m *ContextManager) shadowWorkerLoop() {
	defer func() {
		m.mu.Lock()
		m.shadowWorkerRunning = false
		m.mu.Unlock()
	}()
	ctx := context.Background()
	for {
		job, ok := m.claimNextShadowJob(ctx)
		if !ok {
			return
		}
		m.runShadowJob(ctx, job)
	}
}

func (m *ContextManager) claimNextShadowJob(ctx context.Context) (SemanticShadowJob, bool) {
	if m.pool != nil {
		if job, ok := m.claimDurableShadowJob(ctx); ok {
			return job, true
		}
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if len(m.shadowJobs) == 0 {
		return SemanticShadowJob{}, false
	}
	job := m.shadowJobs[0]
	m.shadowJobs = m.shadowJobs[1:]
	job.Status = "leased"
	job.Attempt++
	return job, true
}

func (m *ContextManager) claimDurableShadowJob(ctx context.Context) (SemanticShadowJob, bool) {
	owner := fmt.Sprintf("shadow-%d", time.Now().UnixNano())
	tx, err := m.pool.Begin(ctx)
	if err != nil {
		return SemanticShadowJob{}, false
	}
	defer func() { _ = tx.Rollback(ctx) }()

	var (
		job     SemanticShadowJob
		payload []byte
	)
	err = tx.QueryRow(ctx, `
		SELECT id, session_id, summary_revision, input_digest, visibility_digest,
		       watermark_to_message_id, prior_revision, payload, attempt
		FROM chat_summary_shadow_jobs
		WHERE status IN ('pending','leased')
		  AND (lease_until IS NULL OR lease_until < now())
		ORDER BY updated_at, id
		FOR UPDATE SKIP LOCKED
		LIMIT 1
	`).Scan(
		&job.ID, &job.SessionID, &job.SummaryRevision, &job.InputDigest, &job.VisibilityDigest,
		&job.WatermarkToMessageID, &job.PriorRevision, &payload, &job.Attempt,
	)
	if err != nil {
		return SemanticShadowJob{}, false
	}
	job.Attempt++
	var body shadowJobPayload
	_ = json.Unmarshal(payload, &body)
	job.Messages = body.Messages
	job.PriorNarrative = body.PriorNarrative
	job.PriorFacts = body.PriorFacts
	job.Status = "leased"
	_, err = tx.Exec(ctx, `
		UPDATE chat_summary_shadow_jobs
		SET status='leased', attempt=$2, lease_owner=$3, lease_until=now()+$4::interval, updated_at=now()
		WHERE id=$1
	`, job.ID, job.Attempt, owner, shadowLeaseTTL.String())
	if err != nil {
		return SemanticShadowJob{}, false
	}
	if err := tx.Commit(ctx); err != nil {
		return SemanticShadowJob{}, false
	}
	return job, true
}

func (m *ContextManager) runShadowJob(ctx context.Context, job SemanticShadowJob) {
	report := &SemanticShadowReport{
		Status:          ShadowStatusFallback,
		InputDigest:     job.InputDigest,
		ToolPermissions: SemanticToolPermissions,
		GeneratedAt:     time.Now(),
	}
	summarizer := m.semanticSummarizerOrDefault()
	res, err := summarizer.Summarize(ctx, SemanticSummaryRequest{
		SessionID:            job.SessionID,
		Messages:             job.Messages,
		InputDigest:          job.InputDigest,
		VisibilityDigest:     job.VisibilityDigest,
		PriorNarrative:       job.PriorNarrative,
		PriorFacts:           job.PriorFacts,
		WatermarkToMessageID: job.WatermarkToMessageID,
		ExpectedRevision:     job.SummaryRevision,
	})
	if err != nil || res == nil || strings.TrimSpace(res.Narrative) == "" {
		reason := "empty_or_error"
		if err != nil {
			reason = "model_error"
		}
		report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: %s; deterministic retained at %s", reason, time.Now().UTC().Format(time.RFC3339))
		_ = m.persistShadowProvenance(ctx, job, report, err)
		return
	}
	if verr := ValidateShadowResult(job, res); verr != nil {
		report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: validation_failed (%v); deterministic retained at %s", verr, time.Now().UTC().Format(time.RFC3339))
		_ = m.persistShadowProvenance(ctx, job, report, verr)
		return
	}
	if err := m.validateShadowAgainstCurrent(ctx, job); err != nil {
		report.FreshnessNote = fmt.Sprintf("semantic_shadow_fallback: stale (%v); deterministic retained at %s", err, time.Now().UTC().Format(time.RFC3339))
		_ = m.persistShadowProvenance(ctx, job, report, err)
		return
	}
	report.Status = ShadowStatusOK
	report.Narrative = res.Narrative
	report.Model = res.Model
	report.Protocol = res.Protocol
	if report.Protocol == "" {
		report.Protocol = semanticProtocolV1
	}
	// Narrative replacement (if enabled) happens in persistShadowProvenance
	// only after the watermark check against the current served summary.
	_ = m.persistShadowProvenance(ctx, job, report, nil)
}

func (m *ContextManager) validateShadowAgainstCurrent(ctx context.Context, job SemanticShadowJob) error {
	current, err := m.GetChatSummary(ctx, job.SessionID)
	if err != nil {
		return err
	}
	if current == nil {
		return fmt.Errorf("summary missing")
	}
	return ValidateShadowWatermark(job, current)
}

// ValidateShadowWatermark rejects shadow results that no longer match the served summary.
func ValidateShadowWatermark(job SemanticShadowJob, current *ChatSummary) error {
	if current == nil {
		return fmt.Errorf("summary missing")
	}
	if current.Revision != job.SummaryRevision {
		return fmt.Errorf("revision mismatch want %d got %d", job.SummaryRevision, current.Revision)
	}
	if current.VisibilityDigest != "" && job.VisibilityDigest != "" && current.VisibilityDigest != job.VisibilityDigest {
		return fmt.Errorf("visibility digest mismatch")
	}
	if job.WatermarkToMessageID > 0 && current.ToMessageID > 0 && current.ToMessageID != job.WatermarkToMessageID {
		return fmt.Errorf("watermark mismatch want %d got %d", job.WatermarkToMessageID, current.ToMessageID)
	}
	return nil
}

// ValidateShadowResult checks optional echoed digests/watermarks from the summarizer.
func ValidateShadowResult(job SemanticShadowJob, res *SemanticSummaryResult) error {
	if res == nil {
		return fmt.Errorf("nil result")
	}
	if strings.TrimSpace(res.Narrative) == "" {
		return fmt.Errorf("empty narrative")
	}
	if res.InputDigest != "" && res.InputDigest != job.InputDigest {
		return fmt.Errorf("input digest mismatch")
	}
	if res.WatermarkToMessageID > 0 && job.WatermarkToMessageID > 0 && res.WatermarkToMessageID != job.WatermarkToMessageID {
		return fmt.Errorf("result watermark mismatch")
	}
	if res.VisibilityDigest != "" && job.VisibilityDigest != "" && res.VisibilityDigest != job.VisibilityDigest {
		return fmt.Errorf("result visibility digest mismatch")
	}
	return nil
}

func (m *ContextManager) persistShadowProvenance(ctx context.Context, job SemanticShadowJob, report *SemanticShadowReport, runErr error) error {
	current, err := m.GetChatSummary(ctx, job.SessionID)
	if err != nil {
		m.finishShadowJob(ctx, job, runErr)
		return err
	}
	if current == nil {
		m.finishShadowJob(ctx, job, fmt.Errorf("summary missing"))
		return fmt.Errorf("summary missing")
	}
	if err := ValidateShadowWatermark(job, current); err != nil {
		m.finishShadowJob(ctx, job, err)
		return err
	}
	current.Shadow = report
	if report != nil && report.FreshnessNote != "" {
		current.FreshnessNote = report.FreshnessNote
	} else if report != nil && report.Status == ShadowStatusOK {
		current.FreshnessNote = ""
	}
	if report != nil && report.Model != "" {
		current.ModelVersion = report.Model
	}
	if report != nil && report.Protocol != "" {
		current.ProtocolVersion = report.Protocol
	}
	current.Source = "deterministic"
	applySemanticReplace(current, report)
	saveErr := m.SaveChatSummary(ctx, current)
	finishErr := runErr
	if saveErr != nil && finishErr == nil {
		finishErr = saveErr
	}
	m.finishShadowJob(ctx, job, finishErr)
	return saveErr
}

// semanticReplaceMaxRunes bounds the model narrative eligible to serve;
// oversized output is treated as a model anomaly and fails closed.
const semanticReplaceMaxRunes = 4096

// applySemanticReplace swaps the served Narrative with the validated shadow
// narrative under the fail-closed gray rollout: master flag on, session in
// the allowlist, report OK, and narrative within bounds. Any miss keeps the
// deterministic narrative. Facts/counts stay deterministic (declared scope:
// narrative only).
func applySemanticReplace(sum *ChatSummary, report *SemanticShadowReport) {
	if sum == nil || report == nil || report.Status != ShadowStatusOK {
		return
	}
	if !SemanticReplaceEnabledFor(sum.SessionID) {
		return
	}
	n := strings.TrimSpace(report.Narrative)
	if n == "" || utf8.RuneCountInString(n) > semanticReplaceMaxRunes {
		report.FreshnessNote = fmt.Sprintf("semantic_replace_fallback: narrative_out_of_bounds; deterministic retained at %s", time.Now().UTC().Format(time.RFC3339))
		sum.FreshnessNote = report.FreshnessNote
		return
	}
	// Keep the displaced deterministic narrative so the next deterministic
	// regeneration merges clean rule-based material instead of the semantic
	// text (rollback leaves no semantic residue in the narrative chain).
	report.DisplacedNarrative = strings.TrimSpace(sum.Narrative)
	sum.Narrative = n
	sum.Source = "semantic_replace"
	sum.FreshnessNote = fmt.Sprintf("semantic_replace: narrative replaced from %s at %s; deterministic regenerates on next turn after rollback", report.Model, time.Now().UTC().Format(time.RFC3339))
}

func (m *ContextManager) finishShadowJob(ctx context.Context, job SemanticShadowJob, runErr error) {
	if m.pool == nil || job.ID == 0 {
		return
	}
	status := "done"
	errText := ""
	if runErr != nil {
		errText = runErr.Error()
		if job.Attempt >= shadowJobMaxAttempt {
			status = "failed"
		} else {
			status = "pending"
		}
	}
	_, _ = m.pool.Exec(ctx, `
		UPDATE chat_summary_shadow_jobs
		SET status=$2, last_error=$3, lease_owner='', lease_until=NULL, updated_at=now()
		WHERE id=$1
	`, job.ID, status, errText)
}

// PendingShadowJobCount returns in-process queued jobs (tests).
func (m *ContextManager) PendingShadowJobCount() int {
	if m == nil {
		return 0
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.shadowJobs)
}

// WaitShadowIdle waits until the in-process worker is idle (tests).
func (m *ContextManager) WaitShadowIdle(timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		m.mu.RLock()
		running := m.shadowWorkerRunning
		pending := len(m.shadowJobs)
		m.mu.RUnlock()
		if !running && pending == 0 {
			return true
		}
		time.Sleep(5 * time.Millisecond)
	}
	return false
}