package runner

import (
	"context"
	"fmt"
	"path/filepath"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/validators"
)

const (
	ModuleStatusSuccess = "success"
	ModuleStatusFailed  = "failed"
)

// ModuleFactory resolves one module instance for one run spec.
type ModuleFactory interface {
	ModuleForSpec(spec contracts.ModuleRunSpec) (contracts.TranscriptModule, error)
}

// Runner executes module instances sequentially against a working transcript.
type Runner struct {
	factory ModuleFactory
}

type ValidationScheduler = contracts.LLMScheduler

// ModuleResult captures deterministic per-module execution output.
type ModuleResult struct {
	ModuleKey          string                      `json:"module_key"`
	ModuleInstance     string                      `json:"module_instance"`
	ReplacementPolicy  proposals.ReplacementPolicy `json:"replacement_policy"`
	Status             string                      `json:"status"`
	ProposalCount      int                         `json:"proposal_count"`
	ValidatorDecisions []ValidatorDecisionRecord   `json:"validator_decisions,omitempty"`
	ValidatorRejected  []ValidatorRejectedChange   `json:"validator_rejected,omitempty"`
	AppliedChanges     []proposals.AppliedChange   `json:"applied_changes,omitempty"`
	SkippedChanges     []proposals.SkippedChange   `json:"skipped_changes,omitempty"`
	ErrorMessage       string                      `json:"error_message,omitempty"`
	StartedAt          time.Time                   `json:"started_at"`
	CompletedAt        time.Time                   `json:"completed_at"`
}

type ValidatorDecisionRecord struct {
	ValidatorName          string `json:"validator_name"`
	ProposalIndex          int    `json:"proposal_index"`
	Approved               bool   `json:"approved"`
	ReasonCode             string `json:"reason_code"`
	Message                string `json:"message"`
	DiagnosticArtifactPath string `json:"diagnostic_artifact_path,omitempty"`
}

type ValidatorRejectedChange struct {
	ValidatorName   string `json:"validator_name"`
	ProposalIndex   int    `json:"proposal_index"`
	ModuleKey       string `json:"module_key"`
	ModuleInstance  string `json:"module_instance"`
	TargetSegmentID int    `json:"target_segment_id"`
	OriginalText    string `json:"original_text"`
	CorrectedText   string `json:"corrected_text"`
	ReasonCode      string `json:"reason_code"`
	Message         string `json:"message"`
}

// RunInput is the deterministic runner input.
type RunInput struct {
	Config                   *config.Config
	Transcript               *schema.Transcript
	Glossary                 *schema.Glossary
	ModuleSpecs              []contracts.ModuleRunSpec
	ProposalLLMClient        contracts.StructuredLLMClient
	ProposalLLMScheduler     contracts.LLMScheduler
	ProposalDiagnosticsDir   string
	ValidationLLMClient      contracts.StructuredLLMClient
	ValidationLLMScheduler   ValidationScheduler
	ValidationDiagnosticsDir string
}

// RunOutput is the deterministic runner output.
type RunOutput struct {
	FinalTranscript *schema.Transcript `json:"-"`
	ModuleResults   []ModuleResult     `json:"module_results"`
}

func New(factory ModuleFactory) *Runner {
	return &Runner{factory: factory}
}

func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
	if r == nil || r.factory == nil {
		return RunOutput{}, fmt.Errorf("runner module factory is required")
	}

	working := cloneTranscript(input.Transcript)
	results := make([]ModuleResult, 0, len(input.ModuleSpecs))

	for _, spec := range input.ModuleSpecs {
		startedAt := time.Now().UTC()
		module, err := r.factory.ModuleForSpec(spec)
		if err != nil {
			failed := ModuleResult{
				ModuleKey:      spec.ModuleKey,
				ModuleInstance: spec.InstanceName,
				Status:         ModuleStatusFailed,
				ErrorMessage:   err.Error(),
				StartedAt:      startedAt,
				CompletedAt:    time.Now().UTC(),
			}
			results = append(results, failed)
			return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q setup failed: %w", spec.InstanceName, err)
		}

		policy := module.ReplacementPolicy()
		sections, err := chunkWorkingTranscript(input.Config, working)
		if err != nil {
			failed := ModuleResult{
				ModuleKey:         spec.ModuleKey,
				ModuleInstance:    spec.InstanceName,
				ReplacementPolicy: policy,
				Status:            ModuleStatusFailed,
				ErrorMessage:      err.Error(),
				StartedAt:         startedAt,
				CompletedAt:       time.Now().UTC(),
			}
			results = append(results, failed)
			return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q chunking failed: %w", spec.InstanceName, err)
		}

		enriched, proposeErr := collectSectionProposals(ctx, collectSectionProposalsInput{
			Module:            module,
			Spec:              spec,
			Policy:            policy,
			Config:            input.Config,
			Glossary:          input.Glossary,
			Sections:          sections,
			ProposalClient:    input.ProposalLLMClient,
			ProposalScheduler: input.ProposalLLMScheduler,
			DiagnosticsDir:    input.ProposalDiagnosticsDir,
		})
		if proposeErr != nil {
			failed := ModuleResult{
				ModuleKey:         spec.ModuleKey,
				ModuleInstance:    spec.InstanceName,
				ReplacementPolicy: policy,
				Status:            ModuleStatusFailed,
				ErrorMessage:      proposeErr.Error(),
				StartedAt:         startedAt,
				CompletedAt:       time.Now().UTC(),
			}
			results = append(results, failed)
			return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, proposeErr)
		}

		validatorDecisions := make([]ValidatorDecisionRecord, 0)
		validatorRejected := make([]ValidatorRejectedChange, 0)
		eligible := enriched
		for _, validator := range module.Validators() {
			var diagnosticsWriter validators.InteractionDiagnosticsWriter
			if input.ValidationDiagnosticsDir != "" {
				diagnosticsWriter = &llmDiagnosticsWriterAdapter{
					writer: llm.NewDiagnosticsWriter(
						filepath.Join(input.ValidationDiagnosticsDir, spec.InstanceName),
						validatorSecrets(input.Config),
					),
				}
			}

			vResult, vErr := validator.Validate(ctx, contracts.ValidationRequest{
				WorkingTranscript: working,
				CandidateProposal: eligible,
				ModuleKey:         spec.ModuleKey,
				ModuleInstance:    spec.InstanceName,
				ReplacementPolicy: policy,
				Glossary:          input.Glossary,
				Config:            input.Config,
				LLMClient:         validationLLMClientAdapter{client: input.ValidationLLMClient},
				Scheduler:         input.ValidationLLMScheduler,
				DiagnosticsWriter: diagnosticsWriter,
			})
			if vErr != nil {
				failed := ModuleResult{
					ModuleKey:          spec.ModuleKey,
					ModuleInstance:     spec.InstanceName,
					ReplacementPolicy:  policy,
					Status:             ModuleStatusFailed,
					ProposalCount:      len(enriched),
					ValidatorDecisions: validatorDecisions,
					ValidatorRejected:  validatorRejected,
					ErrorMessage:       vErr.Error(),
					StartedAt:          startedAt,
					CompletedAt:        time.Now().UTC(),
				}
				results = append(results, failed)
				return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q validator %q failed: %w", spec.InstanceName, validator.Name(), vErr)
			}
			if err := validators.EnforceDecisionCardinality(eligible, vResult.Decisions); err != nil {
				failed := ModuleResult{
					ModuleKey:          spec.ModuleKey,
					ModuleInstance:     spec.InstanceName,
					ReplacementPolicy:  policy,
					Status:             ModuleStatusFailed,
					ProposalCount:      len(enriched),
					ValidatorDecisions: validatorDecisions,
					ValidatorRejected:  validatorRejected,
					ErrorMessage:       err.Error(),
					StartedAt:          startedAt,
					CompletedAt:        time.Now().UTC(),
				}
				results = append(results, failed)
				return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q validator %q cardinality failed: %w", spec.InstanceName, validator.Name(), err)
			}

			nextEligible := make([]proposals.EnrichedCorrectionProposal, 0, len(eligible))
			byIndex := make(map[int]proposals.EnrichedCorrectionProposal, len(eligible))
			for _, p := range eligible {
				byIndex[p.ProposalIndex] = p
			}
			for _, d := range vResult.Decisions {
				validatorDecisions = append(validatorDecisions, ValidatorDecisionRecord{
					ValidatorName:          validator.Name(),
					ProposalIndex:          d.ProposalIndex,
					Approved:               d.Approved,
					ReasonCode:             d.ReasonCode,
					Message:                d.Message,
					DiagnosticArtifactPath: d.DiagnosticArtifactPath,
				})
				if d.Approved {
					nextEligible = append(nextEligible, byIndex[d.ProposalIndex])
					continue
				}
				p := byIndex[d.ProposalIndex]
				validatorRejected = append(validatorRejected, ValidatorRejectedChange{
					ValidatorName:   validator.Name(),
					ProposalIndex:   p.ProposalIndex,
					ModuleKey:       p.ModuleKey,
					ModuleInstance:  p.ModuleInstance,
					TargetSegmentID: p.TargetSegmentID,
					OriginalText:    p.OriginalText,
					CorrectedText:   p.CorrectedText,
					ReasonCode:      d.ReasonCode,
					Message:         d.Message,
				})
			}
			eligible = nextEligible
		}

		applyResult := proposals.ApplyProposals(working, eligible, policy)
		working = applyResult.Transcript

		results = append(results, ModuleResult{
			ModuleKey:          spec.ModuleKey,
			ModuleInstance:     spec.InstanceName,
			ReplacementPolicy:  policy,
			Status:             ModuleStatusSuccess,
			ProposalCount:      len(enriched),
			ValidatorDecisions: validatorDecisions,
			ValidatorRejected:  validatorRejected,
			AppliedChanges:     applyResult.Applied,
			SkippedChanges:     applyResult.Skipped,
			StartedAt:          startedAt,
			CompletedAt:        time.Now().UTC(),
		})
	}

	return RunOutput{FinalTranscript: working, ModuleResults: results}, nil
}

type collectSectionProposalsInput struct {
	Module            contracts.TranscriptModule
	Spec              contracts.ModuleRunSpec
	Policy            proposals.ReplacementPolicy
	Config            *config.Config
	Glossary          *schema.Glossary
	Sections          []chunking.Section
	ProposalClient    contracts.StructuredLLMClient
	ProposalScheduler contracts.LLMScheduler
	DiagnosticsDir    string
}

type sectionProposals struct {
	meta      contracts.SectionMetadata
	corrected []proposals.CorrectionProposal
}

func collectSectionProposals(ctx context.Context, input collectSectionProposalsInput) ([]proposals.EnrichedCorrectionProposal, error) {
	if len(input.Sections) == 0 {
		return []proposals.EnrichedCorrectionProposal{}, nil
	}

	maxWorkers := 1
	if input.Config != nil && input.Config.PrimaryLLM.Concurrency > 1 {
		maxWorkers = input.Config.PrimaryLLM.Concurrency
	}
	if maxWorkers > len(input.Sections) {
		maxWorkers = len(input.Sections)
	}

	sectionResults := make([]sectionProposals, len(input.Sections))
	runCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	sem := make(chan struct{}, maxWorkers)
	var (
		wg       sync.WaitGroup
		errOnce  sync.Once
		firstErr error
	)

	for sectionPos, section := range input.Sections {
		sectionPos := sectionPos
		section := section
		wg.Add(1)
		go func() {
			defer wg.Done()
			select {
			case sem <- struct{}{}:
			case <-runCtx.Done():
				return
			}
			defer func() { <-sem }()

			meta := contracts.SectionMetadataFromSection(section)
			corrected, err := input.Module.Propose(runCtx, contracts.ProposalRequest{
				ExecutionContext: contracts.ExecutionContext{
					Config:            input.Config,
					WorkingTranscript: transcriptFromSection(section),
					Glossary:          input.Glossary,
					Section:           &meta,
					DiagnosticsDir:    input.DiagnosticsDir,
				},
				RunSpec: contracts.ModuleRunSpec{
					ModuleKey:         input.Spec.ModuleKey,
					InstanceName:      input.Spec.InstanceName,
					ReplacementPolicy: input.Policy,
				},
				LLMClient:    input.ProposalClient,
				LLMScheduler: input.ProposalScheduler,
			})
			if err != nil {
				errOnce.Do(func() {
					firstErr = err
					cancel()
				})
				return
			}

			sectionResults[sectionPos] = sectionProposals{
				meta:      meta,
				corrected: corrected,
			}
		}()
	}

	wg.Wait()
	if firstErr != nil {
		return nil, firstErr
	}

	enriched := make([]proposals.EnrichedCorrectionProposal, 0)
	nextProposalIndex := 0
	for _, sectionResult := range sectionResults {
		for _, corrected := range sectionResult.corrected {
			sectionIndex := sectionResult.meta.Index
			enriched = append(enriched, proposals.EnrichedCorrectionProposal{
				CorrectionProposal: corrected,
				ProposalMetadata: proposals.ProposalMetadata{
					ProposalIndex:  nextProposalIndex,
					ModuleKey:      input.Spec.ModuleKey,
					ModuleInstance: input.Spec.InstanceName,
					SectionIndex:   &sectionIndex,
				},
			})
			nextProposalIndex++
		}
	}
	return enriched, nil
}

func chunkWorkingTranscript(cfg *config.Config, transcript *schema.Transcript) ([]chunking.Section, error) {
	if cfg == nil {
		if transcript == nil || len(transcript.Segments) == 0 {
			return []chunking.Section{}, nil
		}
		return []chunking.Section{
			{
				Index:           0,
				EstimatedTokens: estimateTranscriptTokens(transcript),
				Segments:        append([]schema.Segment(nil), transcript.Segments...),
				StartSegmentID:  transcript.Segments[0].ID,
				EndSegmentID:    transcript.Segments[len(transcript.Segments)-1].ID,
			},
		}, nil
	}
	chunker := chunking.NewChunker(chunking.ChunkingConfig{
		MaxSectionTokens: cfg.MaxSectionTokens,
		MinSectionTokens: cfg.MinSectionTokens,
		TargetSections:   cfg.TargetSections,
	})
	return chunker.ChunkTranscript(transcript)
}

func estimateTranscriptTokens(transcript *schema.Transcript) int {
	if transcript == nil || len(transcript.Segments) == 0 {
		return 0
	}
	estimator := chunking.NewSimpleTokenEstimator()
	total := 0
	for _, segment := range transcript.Segments {
		total += estimator.EstimateTokens(segment.Text)
	}
	return total
}

func transcriptFromSection(section chunking.Section) *schema.Transcript {
	segments := make([]schema.Segment, len(section.Segments))
	copy(segments, section.Segments)
	return &schema.Transcript{Segments: segments}
}

func cloneTranscript(t *schema.Transcript) *schema.Transcript {
	if t == nil {
		return &schema.Transcript{}
	}
	segments := make([]schema.Segment, len(t.Segments))
	for i, s := range t.Segments {
		var categories []string
		if s.Categories != nil {
			categories = append([]string(nil), s.Categories...)
		}
		segments[i] = schema.Segment{
			ID:         s.ID,
			Speaker:    s.Speaker,
			Start:      s.Start,
			End:        s.End,
			Text:       s.Text,
			Categories: categories,
		}
	}
	return &schema.Transcript{Segments: segments}
}

type validationLLMClientAdapter struct {
	client contracts.StructuredLLMClient
}

func (a validationLLMClientAdapter) CompleteStructured(ctx context.Context, req validators.StructuredCompletionRequest, out any) (validators.StructuredCompletionResponse, error) {
	if a.client == nil {
		return validators.StructuredCompletionResponse{}, fmt.Errorf("validation structured LLM client is not configured")
	}
	messages := make([]contracts.LLMMessage, len(req.Messages))
	for i, m := range req.Messages {
		messages[i] = contracts.LLMMessage{Role: m.Role, Content: m.Content}
	}
	resp, err := a.client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName: req.StageName,
		Messages:  messages,
		Model:     req.Model,
	}, out)
	if err != nil {
		return validators.StructuredCompletionResponse{}, err
	}
	return validators.StructuredCompletionResponse{
		PromptTokens:     resp.PromptTokens,
		CompletionTokens: resp.CompletionTokens,
		TotalTokens:      resp.TotalTokens,
	}, nil
}

type llmDiagnosticsWriterAdapter struct {
	writer *llm.DiagnosticsWriter
}

func (a *llmDiagnosticsWriterAdapter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (validators.InteractionArtifacts, error) {
	if a == nil || a.writer == nil {
		return validators.InteractionArtifacts{}, fmt.Errorf("diagnostics writer is not configured")
	}
	art, err := a.writer.WriteInteraction(stage, requestMetadata, requestPayload, responsePayload, errorPayload)
	if err != nil {
		return validators.InteractionArtifacts{}, err
	}
	return validators.InteractionArtifacts{
		RequestMetadataPath: art.RequestMetadataPath,
		RequestPayloadPath:  art.RequestPayloadPath,
		ResponsePayloadPath: art.ResponsePayloadPath,
		ErrorPayloadPath:    art.ErrorPayloadPath,
	}, nil
}

func validatorSecrets(cfg *config.Config) []string {
	if cfg == nil {
		return nil
	}
	effective := cfg.EffectiveValidationLLMConfig()
	return []string{
		cfg.PrimaryLLM.APIKey,
		effective.APIKey,
		cfg.ValidationLLM.APIKey,
	}
}