package runner

import (
	"context"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	validatormetadata "gitea.maximumdirect.net/eric/audita/internal/validators/metadata"
)

type UtilizationDiagnostics struct {
	EffectiveConcurrency EffectiveConcurrencyLimits `json:"effective_concurrency"`
	RunTiming            RunTimingSummary           `json:"run_timing"`
	LLMCalls             LLMCallSummary             `json:"llm_calls"`
	Modules              []ModuleTimingSummary      `json:"modules,omitempty"`
	Validators           []ValidatorTimingSummary   `json:"validators,omitempty"`
}

type EffectiveConcurrencyLimits struct {
	TotalLLM      int `json:"total_llm"`
	ProposalLLM   int `json:"proposal_llm"`
	ValidationLLM int `json:"validation_llm"`
}

type RunTimingSummary struct {
	RunWallTimeMS             int64 `json:"run_wall_time_ms"`
	SchedulerQueueWaitMS      int64 `json:"scheduler_queue_wait_ms"`
	LLMExecutionTimeMS        int64 `json:"llm_execution_time_ms"`
	DeterministicValidationMS int64 `json:"deterministic_validation_time_ms"`
	MaxInFlightLLMCalls       int   `json:"max_in_flight_llm_calls"`
	AverageInFlightLLMCalls   int   `json:"average_in_flight_llm_calls"`
}

type LLMCallSummary struct {
	TotalProposalCalls   int `json:"total_proposal_calls"`
	TotalValidationCalls int `json:"total_validation_calls"`
}

type ModuleTimingSummary struct {
	ModuleKey                 string `json:"module_key"`
	ModuleInstance            string `json:"module_instance"`
	ProposalLLMCalls          int    `json:"proposal_llm_calls"`
	ValidationLLMCalls        int    `json:"validation_llm_calls"`
	SchedulerQueueWaitMS      int64  `json:"scheduler_queue_wait_ms"`
	LLMExecutionTimeMS        int64  `json:"llm_execution_time_ms"`
	DeterministicValidationMS int64  `json:"deterministic_validation_time_ms"`
	ModuleWallTimeMS          int64  `json:"module_wall_time_ms"`
}

type ValidatorTimingSummary struct {
	ModuleKey      string `json:"module_key"`
	ModuleInstance string `json:"module_instance"`
	ValidatorKey   string `json:"validator_key"`
	LLMBacked      bool   `json:"llm_backed"`
	Calls          int    `json:"calls"`
	ElapsedMS      int64  `json:"elapsed_ms"`
}

type runInstrumentation struct {
	startedAt time.Time

	mu                               sync.Mutex
	totalProposalCalls               int
	totalValidationCalls             int
	schedulerQueueWait               time.Duration
	llmExecutionTime                 time.Duration
	deterministicValidation          time.Duration
	perModuleProposalCalls           map[string]int
	perModuleValidationCalls         map[string]int
	perModuleSchedulerQueueWait      map[string]time.Duration
	perModuleLLMExecutionTime        map[string]time.Duration
	perModuleDeterministicValidation map[string]time.Duration
	validatorSummaries               map[string]*ValidatorTimingSummary

	inflight                int
	maxInflight             int
	inflightArea            float64
	lastInflightChange      time.Time
	averageInflightComputed bool
}

func newRunInstrumentation(startedAt time.Time) *runInstrumentation {
	return &runInstrumentation{
		startedAt:                        startedAt,
		perModuleProposalCalls:           make(map[string]int),
		perModuleValidationCalls:         make(map[string]int),
		perModuleSchedulerQueueWait:      make(map[string]time.Duration),
		perModuleLLMExecutionTime:        make(map[string]time.Duration),
		perModuleDeterministicValidation: make(map[string]time.Duration),
		validatorSummaries:               make(map[string]*ValidatorTimingSummary),
		lastInflightChange:               startedAt,
	}
}

func (i *runInstrumentation) wrapProposalClient(client contracts.StructuredLLMClient) contracts.StructuredLLMClient {
	if client == nil {
		return nil
	}
	return measuredClient{
		inner:      client,
		collector:  i,
		isProposal: true,
	}
}

func (i *runInstrumentation) wrapValidationClient(client contracts.StructuredLLMClient) contracts.StructuredLLMClient {
	if client == nil {
		return nil
	}
	return measuredClient{
		inner:      client,
		collector:  i,
		isProposal: false,
	}
}

func (i *runInstrumentation) wrapScheduler(s contracts.LLMScheduler, isProposal bool) contracts.LLMScheduler {
	if s == nil {
		return nil
	}
	return measuredScheduler{
		inner:      s,
		collector:  i,
		isProposal: isProposal,
	}
}

func (i *runInstrumentation) recordValidatorTiming(moduleKey, moduleInstance, validatorKey string, llmBacked bool, elapsed time.Duration) {
	i.mu.Lock()
	defer i.mu.Unlock()
	if !llmBacked {
		i.deterministicValidation += elapsed
		i.perModuleDeterministicValidation[moduleInstance] += elapsed
	}
	k := moduleInstance + "::" + validatorKey
	summary, ok := i.validatorSummaries[k]
	if !ok {
		summary = &ValidatorTimingSummary{
			ModuleKey:      moduleKey,
			ModuleInstance: moduleInstance,
			ValidatorKey:   validatorKey,
			LLMBacked:      llmBacked,
		}
		i.validatorSummaries[k] = summary
	}
	summary.Calls++
	summary.ElapsedMS += elapsed.Milliseconds()
}

func (i *runInstrumentation) recordModuleWallTime(moduleInstance string, elapsed time.Duration) {
	i.mu.Lock()
	defer i.mu.Unlock()
	// Derived later from started/completed, but we keep a slot for sanity.
	_ = moduleInstance
	_ = elapsed
}

func (i *runInstrumentation) beginInFlight() {
	now := time.Now().UTC()
	i.mu.Lock()
	defer i.mu.Unlock()
	i.inflightArea += float64(i.inflight) * now.Sub(i.lastInflightChange).Seconds()
	i.lastInflightChange = now
	i.inflight++
	if i.inflight > i.maxInflight {
		i.maxInflight = i.inflight
	}
}

func (i *runInstrumentation) endInFlight() {
	now := time.Now().UTC()
	i.mu.Lock()
	defer i.mu.Unlock()
	i.inflightArea += float64(i.inflight) * now.Sub(i.lastInflightChange).Seconds()
	i.lastInflightChange = now
	if i.inflight > 0 {
		i.inflight--
	}
}

func (i *runInstrumentation) finalize(runOutput *RunOutput, completedAt time.Time, limits EffectiveConcurrencyLimits) *UtilizationDiagnostics {
	i.mu.Lock()
	defer i.mu.Unlock()

	if !i.averageInflightComputed {
		i.inflightArea += float64(i.inflight) * completedAt.Sub(i.lastInflightChange).Seconds()
		i.lastInflightChange = completedAt
		i.averageInflightComputed = true
	}

	runSeconds := completedAt.Sub(i.startedAt).Seconds()
	avgInflight := 0
	if runSeconds > 0 {
		avgInflight = int(i.inflightArea / runSeconds)
	}

	moduleSummaries := make([]ModuleTimingSummary, 0)
	if runOutput != nil {
		moduleSummaries = make([]ModuleTimingSummary, 0, len(runOutput.ModuleResults))
		for _, moduleResult := range runOutput.ModuleResults {
			moduleSummaries = append(moduleSummaries, ModuleTimingSummary{
				ModuleKey:                 moduleResult.ModuleKey,
				ModuleInstance:            moduleResult.ModuleInstance,
				ProposalLLMCalls:          i.perModuleProposalCalls[moduleResult.ModuleInstance],
				ValidationLLMCalls:        i.perModuleValidationCalls[moduleResult.ModuleInstance],
				SchedulerQueueWaitMS:      i.perModuleSchedulerQueueWait[moduleResult.ModuleInstance].Milliseconds(),
				LLMExecutionTimeMS:        i.perModuleLLMExecutionTime[moduleResult.ModuleInstance].Milliseconds(),
				DeterministicValidationMS: i.perModuleDeterministicValidation[moduleResult.ModuleInstance].Milliseconds(),
				ModuleWallTimeMS:          moduleResult.CompletedAt.Sub(moduleResult.StartedAt).Milliseconds(),
			})
		}
	}

	validatorSummaries := make([]ValidatorTimingSummary, 0, len(i.validatorSummaries))
	for _, summary := range i.validatorSummaries {
		validatorSummaries = append(validatorSummaries, *summary)
	}
	sortValidatorSummaries(validatorSummaries)

	return &UtilizationDiagnostics{
		EffectiveConcurrency: limits,
		RunTiming: RunTimingSummary{
			RunWallTimeMS:             completedAt.Sub(i.startedAt).Milliseconds(),
			SchedulerQueueWaitMS:      i.schedulerQueueWait.Milliseconds(),
			LLMExecutionTimeMS:        i.llmExecutionTime.Milliseconds(),
			DeterministicValidationMS: i.deterministicValidation.Milliseconds(),
			MaxInFlightLLMCalls:       i.maxInflight,
			AverageInFlightLLMCalls:   avgInflight,
		},
		LLMCalls: LLMCallSummary{
			TotalProposalCalls:   i.totalProposalCalls,
			TotalValidationCalls: i.totalValidationCalls,
		},
		Modules:    moduleSummaries,
		Validators: validatorSummaries,
	}
}

type measuredClient struct {
	inner      contracts.StructuredLLMClient
	collector  *runInstrumentation
	isProposal bool
}

func (c measuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	if c.collector != nil {
		c.collector.beginInFlight()
	}
	started := time.Now().UTC()
	resp, err := c.inner.CompleteStructured(ctx, req, out)
	elapsed := time.Since(started)
	if c.collector != nil {
		c.collector.endInFlight()
		c.collector.mu.Lock()
		moduleInstance := moduleInstanceFromStage(req.StageName)
		c.collector.llmExecutionTime += elapsed
		c.collector.perModuleLLMExecutionTime[moduleInstance] += elapsed
		if c.isProposal {
			c.collector.totalProposalCalls++
			c.collector.perModuleProposalCalls[moduleInstance]++
		} else {
			c.collector.totalValidationCalls++
			c.collector.perModuleValidationCalls[moduleInstance]++
		}
		c.collector.mu.Unlock()
	}
	return resp, err
}

type measuredScheduler struct {
	inner      contracts.LLMScheduler
	collector  *runInstrumentation
	isProposal bool
}

func (s measuredScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
	started := time.Now().UTC()
	var execDuration time.Duration
	err := s.inner.Run(ctx, func(runCtx context.Context) error {
		execStart := time.Now().UTC()
		callErr := fn(runCtx)
		execDuration += time.Since(execStart)
		return callErr
	})
	totalDuration := time.Since(started)
	waitDuration := totalDuration - execDuration
	if waitDuration < 0 {
		waitDuration = 0
	}
	if s.collector != nil {
		stageModule := moduleInstanceFromContext(ctx)
		s.collector.mu.Lock()
		s.collector.schedulerQueueWait += waitDuration
		if stageModule != "" {
			s.collector.perModuleSchedulerQueueWait[stageModule] += waitDuration
		}
		s.collector.mu.Unlock()
	}
	_ = s.isProposal
	return err
}

func moduleInstanceFromStage(stage string) string {
	stage = strings.TrimSpace(stage)
	if stage == "" {
		return ""
	}
	parts := strings.Split(stage, ":")
	if len(parts) == 0 {
		return stage
	}
	return strings.TrimSpace(parts[0])
}

func moduleInstanceFromContext(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	if v := ctx.Value(moduleInstanceContextKey{}); v != nil {
		if s, ok := v.(string); ok {
			return strings.TrimSpace(s)
		}
	}
	return ""
}

type moduleInstanceContextKey struct{}

func withModuleInstanceContext(ctx context.Context, moduleInstance string) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, moduleInstanceContextKey{}, strings.TrimSpace(moduleInstance))
}

func isLLMBackedValidator(v contracts.Validator) bool {
	return validatormetadata.ClassOf(v) == validatormetadata.ExecutionClassLLMBacked
}

func sortValidatorSummaries(in []ValidatorTimingSummary) {
	sort.SliceStable(in, func(i, j int) bool {
		if in[i].ModuleInstance != in[j].ModuleInstance {
			return in[i].ModuleInstance < in[j].ModuleInstance
		}
		return in[i].ValidatorKey < in[j].ValidatorKey
	})
}