package llm

import (
	"sort"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

// TokenCount retains the distinction between zero and unavailable usage.
type TokenCount struct {
	Value    int  `json:"value"`
	Known    bool `json:"known"`
	Complete bool `json:"complete"`
}

type Usage struct {
	Prompt     TokenCount `json:"prompt"`
	Completion TokenCount `json:"completion"`
	Total      TokenCount `json:"total"`
	Cached     TokenCount `json:"cached"`
	CacheWrite TokenCount `json:"cache_write"`
}

// ExecutionObservation contains execution metadata only. It never contains
// endpoint URLs, credentials, messages, provider errors, or generated text.
type ExecutionObservation 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"`
	CapacityWait       time.Duration                      `json:"capacity_wait_ns"`
	GenerationDuration time.Duration                      `json:"generation_duration_ns"`
	RetryDelay         time.Duration                      `json:"retry_delay_ns"`
	Duration           time.Duration                      `json:"duration_ns"`
	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              Usage                              `json:"usage"`
}

type LogicalCall 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"`
	WallDuration time.Duration                      `json:"wall_duration_ns"`
	Failure      structuredoutput.FailureKind       `json:"failure,omitempty"`
}

type ObservationSnapshot struct {
	Calls  []LogicalCall          `json:"calls"`
	Events []ExecutionObservation `json:"events"`
}

type ExecutionSummary struct {
	LogicalCalls      uint64 `json:"logical_calls"`
	ProviderAttempts  uint64 `json:"provider_attempts"`
	TransientRetries  uint64 `json:"transient_retries"`
	StructuralRepairs uint64 `json:"structural_repairs"`
	Usage             Usage  `json:"usage"`
}

// ObservationCollector accepts short synchronous callbacks from parallel calls.
// Snapshot should be taken after the runner has joined active section work.
type ObservationCollector struct {
	mu     sync.Mutex
	calls  map[uint64]LogicalCall
	events []ExecutionObservation
}

func NewObservationCollector() *ObservationCollector {
	return &ObservationCollector{calls: make(map[uint64]LogicalCall)}
}

func (c *ObservationCollector) Hooks() CallHooks {
	return CallHooks{OnCall: c.recordCall, OnEventWithCall: c.recordEvent}
}

func copyContext(value structuredoutput.DiagnosticContext) structuredoutput.DiagnosticContext {
	if value.SectionIndex != nil {
		index := *value.SectionIndex
		value.SectionIndex = &index
	}
	if value.BatchIndex != nil {
		index := *value.BatchIndex
		value.BatchIndex = &index
	}
	return value
}

func (c *ObservationCollector) recordCall(record structuredoutput.CallRecord) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	call := LogicalCall{CallID: record.CallID, Context: copyContext(record.Context), Prompt: record.Prompt, Started: record.Started, Finished: record.Finished, Failure: record.Failure}
	if !record.Finished.IsZero() {
		call.WallDuration = record.Finished.Sub(record.Started)
	}
	c.calls[record.CallID] = call
}

func (c *ObservationCollector) recordEvent(callID uint64, context structuredoutput.DiagnosticContext, event pk.ExecutionEvent) {
	if c == nil {
		return
	}
	count := func(value pk.ObservedTokenCount) TokenCount {
		return TokenCount{Value: value.Value, Known: value.Known, Complete: value.Complete}
	}
	observation := ExecutionObservation{
		CallID: callID, Context: copyContext(context), Kind: string(event.Kind), RunID: event.RunID,
		Time: event.Time, BackendID: event.BackendID, ProfileID: event.ProfileID, Model: event.Model,
		Generation: event.Generation, Attempt: event.Attempt, Dispatched: event.Dispatched,
		CapacityWait: event.CapacityWait, GenerationDuration: event.GenerationDuration, RetryDelay: event.RetryDelay, Duration: event.Duration,
		Outcome: string(event.Outcome), Failure: string(event.Failure), StatusCode: event.StatusCode, RetriesExhausted: event.RetriesExhausted,
		Attempts: event.Attempts, TransientRetries: event.TransientRetries, StructuralRepairs: event.RepairCalls,
		Usage: Usage{Prompt: count(event.Usage.PromptTokens), Completion: count(event.Usage.CompletionTokens), Total: count(event.Usage.TotalTokens), Cached: count(event.Usage.CachedTokens), CacheWrite: count(event.Usage.CacheWriteTokens)},
	}
	c.mu.Lock()
	c.events = append(c.events, observation)
	c.mu.Unlock()
}

func (c *ObservationCollector) Snapshot() ObservationSnapshot {
	if c == nil {
		return ObservationSnapshot{}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	snapshot := ObservationSnapshot{Calls: make([]LogicalCall, 0, len(c.calls)), Events: make([]ExecutionObservation, len(c.events))}
	for _, call := range c.calls {
		call.Context = copyContext(call.Context)
		snapshot.Calls = append(snapshot.Calls, call)
	}
	copy(snapshot.Events, c.events)
	for i := range snapshot.Events {
		snapshot.Events[i].Context = copyContext(snapshot.Events[i].Context)
	}
	sort.Slice(snapshot.Calls, func(i, j int) bool { return snapshot.Calls[i].CallID < snapshot.Calls[j].CallID })
	sort.SliceStable(snapshot.Events, func(i, j int) bool {
		if snapshot.Events[i].CallID != snapshot.Events[j].CallID {
			return snapshot.Events[i].CallID < snapshot.Events[j].CallID
		}
		return snapshot.Events[i].Time.Before(snapshot.Events[j].Time)
	})
	return snapshot
}

// Summarize counts terminal usage once per logical call. PromptKit's terminal
// usage already aggregates its attempts, retries, and repairs.
func Summarize(snapshot ObservationSnapshot) ExecutionSummary {
	result := ExecutionSummary{LogicalCalls: uint64(len(snapshot.Calls))}
	terminal := make(map[uint64]ExecutionObservation)
	for _, event := range snapshot.Events {
		if event.Kind == string(pk.ExecutionFinished) {
			terminal[event.CallID] = event
		}
	}
	dispatched := false
	aggregate := func(dst *TokenCount, source TokenCount) {
		if source.Known {
			dst.Value += source.Value
			dst.Known = true
		}
		dst.Complete = dst.Complete && source.Complete
	}
	result.Usage = Usage{Prompt: TokenCount{Complete: true}, Completion: TokenCount{Complete: true}, Total: TokenCount{Complete: true}, Cached: TokenCount{Complete: true}, CacheWrite: TokenCount{Complete: true}}
	for _, call := range snapshot.Calls {
		event, ok := terminal[call.CallID]
		if !ok {
			continue
		}
		result.ProviderAttempts += event.Attempts
		result.TransientRetries += event.TransientRetries
		result.StructuralRepairs += event.StructuralRepairs
		if event.Attempts == 0 {
			continue
		}
		dispatched = true
		aggregate(&result.Usage.Prompt, event.Usage.Prompt)
		aggregate(&result.Usage.Completion, event.Usage.Completion)
		aggregate(&result.Usage.Total, event.Usage.Total)
		aggregate(&result.Usage.Cached, event.Usage.Cached)
		aggregate(&result.Usage.CacheWrite, event.Usage.CacheWrite)
	}
	if !dispatched {
		result.Usage.Prompt.Known = true
		result.Usage.Completion.Known = true
		result.Usage.Total.Known = true
		result.Usage.Cached.Known = true
		result.Usage.CacheWrite.Known = true
	}
	return result
}