package processreport

import (
	"time"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/core/reporting"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/runner"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

func mapLLMSummary(source llm.ExecutionSummary) reporting.LLMSummary {
	metric := func(value llm.TokenCount) reporting.TokenMetric {
		return reporting.TokenMetric{Value: value.Value, Known: value.Known, Complete: value.Complete}
	}
	return reporting.LLMSummary{
		LogicalCalls: source.LogicalCalls, ProviderAttempts: source.ProviderAttempts,
		TransientRetries: source.TransientRetries, StructuralRepairs: source.StructuralRepairs,
		Usage: reporting.LLMUsage{Prompt: metric(source.Usage.Prompt), Completion: metric(source.Usage.Completion), Total: metric(source.Usage.Total), Cached: metric(source.Usage.Cached), CacheWrite: metric(source.Usage.CacheWrite)},
	}
}

// BuildUtilization groups complete logical calls by their selected backend.
func BuildUtilization(output *runner.RunOutput, backends map[string]config.Backend) reporting.UtilizationDiagnostics {
	result := reporting.UtilizationDiagnostics{Backends: make(map[string]reporting.BackendUtilization, len(backends))}
	if output == nil {
		return result
	}
	result.RunWallTimeMS = output.RunWallTimeMS
	groups := make(map[string]*llm.ObservationSnapshot, len(backends))
	callBackends := make(map[uint64]string, len(output.Observations.Calls))
	moduleIndices := make(map[string]int)
	moduleIndex := func(key, instance string) int {
		if i, ok := moduleIndices[instance]; ok {
			return i
		}
		i := len(result.Modules)
		moduleIndices[instance] = i
		result.Modules = append(result.Modules, reporting.ModuleUtilization{ModuleKey: key, ModuleInstance: instance})
		return i
	}
	// Seed configured order, including modules that made no LLM calls.
	for _, module := range output.ModuleResults {
		i := moduleIndex(module.ModuleKey, module.ModuleInstance)
		result.Modules[i].ModuleWallTimeMS = module.CompletedAt.Sub(module.StartedAt).Milliseconds()
	}
	for _, call := range output.Observations.Calls {
		id := call.Prompt.BackendID
		if groups[id] == nil {
			groups[id] = &llm.ObservationSnapshot{}
		}
		groups[id].Calls = append(groups[id].Calls, call)
		callBackends[call.CallID] = id
		if call.Context.ModuleInstance != "" {
			i := moduleIndex(call.Context.ModuleKey, call.Context.ModuleInstance)
			switch call.Context.Role {
			case "proposal":
				result.Modules[i].ProposalLLMCalls++
			case "validation":
				result.Modules[i].ValidationLLMCalls++
			}
		}
	}
	for _, event := range output.Observations.Events {
		if id, ok := callBackends[event.CallID]; ok {
			groups[id].Events = append(groups[id].Events, event)
		}
	}
	for id, backend := range backends {
		group := groups[id]
		if group == nil {
			group = &llm.ObservationSnapshot{}
		}
		var wait, generation, retry, wall time.Duration
		for _, call := range group.Calls {
			wall += call.WallDuration
		}
		for _, event := range group.Events {
			if event.Kind == string(pk.ExecutionFinished) {
				wait += event.CapacityWait
				generation += event.GenerationDuration
				retry += event.RetryDelay
			}
		}
		result.Backends[id] = reporting.BackendUtilization{
			Concurrency: backend.Concurrency, QueueCapacity: backend.QueueCapacity,
			LLMSummary:     mapLLMSummary(llm.Summarize(*group)),
			CapacityWaitMS: wait.Milliseconds(), GenerationMS: generation.Milliseconds(),
			RetryDelayMS: retry.Milliseconds(), LogicalWallMS: wall.Milliseconds(),
		}
	}
	for _, timing := range output.ValidatorTimings {
		result.Validators = append(result.Validators, reporting.ValidatorUtilization{ModuleKey: timing.ModuleKey, ModuleInstance: timing.ModuleInstance, ValidatorKey: timing.ValidatorName, LLMBacked: timing.LLMBacked, Calls: timing.Calls, ElapsedMS: timing.ElapsedMS})
	}
	return result
}

type executionCall struct {
	CallID    uint64                             `json:"call_id"`
	Context   structuredoutput.DiagnosticContext `json:"context"`
	Prompt    structuredoutput.PromptMetadata    `json:"prompt"`
	Started   time.Time                          `json:"started"`
	Finished  *time.Time                         `json:"finished,omitempty"`
	ElapsedMS int64                              `json:"elapsed_ms"`
	Failure   structuredoutput.FailureKind       `json:"failure,omitempty"`
	Summary   reporting.LLMSummary               `json:"summary"`
}

type executionEvent struct {
	CallID            uint64                             `json:"call_id"`
	Context           structuredoutput.DiagnosticContext `json:"context"`
	Kind              string                             `json:"kind"`
	RunID             string                             `json:"run_id,omitempty"`
	Time              time.Time                          `json:"time"`
	BackendID         string                             `json:"backend_id,omitempty"`
	ProfileID         string                             `json:"profile_id,omitempty"`
	Model             string                             `json:"model,omitempty"`
	Generation        *uint64                            `json:"generation,omitempty"`
	Attempt           *uint64                            `json:"attempt,omitempty"`
	Dispatched        bool                               `json:"dispatched"`
	CapacityWaitMS    int64                              `json:"capacity_wait_ms"`
	GenerationMS      int64                              `json:"generation_ms"`
	RetryDelayMS      int64                              `json:"retry_delay_ms"`
	DurationMS        int64                              `json:"duration_ms"`
	Outcome           string                             `json:"outcome,omitempty"`
	Failure           string                             `json:"failure,omitempty"`
	StatusCode        int                                `json:"status_code,omitempty"`
	RetriesExhausted  bool                               `json:"retries_exhausted,omitempty"`
	Attempts          uint64                             `json:"attempts"`
	TransientRetries  uint64                             `json:"transient_retries"`
	StructuralRepairs uint64                             `json:"structural_repairs"`
	Usage             reporting.LLMUsage                 `json:"usage"`
}

// BuildLLMExecution serializes only safe observation fields and millisecond durations.
func BuildLLMExecution(snapshot llm.ObservationSnapshot) any {
	calls := make([]executionCall, 0, len(snapshot.Calls))
	events := make([]executionEvent, 0, len(snapshot.Events))
	eventsByCall := make(map[uint64][]llm.ExecutionObservation, len(snapshot.Calls))
	for _, event := range snapshot.Events {
		eventsByCall[event.CallID] = append(eventsByCall[event.CallID], event)
	}
	for _, call := range snapshot.Calls {
		group := llm.ObservationSnapshot{Calls: []llm.LogicalCall{call}, Events: eventsByCall[call.CallID]}
		var finished *time.Time
		if !call.Finished.IsZero() {
			value := call.Finished
			finished = &value
		}
		calls = append(calls, executionCall{CallID: call.CallID, Context: call.Context, Prompt: call.Prompt, Started: call.Started, Finished: finished, ElapsedMS: call.WallDuration.Milliseconds(), Failure: call.Failure, Summary: mapLLMSummary(llm.Summarize(group))})
	}
	for _, event := range snapshot.Events {
		var generation, attempt *uint64
		if event.Generation != 0 {
			value := event.Generation
			generation = &value
		}
		if event.Attempt != 0 {
			value := event.Attempt
			attempt = &value
		}
		events = append(events, executionEvent{CallID: event.CallID, Context: event.Context, Kind: event.Kind, RunID: event.RunID, Time: event.Time, BackendID: event.BackendID, ProfileID: event.ProfileID, Model: event.Model, Generation: generation, Attempt: attempt, Dispatched: event.Dispatched, CapacityWaitMS: event.CapacityWait.Milliseconds(), GenerationMS: event.GenerationDuration.Milliseconds(), RetryDelayMS: event.RetryDelay.Milliseconds(), DurationMS: event.Duration.Milliseconds(), Outcome: event.Outcome, Failure: event.Failure, StatusCode: event.StatusCode, RetriesExhausted: event.RetriesExhausted, Attempts: event.Attempts, TransientRetries: event.TransientRetries, StructuralRepairs: event.StructuralRepairs, Usage: mapLLMSummary(llm.ExecutionSummary{Usage: event.Usage}).Usage})
	}
	return struct {
		Calls  []executionCall  `json:"calls"`
		Events []executionEvent `json:"events"`
	}{calls, events}
}