package pipeline

import (
	"context"
	"fmt"
	"path"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

type chunkPlanExecution struct {
	accepted    bool
	chunks      []source.Chunk
	plan        *source.ChunkPlan
	diagnostics []contracts.DiagnosticGroup
	rejection   *contracts.RejectedOutput
	lookup      ChunkPlanDecision
	record      *ChunkPlanRecord
	action      string
	summary     artifacts.ChunkPlanSummary
	validation  *artifacts.ValidationSummary
}

type generatedChunkPlanCandidate struct {
	plan                source.ChunkPlan
	chunks              []source.Chunk
	record              ChunkPlanRecord
	producerDiagnostics []contracts.ProducerDiagnostic
	terminal            *attemptTerminalRecorder
}

func effectiveChunkCacheMode(mode ChunkCacheMode) ChunkCacheMode {
	if mode == "" {
		return ChunkCacheBypass
	}
	parsed, _ := ParseChunkCacheMode(string(mode))
	return parsed
}

// Chunk-plan lookup and publication happen serially before lane workers start.
// The runner therefore does not add synchronization around the store.
func (r *Runner) runChunkPlan(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string) (chunkPlanExecution, error) {
	mode := effectiveChunkCacheMode(input.ChunkCacheMode)
	chunker := input.Prepared.chunker
	result := chunkPlanExecution{
		lookup: ChunkPlanDecision{Reason: "chunk plan lookup skipped"},
		summary: artifacts.ChunkPlanSummary{
			Mode: string(mode), SourceDigest: doc.Digest, RequestedModule: chunker.Key(), LookupStatus: "skipped",
			LookupReason: "chunk plan lookup skipped", ValidationStatus: "not_run", PublicationStatus: "not_requested",
		},
	}

	if mode == ChunkCacheAuto {
		record, decision, err := input.ChunkPlans.Load(doc.Digest)
		result.lookup = decision
		result.summary.LookupStatus = chunkPlanLookupStatus(decision.Status)
		result.summary.LookupReason = chunkPlanLookupReason(decision.Status)
		if err != nil {
			result.summary.LookupStatus = "skipped"
			result.summary.LookupReason = chunkPlanLookupReason("")
			return result, fmt.Errorf("load chunk plan: %w", err)
		}
		switch decision.Status {
		case ChunkPlanHit:
			plan, chunks, validationErr := validateAndMaterializeChunkPlan(doc, record.Plan)
			if validationErr == nil {
				report, err := r.validateChunkReport(ctx, doc, chunker.Key(), chunks, sourceInput, sessionID, input.pipeline.ChunkReferences.ReferenceSet, input.Metadata, input.Prepared.chunkValidators, 1, input.Debug)
				if err != nil {
					result.setValidation(nil, err)
					return result, err
				}
				rejection := report.FirstRejection()
				incomplete := firstIncompleteValidation(report)
				if rejection == nil && incomplete == nil {
					if err := result.setCandidate(record, "reused"); err != nil {
						return result, fmt.Errorf("clone reused chunk plan record: %w", err)
					}
					result.plan = &plan
					result.chunks = chunks
					result.accepted = true
					result.setValidation(nil, nil)
					cachedTerminal := producerAttemptTerminal{Action: producerTerminalAccepted, Diagnostics: contracts.CloneProducerDiagnostics(record.Diagnostics), Validation: report}
					diagnostics, diagnosticErr := terminalDiagnosticGroups(cachedTerminal, contracts.DiagnosticOrigin{Stage: contracts.DiagnosticOriginStageChunk, ModuleKey: chunker.Key()}, nil)
					if diagnosticErr != nil {
						return result, fmt.Errorf("promote reused chunk diagnostics: %w", diagnosticErr)
					}
					result.diagnostics = diagnostics
					summary := validationSummary(cachedTerminal, StageChunk, "", "", chunker.Key(), "", 0)
					result.validation = &summary
					return result, nil
				}
				if incomplete != nil && input.pipeline.ChunkValidationPolicy.ValidatorFailure == ValidatorFailureFailRun {
					failure := validatorFailureError(*incomplete)
					result.setValidation(nil, failure)
					failedTerminal := producerAttemptTerminal{Action: producerTerminalFailed, Validation: report, ValidationIncomplete: true}
					summary := validationSummary(failedTerminal, StageChunk, "", "", chunker.Key(), "", 0)
					result.validation = &summary
					return result, failure
				}
				// A cache hit is not model material. Its rejection is discarded and
				// generation begins with the ordinary initial request below. Diagnostics
				// from this discarded candidate are intentionally not promoted.
			}
			result.lookup = ChunkPlanDecision{Status: ChunkPlanInvalid, Reason: chunkPlanLookupReason(ChunkPlanInvalid)}
			result.summary.LookupStatus = "invalid"
			result.summary.LookupReason = chunkPlanLookupReason(ChunkPlanInvalid)
		case ChunkPlanMissing, ChunkPlanInvalid:
			// Generate below.
		default:
			return result, fmt.Errorf("load chunk plan returned unsupported decision status")
		}
	}
	if mode == ChunkCacheAuto || mode == ChunkCacheRefresh {
		result.summary.PublicationStatus = "not_published"
	}

	terminal, err := runProducerAttempts(ctx, producerAttemptConfig{
		Retries:              input.pipeline.Chunk.Retries,
		Policy:               input.pipeline.ChunkValidationPolicy,
		AllowStructuralRetry: input.pipeline.ChunkExecutionClass == contracts.ExecutionClassLLMBacked,
	}, func(attemptCtx context.Context, request producerAttemptRequest) (producerAttemptOutput, error) {
		attempt := request.Number
		attemptStarted := time.Now().UTC()
		attemptPath := path.Join("chunk", fmt.Sprintf("attempt-%02d", attempt))
		attemptCtx, llmScope := withDebugLLMScope(attemptCtx, attemptPath)
		attemptTerminal := newAttemptTerminalRecorder(input.Debug, attemptPath, "chunk", llmScope, debugTimedEnvelope{Stage: string(StageChunk), ModuleKey: chunker.Key(), Attempt: attempt, AttemptKind: string(request.Kind), StartedAt: attemptStarted})
		requestMetadata, metadataErr := cloneMetadata(input.Metadata)
		if metadataErr != nil {
			return producerAttemptOutput{}, attemptTerminal.record(nil, fmt.Errorf("clone chunk request metadata: %w", metadataErr))
		}
		chunkResult, callErr := chunker.Plan(attemptCtx, contracts.ChunkRequest{
			Source: doc, SourceInput: sourceInput.Clone(), SessionID: sessionID,
			References: CloneReferenceSet(input.pipeline.ChunkReferences.ReferenceSet),
			LLMProfile: input.pipeline.Chunk.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(input.pipeline.Chunk.StructuredOutputRepairAttempts), Correction: request.Correction, Metadata: requestMetadata,
		})
		if callErr != nil {
			return producerAttemptOutput{}, attemptTerminal.record(nil, fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), callErr))
		}
		plan, chunks, validationErr := validateAndMaterializeChunkPlan(doc, chunkResult.Plan)
		if validationErr != nil {
			attemptErr := fmt.Errorf("validate chunk plan from chunker %q: %w", chunker.Key(), validationErr)
			payload := map[string]any{"plan": debugChunkPlanEnvelope(chunkResult.Plan)}
			return producerAttemptOutput{}, attemptTerminal.record(payload, fmt.Errorf("%w: %v", contracts.ErrInvalidStructuredOutput, attemptErr))
		}
		planDigest, digestErr := source.DigestChunkPlan(plan)
		if digestErr != nil {
			return producerAttemptOutput{}, attemptTerminal.record(nil, fmt.Errorf("digest generated chunk plan: %w", digestErr))
		}
		producerMetadata, _, metadataErr := moduleManifestMetadata(chunker)
		if metadataErr != nil {
			return producerAttemptOutput{}, attemptTerminal.record(nil, fmt.Errorf("clone chunker manifest metadata: %w", metadataErr))
		}
		profile := ""
		if input.pipeline.ChunkExecutionClass == contracts.ExecutionClassLLMBacked {
			profile = input.pipeline.Chunk.LLMProfile
		}
		candidate := ChunkPlanRecord{
			SchemaVersion: ChunkPlanSchemaVersion, SourceDigest: doc.Digest, PlanDigest: planDigest,
			Plan: source.CloneChunkPlan(plan),
			Producer: ChunkPlanProducer{
				InputModule: input.Prepared.input.Key(), ChunkModule: chunker.Key(), LLMProfile: profile,
				References: append([]artifacts.ReferenceProvenance(nil), referenceTargetProvenance(input.pipeline.ChunkReferences)...),
				Metadata:   producerMetadata,
			},
			CreatedAt: time.Now().UTC(),
		}
		return producerAttemptOutput{Value: generatedChunkPlanCandidate{plan: plan, chunks: chunks, record: candidate, producerDiagnostics: contracts.CloneProducerDiagnostics(chunkResult.Diagnostics), terminal: &attemptTerminal}, Candidate: chunkResult.ModelCandidate, Diagnostics: contracts.CloneProducerDiagnostics(chunkResult.Diagnostics)}, nil
	}, func(validationCtx context.Context, output producerAttemptOutput) (validationReport, error) {
		candidate, ok := output.Value.(generatedChunkPlanCandidate)
		if !ok {
			return validationReport{}, fmt.Errorf("chunk attempt candidate has incompatible type")
		}
		report, validationErr := r.validateChunkReport(validationCtx, doc, chunker.Key(), candidate.chunks, sourceInput, sessionID, input.pipeline.ChunkReferences.ReferenceSet, input.Metadata, input.Prepared.chunkValidators, candidate.terminal.envelope.Attempt, input.Debug)
		payload := map[string]any{
			"plan": debugChunkPlanEnvelope(candidate.plan), "materialized_chunks": debugSourceChunkEnvelopes(candidate.chunks),
			"rejection": debugRejectedOutputPtr(chunkRejection(report, candidate.terminal.envelope.Attempt, chunker.Key())),
		}
		if validationErr != nil {
			return report, candidate.terminal.record(payload, validationErr)
		}
		return report, candidate.terminal.record(payload, nil)
	})
	terminalSummary := validationSummary(terminal, StageChunk, "", "", chunker.Key(), "", 0)
	if debugErr := writeProducerTerminalDebug(input.Debug, "chunk/terminal.json", terminal, input.pipeline.ChunkValidationPolicy, terminalSummary); debugErr != nil {
		return result, debugErr
	}
	if err != nil {
		if result.summary.ValidationStatus == "not_run" {
			result.summary.ValidationStatus = "error"
		}
		return result, err
	}
	diagnostics, diagnosticErr := terminalDiagnosticGroups(terminal, contracts.DiagnosticOrigin{Stage: contracts.DiagnosticOriginStageChunk, ModuleKey: chunker.Key()}, nil)
	if diagnosticErr != nil {
		return result, fmt.Errorf("promote chunk diagnostics: %w", diagnosticErr)
	}
	result.diagnostics = diagnostics
	if terminal.Action == producerTerminalRejected {
		result.rejection = terminal.Rejection
		if result.rejection != nil {
			result.rejection.Stage = string(StageChunk)
			result.rejection.ModuleKey = chunker.Key()
		}
		result.validation = &terminalSummary
		if result.rejection != nil {
			result.rejection.Validation = cloneValidationSummaryPtr(result.validation)
		}
		result.setValidation(result.rejection, nil)
		return result, nil
	}
	candidate, ok := terminal.Value.(generatedChunkPlanCandidate)
	if !ok {
		return result, fmt.Errorf("chunk attempt terminal has incompatible value")
	}
	if err := result.setCandidate(candidate.record, "generated"); err != nil {
		return result, fmt.Errorf("clone generated chunk plan record: %w", err)
	}
	if mode == ChunkCacheRefresh {
		result.action = "refreshed"
		result.summary.Action = "refreshed"
	}
	if mode == ChunkCacheBypass {
		result.action = "bypassed"
		result.summary.Action = "bypassed"
	}
	result.accepted = true
	result.plan = &candidate.plan
	result.chunks = candidate.chunks
	result.validation = &terminalSummary
	result.setValidation(nil, nil)
	if terminal.ValidationIncomplete {
		result.summary.ValidationStatus = "incomplete"
	}

	if (mode == ChunkCacheAuto || mode == ChunkCacheRefresh) && !terminal.ValidationIncomplete {
		record, cloneErr := cloneChunkPlanRecord(*result.record)
		if cloneErr != nil {
			return result, fmt.Errorf("clone chunk plan record for publication: %w", cloneErr)
		}
		record.Diagnostics = contracts.CloneProducerDiagnostics(candidate.producerDiagnostics)
		if err := input.ChunkPlans.Save(record); err != nil {
			return result, fmt.Errorf("save chunk plan: %w", err)
		}
		result.summary.PublicationStatus = "published"
	}
	return result, nil
}

func chunkRejection(report validationReport, attempt int, moduleKey string) *contracts.RejectedOutput {
	rejection := report.FirstRejection()
	if rejection == nil {
		return nil
	}
	return &contracts.RejectedOutput{Stage: string(StageChunk), ModuleKey: moduleKey, ValidatorName: rejection.validatorName, ReasonCode: rejection.reasonCode, Message: rejection.message, AttemptCount: attempt, DiagnosticArtifactPath: rejection.diagnosticPath}
}

func chunkPlanLookupStatus(status ChunkPlanStatus) string {
	switch status {
	case ChunkPlanHit:
		return "hit"
	case ChunkPlanMissing:
		return "missing"
	case ChunkPlanInvalid:
		return "invalid"
	default:
		return "skipped"
	}
}

func chunkPlanLookupReason(status ChunkPlanStatus) string {
	switch status {
	case ChunkPlanHit:
		return "stored chunk plan is valid"
	case ChunkPlanMissing:
		return "chunk plan not found"
	case ChunkPlanInvalid:
		return "stored chunk plan is invalid"
	default:
		return "chunk plan lookup skipped"
	}
}

func (result *chunkPlanExecution) setCandidate(record ChunkPlanRecord, action string) error {
	cloned, err := cloneChunkPlanRecord(record)
	if err != nil {
		return err
	}
	result.record = &cloned
	result.action = action
	result.summary.Action = action
	result.summary.SourceDigest = record.SourceDigest
	result.summary.CandidateDigest = record.PlanDigest
	return nil
}

func (result *chunkPlanExecution) setValidation(rejection *contracts.RejectedOutput, err error) {
	switch {
	case err != nil:
		result.summary.ValidationStatus = "error"
	case rejection != nil:
		result.summary.ValidationStatus = "rejected"
	default:
		result.summary.ValidationStatus = "approved"
	}
}

func cloneChunkPlanRecord(record ChunkPlanRecord) (ChunkPlanRecord, error) {
	record.Plan = source.CloneChunkPlan(record.Plan)
	record.Producer.References = append([]artifacts.ReferenceProvenance(nil), record.Producer.References...)
	metadata, err := cloneMetadata(record.Producer.Metadata)
	if err != nil {
		return ChunkPlanRecord{}, fmt.Errorf("clone chunk plan producer metadata: %w", err)
	}
	record.Producer.Metadata = metadata
	record.Diagnostics = contracts.CloneProducerDiagnostics(record.Diagnostics)
	return record, nil
}

func applyChunkPlanExecution(output *RunOutput, result chunkPlanExecution) error {
	if output == nil {
		return nil
	}
	summary := result.summary
	output.ChunkPlan = &summary
	if output.Manifest.ChunkPlan == nil {
		return nil
	}
	manifest := output.Manifest.ChunkPlan
	manifest.Action = result.action
	if result.record == nil {
		return nil
	}
	record := result.record
	manifest.SourceDigest = record.SourceDigest
	manifest.PlanDigest = record.PlanDigest
	manifest.PlanSchemaVersion = record.SchemaVersion
	manifest.ProducerInputModule = record.Producer.InputModule
	manifest.ProducerModule = record.Producer.ChunkModule
	manifest.ProducerLLMProfile = record.Producer.LLMProfile
	manifest.ProducerReferences = append([]artifacts.ReferenceProvenance(nil), record.Producer.References...)
	metadata, err := cloneMetadata(record.Producer.Metadata)
	if err != nil {
		return fmt.Errorf("clone chunk plan manifest producer metadata: %w", err)
	}
	manifest.ProducerMetadata = metadata
	createdAt := record.CreatedAt
	manifest.CreatedAt = &createdAt
	return nil
}