package runner

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

	"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 retains compact facts for a single semantic request.
type RepairAttempt struct {
	Artifacts     llm.InteractionArtifacts
	Revision      *proposals.EnrichedCorrectionProposal
	Decisions     []ValidatorDecisionRecord
	Rejected      []ValidatorRejectedChange
	Counts        outcomes.ExecutionCounts
	FailureReason outcomes.ReasonCode
	Succeeded     bool
	Approved      bool
	Operations    []outcomes.Operation
	Filtered      []outcomes.FilteredProposal
}

// RepairOutcome keeps the proposal identity through revision and application.
type RepairOutcome struct {
	ProposalIndex int
	Attempts      []RepairAttempt
	Outcome       string
	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
	}
	var branches []RepairOutcome
	for position, original := range input.SectionEnriched {
		if !actionable(original.ProposalIndex, initial.rejected) {
			continue
		}
		branch := RepairOutcome{ProposalIndex: original.ProposalIndex, Original: original, Final: original, Outcome: "interrupted"}
		if err := ctx.Err(); err != nil {
			branches = append(branches, branch)
			return initial, branches, err
		}
		if input.Spec.RepairClient == nil || input.Exchange == nil || promptID == "" {
			branches = append(branches, branch)
			return initial, branches, fmt.Errorf("correction repair setup is unavailable")
		}
		localRef := 0
		if position < len(references) {
			localRef = references[position]
		} else if input.WorkingTranscript != nil {
			for i, s := range input.WorkingTranscript.Segments {
				if s.ID == original.TargetSegmentID {
					localRef = i + 1
					break
				}
			}
		}
		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(initial.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 initial, branches, 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)
			record := RepairAttempt{Counts: completion.Execution}
			save := func(err error) {
				if input.DiagnosticsDir == "" {
					return
				}
				secrets := repairSecrets(ctx, input)
				section := 0
				if input.SectionIndex != nil {
					section = *input.SectionIndex
				}
				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), secrets)
				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": record.Rejected}, failure)
				if writeErr != nil {
					if input.ReportDiagnosticError != nil {
						input.ReportDiagnosticError(fmt.Errorf("repair interaction %s: %w", name, writeErr))
					}
					return
				}
				record.Artifacts = artifacts
			}
			if callErr != nil {
				var failure *structuredoutput.Failure
				if errors.As(callErr, &failure) {
					record.FailureReason = failure.Reason
				} else {
					record.FailureReason = outcomes.ReasonInternalFailure
				}
				save(callErr)
				branch.Attempts = append(branch.Attempts, record)
				if structuredoutput.IsRecoverable(callErr) {
					branch.Outcome = "unavailable"
					break
				}
				branches = append(branches, branch)
				return initial, branches, callErr
			}
			decoded, err := proposal_generation.DecodeRepairResponse(completion.Content, position+1)
			if err != nil {
				record.FailureReason = outcomes.ReasonInvalidOutput
				save(err)
				branch.Attempts = append(branch.Attempts, record)
				branch.Outcome = "unavailable"
				break
			}
			record.Succeeded = true
			replay.Turns = append(replay.Turns, completion.RepairInstruction...)
			replay.Turns = append(replay.Turns, structuredoutput.Message{Role: "assistant", Content: completion.FinalContent})
			if decoded.Action == "withdraw" {
				save(nil)
				branch.Attempts = append(branch.Attempts, record)
				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}
			branch.Final = current
			next := input
			next.RepairProposalIndex, next.RepairAttempt = &proposalIndex, &attemptNumber
			next.SectionEnriched = []proposals.EnrichedCorrectionProposal{current}
			validated, validationErr := validateSectionCandidates(ctx, next)
			revisionSnapshot := current
			record.Revision = &revisionSnapshot
			record.Decisions = validated.decisions
			record.Rejected = validated.rejected
			if len(validated.filtered) > 0 {
				// A repair response contains one proposal. Attribute its first
				// defect once, as for proposals in the original response.
				record.Filtered = validated.filtered[:1]
			}
			for i := range validated.operations {
				proposalIndex, attemptNumber := original.ProposalIndex, attempt+1
				validated.operations[i].Reference.ProposalIndex = &proposalIndex
				validated.operations[i].Reference.RepairAttempt = &attemptNumber
			}
			record.Operations = validated.operations
			save(validationErr)
			branch.Attempts = append(branch.Attempts, record)
			initial.operations = append(initial.operations, validated.operations...)
			if validationErr != nil {
				branches = append(branches, branch)
				return initial, branches, validationErr
			}
			if len(validated.approved) == 1 {
				branch.Attempts[len(branch.Attempts)-1].Approved = true
				branch.Outcome = "approved"
				initial.approved = append(initial.approved, current)
				retained := initial.filtered[:0]
				for _, item := range initial.filtered {
					if item.ProposalIndex != original.ProposalIndex {
						retained = append(retained, item)
					}
				}
				initial.filtered = retained
				break
			}
			findings = safeFindings(validated.rejected)
			if actionable(original.ProposalIndex, validated.rejected) {
				if attempt+1 == input.Spec.Settings.LLM.CorrectionRepairAttempts {
					branch.Outcome = "exhausted"
				}
				continue
			}
			branch.Outcome = "unavailable"
			for _, d := range validated.rejected {
				if d.Disposition != validators.DispositionUnavailable {
					branch.Outcome = "rejected"
					break
				}
			}
			break
		}
		branches = append(branches, branch)
	}
	return initial, branches, nil
}

func repairOperations(input validateSectionCandidatesInput, branches []RepairOutcome) []outcomes.Operation {
	var result []outcomes.Operation
	section := 0
	if input.SectionIndex != nil {
		section = *input.SectionIndex
	}
	for _, branch := range branches {
		for i, attempt := range branch.Attempts {
			index, number := branch.ProposalIndex, i+1
			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: &index, RepairAttempt: &number}, Counts: attempt.Counts, Succeeded: attempt.Succeeded, FailureReason: attempt.FailureReason, ProposalIndices: []int{index}}
			if attempt.Succeeded {
				operation.RecoveryReasons = outcomes.ExecutionRecoveries(attempt.Counts)
				if len(attempt.Filtered) > 0 {
					operation.Filtered = attempt.Filtered
					operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryProposalsFiltered)
				}
				if attempt.Approved {
					operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryCorrectionRevised)
				}
			}
			result = append(result, operation)
		}
	}
	return result
}

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
}