package runner

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"slices"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/redaction"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/proposal_generation"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/validators"
)

// RepairAttempt holds optional debug evidence for one semantic request.
type RepairAttempt struct {
	Artifacts     llm.InteractionArtifacts
	Revision      *proposals.EnrichedCorrectionProposal
	Decisions     []ValidatorDecisionRecord
	FailureReason outcomes.ReasonCode
}

// RepairOutcome summarizes the branch and carries its lineage when debug is enabled.
type RepairOutcome struct {
	ProposalIndex int
	AttemptCount  int
	Outcome       string
	Attempts      []RepairAttempt
	Original      proposals.EnrichedCorrectionProposal
	Final         proposals.EnrichedCorrectionProposal
}

func actionable(index int, rejected []ValidatorRejectedChange) bool {
	found := false
	for _, decision := range rejected {
		if decision.ProposalIndex != index {
			continue
		}
		if decision.ReasonCode == validators.ReasonNoEffect {
			return false
		}
		if decision.Disposition != validators.DispositionUnavailable {
			found = true
		}
	}
	return found
}

// repairCandidates branches each rejected proposal from the same section exchange.
func repairCandidates(ctx context.Context, input validateSectionCandidatesInput, references []int, initial validateSectionCandidatesResult, promptID string) (validateSectionCandidatesResult, []RepairOutcome, error) {
	if input.Spec.Settings.LLM.CorrectionRepairAttempts <= 0 {
		return initial, nil, nil
	}
	initialFindings := make(map[int][]ValidatorRejectedChange)
	for _, decision := range initial.rejected {
		initialFindings[decision.ProposalIndex] = append(initialFindings[decision.ProposalIndex], decision)
	}
	superseded := make(map[int]bool)
	var branches []RepairOutcome
	var repairOps []outcomes.Operation
	debug := input.DiagnosticsDir != ""
	section := 0
	if input.SectionIndex != nil {
		section = *input.SectionIndex
	}
	finish := func(err error) (validateSectionCandidatesResult, []RepairOutcome, error) {
		if len(superseded) > 0 {
			initial.filtered = slices.DeleteFunc(initial.filtered, func(item outcomes.FilteredProposal) bool {
				return superseded[item.ProposalIndex]
			})
		}
		initial.operations = append(initial.operations, repairOps...)
		return initial, branches, err
	}

	for position, original := range input.SectionEnriched {
		rejected := initialFindings[original.ProposalIndex]
		if !actionable(original.ProposalIndex, rejected) {
			continue
		}
		branch := RepairOutcome{ProposalIndex: original.ProposalIndex, Outcome: "interrupted"}
		if debug {
			branch.Original, branch.Final = original, original
		}
		if err := ctx.Err(); err != nil {
			branches = append(branches, branch)
			return finish(err)
		}
		if input.Spec.RepairClient == nil || input.Exchange == nil || promptID == "" {
			branches = append(branches, branch)
			return finish(fmt.Errorf("correction repair setup is unavailable"))
		}
		localRef := 0
		if position < len(references) {
			localRef = references[position]
		} else {
			if input.Section == nil {
				input.Section = validators.NewSectionLookup(input.WorkingTranscript)
			}
			if pos, ok := input.Section.Position(original.TargetSegmentID); ok {
				localRef = pos + 1
			}
		}
		current := original
		safeFindings := func(rejected []ValidatorRejectedChange) []proposal_generation.RepairFinding {
			findings := repairFindings(original.ProposalIndex, rejected)
			redactor := redaction.New(repairSecrets(ctx, input))
			for i := range findings {
				findings[i].Explanation = redactor.String(findings[i].Explanation)
			}
			return findings
		}
		findings := safeFindings(rejected)
		replay := &structuredoutput.Replay{Exchange: input.Exchange}
		for attempt := 0; attempt < input.Spec.Settings.LLM.CorrectionRepairAttempts; attempt++ {
			payload, err := (proposal_generation.RepairRequest{ProposalRef: position + 1, Correction: proposal_generation.StructuredCorrectionProposal{SegmentReference: localRef, OriginalText: current.OriginalText, CorrectedText: current.CorrectedText, Confidence: current.Confidence}, Findings: findings}).JSON()
			if err != nil {
				branches = append(branches, branch)
				return finish(err)
			}
			var response proposal_generation.RepairResponse
			proposalIndex, attemptNumber := original.ProposalIndex, attempt+1
			completion, callErr := input.Spec.RepairClient.CompleteStructured(ctx, structuredoutput.Request{PromptID: promptID, Variables: map[string]string{"RepairRequestJSON": payload}, Schema: responseschema.CorrectionRepairKey, Replay: replay, Context: structuredoutput.DiagnosticContext{StageIndex: &input.Spec.Index, StageID: input.Spec.ID, Role: "correction_repair", ModuleKey: input.Spec.ModuleKey, ModuleInstance: input.Spec.InstanceName, SectionIndex: input.SectionIndex, ProposalIndex: &proposalIndex, RepairAttempt: &attemptNumber}}, &response)
			operation := outcomes.Operation{
				Reference: outcomes.Reference{
					StageIndex: &input.Spec.Index, StageID: input.Spec.ID, ModuleKey: input.Spec.ModuleKey,
					ModuleInstance: input.Spec.InstanceName, SectionIndex: section, Kind: outcomes.OperationCorrectionRepair,
					ProposalIndex: &proposalIndex, RepairAttempt: &attemptNumber,
				},
				Counts: completion.Execution, ProposalIndices: []int{proposalIndex},
			}
			finishAttempt := func(err error, revision *proposals.EnrichedCorrectionProposal, validated validateSectionCandidatesResult) {
				branch.AttemptCount++
				repairOps = append(repairOps, operation)
				if !debug {
					return
				}
				record := RepairAttempt{Decisions: validated.decisions, FailureReason: operation.FailureReason}
				if revision != nil {
					snapshot := *revision
					record.Revision = &snapshot
				}
				name := fmt.Sprintf("%s:section-%04d:proposal-%04d:attempt-%04d:repair", input.Spec.InstanceName, section, proposalIndex, attemptNumber)
				var failure any
				if err != nil {
					failure = map[string]string{"error": err.Error()}
				}
				writer := llm.NewDiagnosticsWriter(filepath.Join(input.DiagnosticsDir, input.Spec.InstanceName), repairSecrets(ctx, input))
				artifacts, writeErr := writer.WriteInteraction(name, map[string]any{"stage_index": input.Spec.Index, "stage_id": input.Spec.ID, "proposal_index": proposalIndex, "repair_attempt": attemptNumber, "prompt_id": promptID, "prompt": completion.Prompt, "feedback": findings}, completion.Prepared, map[string]any{"decoded": response, "final_content": completion.FinalContent, "revision": record.Revision, "decisions": record.Decisions, "rejections": validated.rejected}, failure)
				if writeErr == nil {
					record.Artifacts = artifacts
				} else if input.ReportDiagnosticError != nil {
					input.ReportDiagnosticError(fmt.Errorf("repair interaction %s: %w", name, writeErr))
				}
				branch.Attempts = append(branch.Attempts, record)
			}
			if callErr != nil {
				var failure *structuredoutput.Failure
				if errors.As(callErr, &failure) {
					operation.FailureReason = failure.Reason
				} else {
					operation.FailureReason = outcomes.ReasonInternalFailure
				}
				finishAttempt(callErr, nil, validateSectionCandidatesResult{})
				if structuredoutput.IsRecoverable(callErr) {
					branch.Outcome = "unavailable"
					break
				}
				branches = append(branches, branch)
				return finish(callErr)
			}
			decoded, err := proposal_generation.DecodeRepairResponse(completion.Content, position+1)
			if err != nil {
				operation.FailureReason = outcomes.ReasonInvalidOutput
				finishAttempt(err, nil, validateSectionCandidatesResult{})
				branch.Outcome = "unavailable"
				break
			}
			operation.Succeeded = true
			operation.RecoveryReasons = outcomes.ExecutionRecoveries(operation.Counts)
			replay.Turns = append(replay.Turns, completion.RepairInstruction...)
			replay.Turns = append(replay.Turns, structuredoutput.Message{Role: "assistant", Content: completion.FinalContent})
			if decoded.Action == "withdraw" {
				finishAttempt(nil, nil, validateSectionCandidatesResult{})
				branch.Outcome = "withdrawn"
				break
			}
			revision := decoded.Correction
			localRef = revision.SegmentReference
			target := 0
			if input.WorkingTranscript != nil && localRef > 0 && localRef <= len(input.WorkingTranscript.Segments) {
				target = input.WorkingTranscript.Segments[localRef-1].ID
			}
			current = original
			current.CorrectionProposal = proposals.CorrectionProposal{TargetSegmentID: target, OriginalText: revision.OriginalText, CorrectedText: revision.CorrectedText, Confidence: revision.Confidence}
			if debug {
				branch.Final = current
			}
			next := input
			next.RepairProposalIndex, next.RepairAttempt = &proposalIndex, &attemptNumber
			next.SectionEnriched = []proposals.EnrichedCorrectionProposal{current}
			validated, validationErr := validateSectionCandidates(ctx, next)
			if len(validated.filtered) > 0 {
				// Each repair response contains only one proposal.
				operation.Filtered = validated.filtered[:1]
				operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryProposalsFiltered)
			}
			for i := range validated.operations {
				proposalIndex, attemptNumber := original.ProposalIndex, attempt+1
				validated.operations[i].Reference.ProposalIndex = &proposalIndex
				validated.operations[i].Reference.RepairAttempt = &attemptNumber
			}
			approved := validationErr == nil && len(validated.approved) == 1
			if approved {
				operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryCorrectionRevised)
			}
			finishAttempt(validationErr, &current, validated)
			initial.operations = append(initial.operations, validated.operations...)
			if validationErr != nil {
				branches = append(branches, branch)
				return finish(validationErr)
			}
			if approved {
				branch.Outcome = "approved"
				initial.approved = append(initial.approved, current)
				superseded[original.ProposalIndex] = true
				break
			}
			if actionable(original.ProposalIndex, validated.rejected) {
				if attempt+1 == input.Spec.Settings.LLM.CorrectionRepairAttempts {
					branch.Outcome = "exhausted"
					break
				}
				findings = safeFindings(validated.rejected)
				continue
			}
			branch.Outcome = "unavailable"
			for _, d := range validated.rejected {
				if d.Disposition != validators.DispositionUnavailable {
					branch.Outcome = "rejected"
					break
				}
			}
			break
		}
		branches = append(branches, branch)
	}
	return finish(nil)
}

func repairSecrets(ctx context.Context, input validateSectionCandidatesInput) []string {
	var secrets []string
	for _, client := range []contracts.StructuredLLMClient{input.Spec.RepairClient, input.ValidationLLMClient} {
		if source, ok := client.(interface {
			ConfiguredSecrets(context.Context) []string
		}); ok {
			secrets = append(secrets, source.ConfiguredSecrets(ctx)...)
		}
	}
	return secrets
}