package ai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/jackc/pgx/v5"

	"github.com/missionworkforce/missionmatch/internal/platform/audit"
	"github.com/missionworkforce/missionmatch/internal/platform/inference"
	"github.com/missionworkforce/missionmatch/internal/platform/tenant"
)

// Route is the strategy the router picks for a question. The research recommends choosing the cheapest
// sufficient strategy per query (adaptive RAG), and routing aggregation questions to the database rather
// than to vector RAG (docs/AI_RESEARCH.md, routing + TAG).
type Route string

const (
	RouteRefuse Route = "refuse" // unsafe input (eeo, injection)
	RouteTAG    Route = "tag"    // the database must compute (count, sum, aggregate)
	RouteRAG    Route = "rag"    // retrieve the tenant's candidates and narrate
	RouteDirect Route = "direct" // answer from stable rules in context (CAG), no retrieval
)

// RationaleSource records whether the answer prose was confirmed by the faithfulness gate (ai_faithful) or
// fell back to the deterministic rendering of the evidence (deterministic).
const (
	SourceAIFaithful    = "ai_faithful"
	SourceDeterministic = "deterministic"
)

// ErrNoModel marks that no inference client is configured (or it was not needed). Generation then falls
// back to the deterministic answer; it is never a hang and never a panic.
var ErrNoModel = errors.New("ai: no inference client configured")

// ErrUnknownIntent marks a TAG request for an intent that is not in the curated registry. The handler maps
// it to a 422, never a raw SQL error.
var ErrUnknownIntent = errors.New("ai: unknown query intent")

// Response is an answer plus the provenance the self correction gate and the audit need. The answer is
// ALWAYS grounded: it is the model's prose only when the faithfulness gate confirms it, otherwise the
// deterministic rendering of the same evidence. A refusal (EEO, unanswerable) carries no answer.
type Response struct {
	Answer          string   `json:"answer"`
	Route           Route    `json:"route"`
	RationaleSource string   `json:"rationale_source"` // ai_faithful | deterministic
	Citations       []string `json:"citations"`        // the evidence the answer was grounded on
	Faithful        bool     `json:"faithful"`
	FaithScore      float64  `json:"faithfulness_score"`
	EvidenceCount   int      `json:"evidence_count"`
	Refused         bool     `json:"refused"`
	RefusalReason   string   `json:"refusal_reason,omitempty"`
	FallbackReason  string   `json:"fallback_reason,omitempty"` // why the deterministic fallback was used; NEVER in the audit payload
	MissionScope    string   `json:"mission_scope,omitempty"`   // the mission area retrieval was scoped to (WS6), empty when unscoped
}

// Orchestrator routes a question across the augmented generation strategies, runs the faithfulness gate
// before any model prose ships, and audits every action. It holds the tenant database (RLS scoped) and the
// stateless inference client; it stores nothing itself.
type Orchestrator struct {
	db     *tenant.DB
	client *inference.Client
	aud    *audit.Writer
	mem    *Memory
}

// NewOrchestrator builds the orchestrator. client may be nil (every answer is then deterministic, the
// graceful no plane posture); aud may be nil in tests.
func NewOrchestrator(db *tenant.DB, client *inference.Client, aud *audit.Writer) *Orchestrator {
	return &Orchestrator{db: db, client: client, aud: aud}
}

// WithMemory wires the per tenant memory so RAG answers are memory augmented: the tenant's own recalled
// craft notes join the retrieved candidates as evidence. Memory is RLS scoped, so one tenant's memory can
// never surface in another's answer. Returns the orchestrator for chaining.
func (o *Orchestrator) WithMemory(m *Memory) *Orchestrator {
	o.mem = m
	return o
}

// Answer is the top level dispatcher: a structured intent runs the safe TAG path; otherwise the rule based
// router picks CAG (stable rules), RAG (retrieve and narrate), or a refusal. KAG skill adjacency is a live
// capability reached directly (ai.SkillAdjacency); GraphRAG is scaffolded and is never routed to. A free
// form aggregation question (RouteTAG with no intent) is served best effort by RAG, because the brain never
// authors free form SQL; an exact aggregation over tenant data uses the structured intent path.
func (o *Orchestrator) Answer(ctx context.Context, question, intentName string, args []any) (*Response, error) {
	if strings.TrimSpace(intentName) != "" {
		return o.AnswerTAG(ctx, intentName, args)
	}
	switch Classify(question) {
	case RouteRefuse:
		return o.refuse(ctx, RouteRefuse, "the question references prohibited demographic data"), nil
	case RouteRAG, RouteTAG:
		return o.AnswerRAG(ctx, question, 8)
	default:
		return o.AnswerDirect(ctx, question)
	}
}

// Classify is the lightweight router. A trained classifier is the documented upgrade; rule based routing
// over question shape is competitive at this scale and is transparent (docs/AI_RESEARCH.md, routing).
func Classify(question string) Route {
	q := strings.ToLower(question)
	// Refuse a question that names the eeo schema OR that asks about a protected attribute in plain words.
	if ForbidEEO(q) != nil || MentionsProtectedAttribute(q) {
		return RouteRefuse
	}
	for _, kw := range []string{"how many", "count", "number of", "total", "sum", "average", "percentage", "ratio", "mix", "breakdown"} {
		if strings.Contains(q, kw) {
			return RouteTAG
		}
	}
	for _, kw := range []string{"who", "which candidate", "find", "show me", "list", "top ", "best ", "match", "search"} {
		if strings.Contains(q, kw) {
			return RouteRAG
		}
	}
	return RouteDirect
}

// AnswerRAG retrieves the tenant's candidates for the question, narrates strictly from that evidence, and
// runs the faithfulness gate before returning. The deterministic engine remains the floor: this narrates
// over retrieved rows, it never scores. Retrieval needs the plane to embed the query, so a plane outage
// surfaces as a typed error (through the circuit breaker), never a hang.
func (o *Orchestrator) AnswerRAG(ctx context.Context, question string, limit int) (*Response, error) {
	if ForbidEEO(question) != nil || MentionsProtectedAttribute(question) {
		return o.refuse(ctx, RouteRAG, "the question references prohibited demographic data"), nil
	}
	if o.client == nil {
		return nil, ErrNoModel // retrieval cannot embed the query without the plane
	}
	// Mission scoping (WS6 Option A): if the question names a mission area from the curated mission_spine
	// vocabulary, constrain retrieval to candidates whose full_name, headline, or location text matches it.
	// An unnamed mission is "", which runs the unscoped retrieval byte for byte as before.
	mission := MissionFromQuestion(question)
	hits, err := HybridSearch(ctx, o.db, o.client, question, mission, limit)
	if err != nil {
		return nil, err
	}
	if len(hits) == 0 {
		reason := "no candidates matched the question"
		if mission != "" {
			reason = "no candidates matched the question within the " + mission + " mission scope"
		}
		return o.refuse(ctx, RouteRAG, reason), nil
	}
	evidence := make([]string, len(hits))
	names := make([]string, len(hits))
	for i, h := range hits {
		evidence[i] = fmt.Sprintf("%s (relevance %.4f)", h.FullName, h.Score)
		names[i] = h.FullName
	}
	// Memory augmented generation: the tenant's own recalled craft notes join the evidence (RLS scoped, so
	// never another tenant's). Best effort, never fatal: a recall error or no client just omits the notes.
	if o.mem != nil {
		if notes, mErr := o.mem.Recall(ctx, question, 3); mErr == nil {
			for _, n := range notes {
				evidence = append(evidence, "memory note: "+n)
			}
		}
	}
	deterministic := "Top matching candidates: " + strings.Join(names, ", ") + "."
	model, genErr := o.generate(ctx,
		"Answer the question using ONLY the listed candidates as evidence. Cite candidates by name. If the evidence does not answer the question, say so.",
		"QUESTION: "+question+"\n\nCANDIDATES:\n"+strings.Join(evidence, "\n"))
	resp := o.finalize(ctx, RouteRAG, model, genErr, deterministic, evidence)
	resp.MissionScope = mission
	return resp, nil
}

// AnswerTAG runs a curated, parameterized, read only intent and narrates the database's own result. This
// is the path for aggregation questions; the model never authors SQL. The data is the database's, so a
// plane outage degrades only the prose: the deterministic rendering of the rows is always returned.
func (o *Orchestrator) AnswerTAG(ctx context.Context, intentName string, args []any) (*Response, error) {
	intent, ok := IntentByName(intentName)
	if !ok {
		return nil, fmt.Errorf("%w: %q", ErrUnknownIntent, intentName)
	}
	res, err := ExecuteIntent(ctx, o.db, intent, args)
	if err != nil {
		return nil, err
	}
	rendered, _ := json.Marshal(map[string]any{"columns": res.Columns, "rows": res.Rows})
	evidence := []string{string(rendered)}
	deterministic := renderTable(intent.Description, res.Columns, res.Rows)
	model, genErr := o.generate(ctx,
		"State the result plainly from the table data. Do not invent numbers; report exactly what the rows say.",
		"INTENT: "+intent.Description+"\n\nRESULT:\n"+string(rendered))
	return o.finalize(ctx, RouteTAG, model, genErr, deterministic, evidence), nil
}

// AnswerDirect answers from the stable public rules in context (CAG): no retrieval, no tenant data. The
// rules are contents free, so this never touches tenant rows; a plane outage returns the rules themselves.
func (o *Orchestrator) AnswerDirect(ctx context.Context, question string) (*Response, error) {
	if ForbidEEO(question) != nil || MentionsProtectedAttribute(question) {
		return o.refuse(ctx, RouteDirect, "the question references prohibited demographic data"), nil
	}
	rules := PublicRules()
	evidence := []string{rules}
	model, genErr := o.generate(ctx,
		"Answer using ONLY these platform rules as evidence. If the rules do not answer the question, say so. Do not invent a number or a weight.",
		"QUESTION: "+question+"\n\nRULES:\n"+rules)
	return o.finalize(ctx, RouteDirect, model, genErr, rules, evidence), nil
}

// generate calls the plane (the injection guard is wrapped on the gateway side). A nil client is not an
// error to the caller: it returns ErrNoModel so finalize takes the deterministic path.
func (o *Orchestrator) generate(ctx context.Context, system, prompt string) (string, error) {
	if o.client == nil {
		return "", ErrNoModel
	}
	out, err := o.client.Generate(ctx, inference.GenerateRequest{System: system, Prompt: prompt, MaxTokens: 400, Temperature: 0.1})
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(out.Text), nil
}

// finalize is the self correction step. It returns the model prose ONLY when the faithfulness gate confirms
// it is supported by the evidence; otherwise it returns the deterministic rendering of the same evidence
// (a 200, rationale_source deterministic), recording why. The audit records the route, the rationale source,
// and the evidence count, never the answer text, the question, the fallback reason, or any candidate.
func (o *Orchestrator) finalize(ctx context.Context, route Route, modelAnswer string, genErr error, deterministic string, citations []string) *Response {
	resp := &Response{Route: route, Citations: citations, EvidenceCount: len(citations)}
	if genErr != nil || strings.TrimSpace(modelAnswer) == "" {
		resp.Answer = deterministic
		resp.RationaleSource = SourceDeterministic
		resp.FallbackReason = fallbackReason(genErr)
		o.audit(ctx, "aedt.ai.answered", map[string]any{"route": route, "rationale_source": SourceDeterministic, "evidence": len(citations)})
		return resp
	}
	verdict, err := o.client.CheckFaithfulness(ctx, modelAnswer, citations)
	if err == nil && verdict.Supported {
		resp.Answer = modelAnswer
		resp.RationaleSource = SourceAIFaithful
		resp.Faithful = true
		resp.FaithScore = verdict.Score
	} else {
		resp.Answer = deterministic
		resp.RationaleSource = SourceDeterministic
		if err != nil {
			resp.FallbackReason = fallbackReason(err)
		} else {
			resp.FaithScore = verdict.Score
			resp.FallbackReason = "could_not_confirm"
		}
	}
	o.audit(ctx, "aedt.ai.answered", map[string]any{"route": route, "rationale_source": resp.RationaleSource, "faithful": resp.Faithful, "evidence": len(citations)})
	return resp
}

func fallbackReason(err error) string {
	switch {
	case err == nil:
		return "could_not_confirm"
	case errors.Is(err, ErrNoModel):
		return "no_model"
	case errors.Is(err, inference.ErrCircuitOpen):
		return "circuit_open"
	case errors.Is(err, context.DeadlineExceeded):
		return "plane_timeout"
	default:
		return "plane_error"
	}
}

// renderTable builds a deterministic, readable answer from a TAG result, so the aggregation question is
// answered exactly from the database even when the model prose is withheld.
func renderTable(description string, columns []string, rows [][]any) string {
	if len(rows) == 0 {
		return description + ": no rows."
	}
	var b strings.Builder
	b.WriteString(description + ": ")
	parts := make([]string, 0, len(rows))
	for _, row := range rows {
		cells := make([]string, 0, len(row))
		for i, v := range row {
			col := ""
			if i < len(columns) {
				col = columns[i] + "="
			}
			cells = append(cells, fmt.Sprintf("%s%v", col, v))
		}
		parts = append(parts, strings.Join(cells, " "))
	}
	b.WriteString(strings.Join(parts, "; "))
	b.WriteString(".")
	return b.String()
}

func (o *Orchestrator) refuse(ctx context.Context, route Route, reason string) *Response {
	o.audit(ctx, "aedt.ai.refused", map[string]any{"route": route, "reason": reason})
	return &Response{Route: route, Refused: true, RefusalReason: reason}
}

// audit writes one append only AEDT event for an AI action. Payloads carry the route and verdict only,
// never the answer text, the question, the fallback reason, or any candidate, so the audit holds no PII.
func (o *Orchestrator) audit(ctx context.Context, action string, payload map[string]any) {
	if o.aud == nil || o.db == nil {
		return
	}
	body, _ := json.Marshal(payload)
	_ = o.db.RunInTenantTx(ctx, func(ctx context.Context, tx pgx.Tx) error {
		return o.aud.OnTx(ctx, tx, audit.Event{Action: action, EntityType: "ai", NewValues: body})
	})
}