// Package observation owns invocation-local execution identities, metadata
// accounting, and synchronous event delivery.
package observation

import (
	"context"
	"crypto/rand"
	"fmt"
	"math"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

// Recorder accumulates one execution's measured attempts and usage. Its owner
// calls methods in execution order; it has no shared history or background work.
type Recorder struct {
	runID                string
	startedAt            time.Time
	now                  func() time.Time
	observer             domain.ExecutionObserver
	selection            Selection
	attempts             uint64
	retries              uint64
	repairCalls          uint64
	lastRepairGeneration uint64
	finished             bool
	capacityWait         time.Duration
	generationDuration   time.Duration
	retryDelay           time.Duration
	usage                usageTotals
}

// Selection contains resolved identities only. Unresolved fields stay empty.
type Selection struct {
	PromptID      string
	PromptVersion string
	ProfileID     string
	BackendID     string
	Model         string
}

// AttemptData describes one completed active-permit acquisition and optional
// generation dispatch. Non-dispatched attempts contribute no token usage.
type AttemptData struct {
	Generation         uint64
	Attempt            uint64
	Dispatched         bool
	CapacityWait       time.Duration
	GenerationDuration time.Duration
	Outcome            domain.ExecutionOutcome
	Failure            domain.ExecutionFailureKind
	StatusCode         int
	Usage              domain.TokenUsage
	UsagePresence      *domain.TokenUsagePresence
}

// DelayData describes one completed or interrupted retry wait.
type DelayData struct {
	Generation uint64
	Attempt    uint64
	Elapsed    time.Duration
	Failure    domain.ExecutionFailureKind
}

// FinishData describes the terminal outcome after resource release.
type FinishData struct {
	Outcome          domain.ExecutionOutcome
	Failure          domain.ExecutionFailureKind
	StatusCode       int
	RetriesExhausted bool
}

// New creates an invocation-local recorder and run ID. It returns the recorder
// even if ID creation fails, allowing an error terminal event with an empty ID.
func New(observer domain.ExecutionObserver) (*Recorder, error) {
	runID, err := NewRunID()
	return newRecorder(runID, time.Now, observer), err
}

func newRecorder(runID string, now func() time.Time, observer domain.ExecutionObserver) *Recorder {
	return &Recorder{runID: runID, startedAt: now(), now: now, observer: observer,
		usage: newUsageTotals()}
}

// NewRunID creates the identity shared by observations and successful results.
func NewRunID() (string, error) {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", err
	}
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
}

// RunID returns this invocation's identity.
func (r *Recorder) RunID() string { return r.runID }

// Start delivers the opening event with only identities already resolved.
func (r *Recorder) Start() {
	event := r.event(domain.ExecutionStarted)
	r.deliver(event)
}

// UpdateSelection replaces the known resolved identities for later events.
func (r *Recorder) UpdateSelection(selection Selection) { r.selection = selection }

// Attempt records one completed acquisition and optional dispatch.
func (r *Recorder) Attempt(data AttemptData) {
	r.capacityWait = addDuration(r.capacityWait, data.CapacityWait)
	if data.Dispatched {
		r.attempts++
		if data.Attempt > 1 {
			r.retries++
		}
		if data.Generation > 1 && data.Generation != r.lastRepairGeneration {
			r.repairCalls++
			r.lastRepairGeneration = data.Generation
		}
		r.generationDuration = addDuration(r.generationDuration, data.GenerationDuration)
		r.usage.add(data.Usage, data.UsagePresence)
	}
	event := r.event(domain.ExecutionAttempt)
	event.Generation = data.Generation
	event.Attempt = data.Attempt
	event.Dispatched = data.Dispatched
	event.CapacityWait = data.CapacityWait
	if data.Dispatched {
		event.GenerationDuration = data.GenerationDuration
	}
	event.Outcome = data.Outcome
	event.Failure = data.Failure
	event.StatusCode = data.StatusCode
	if data.Dispatched {
		event.Usage = observedUsage(data.Usage, data.UsagePresence)
	}
	r.deliver(event)
}

// Delay records the measured time spent waiting after one failed attempt.
func (r *Recorder) Delay(data DelayData) {
	r.retryDelay = addDuration(r.retryDelay, data.Elapsed)
	event := r.event(domain.ExecutionRetryDelay)
	event.Generation = data.Generation
	event.Attempt = data.Attempt
	event.RetryDelay = data.Elapsed
	event.Failure = data.Failure
	r.deliver(event)
}

// Finish delivers terminal totals once on a normal return. Its caller owns the
// operation's resource cleanup and failure classification.
func (r *Recorder) Finish(data FinishData) {
	if r.finished {
		return
	}
	r.finished = true
	event := r.event(domain.ExecutionFinished)
	event.Outcome = data.Outcome
	event.Failure = data.Failure
	event.StatusCode = data.StatusCode
	event.RetriesExhausted = data.RetriesExhausted
	event.Attempts = r.attempts
	event.TransientRetries = r.retries
	event.RepairCalls = r.repairCalls
	event.CapacityWait = r.capacityWait
	event.GenerationDuration = r.generationDuration
	event.RetryDelay = r.retryDelay
	event.Duration = event.Time.Sub(r.startedAt)
	if event.Duration < 0 {
		event.Duration = 0
	}
	event.Usage = r.usage.snapshot(r.attempts)
	r.deliver(event)
}

func (r *Recorder) event(kind domain.ExecutionEventKind) domain.ExecutionEvent {
	return domain.ExecutionEvent{
		Kind: kind, RunID: r.runID, Time: r.now(), PromptID: r.selection.PromptID,
		PromptVersion: r.selection.PromptVersion, ProfileID: r.selection.ProfileID,
		BackendID: r.selection.BackendID, Model: r.selection.Model,
	}
}

func (r *Recorder) deliver(event domain.ExecutionEvent) {
	if r.observer != nil {
		r.observer(event)
	}
}

type contextKey struct{}

// WithoutExecutionState starts a public operation without inheriting another
// invocation's recorder or prepared-claim sink. Other context values and all
// cancellation/deadline behavior remain with the caller's context.
func WithoutExecutionState(ctx context.Context) context.Context {
	return executionBoundary{Context: ctx}
}

type executionBoundary struct{ context.Context }

func (c executionBoundary) Value(key any) any {
	switch key.(type) {
	case contextKey, claimSinkKey:
		return nil
	default:
		return c.Context.Value(key)
	}
}

// WithRecorder carries one invocation's recorder through internal client calls.
func WithRecorder(ctx context.Context, recorder *Recorder) context.Context {
	return context.WithValue(ctx, contextKey{}, recorder)
}

// FromContext returns the invocation-local recorder, if one was installed.
func FromContext(ctx context.Context) *Recorder {
	recorder, _ := ctx.Value(contextKey{}).(*Recorder)
	return recorder
}

type claimSinkKey struct{}

// WithClaimSink lets the facade learn which invocation actually claimed a
// prepared handle, so it can emit the terminal event after public mapping.
func WithClaimSink(ctx context.Context, sink **Recorder) context.Context {
	return context.WithValue(ctx, claimSinkKey{}, sink)
}

// RegisterClaim records an invocation after its prepared handle is claimed.
func RegisterClaim(ctx context.Context, recorder *Recorder) {
	if sink, ok := ctx.Value(claimSinkKey{}).(**Recorder); ok && sink != nil {
		*sink = recorder
	}
}

type usageField struct {
	value           int
	known, complete bool
}
type usageTotals struct {
	prompt, completion, total, cached, cacheWrite usageField
}

func newUsageTotals() usageTotals {
	return usageTotals{prompt: usageField{complete: true}, completion: usageField{complete: true},
		total: usageField{complete: true}, cached: usageField{complete: true}, cacheWrite: usageField{complete: true}}
}

func (u *usageTotals) add(usage domain.TokenUsage, presence *domain.TokenUsagePresence) {
	known := usagePresence(usage, presence)
	u.prompt.add(usage.PromptTokens, known.PromptTokens)
	u.completion.add(usage.CompletionTokens, known.CompletionTokens)
	u.total.add(usage.TotalTokens, known.TotalTokens)
	u.cached.add(usage.CachedTokens, known.CachedTokens)
	u.cacheWrite.add(usage.CacheWriteTokens, known.CacheWriteTokens)
}

func (u usageTotals) snapshot(attempts uint64) domain.ObservedTokenUsage {
	return domain.ObservedTokenUsage{
		PromptTokens: u.prompt.snapshot(attempts), CompletionTokens: u.completion.snapshot(attempts),
		TotalTokens: u.total.snapshot(attempts), CachedTokens: u.cached.snapshot(attempts),
		CacheWriteTokens: u.cacheWrite.snapshot(attempts),
	}
}

func (f *usageField) add(value int, known bool) {
	if known {
		f.value += value
		f.known = true
	} else {
		f.complete = false
	}
}

func (f usageField) snapshot(attempts uint64) domain.ObservedTokenCount {
	if attempts == 0 {
		return domain.ObservedTokenCount{Known: true, Complete: true}
	}
	return domain.ObservedTokenCount{Value: f.value, Known: f.known, Complete: f.complete}
}

func observedUsage(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.ObservedTokenUsage {
	known := usagePresence(usage, presence)
	return domain.ObservedTokenUsage{
		PromptTokens:     observedCount(usage.PromptTokens, known.PromptTokens),
		CompletionTokens: observedCount(usage.CompletionTokens, known.CompletionTokens),
		TotalTokens:      observedCount(usage.TotalTokens, known.TotalTokens),
		CachedTokens:     observedCount(usage.CachedTokens, known.CachedTokens),
		CacheWriteTokens: observedCount(usage.CacheWriteTokens, known.CacheWriteTokens),
	}
}

func usagePresence(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.TokenUsagePresence {
	if presence != nil {
		return *presence
	}
	return domain.TokenUsagePresence{
		PromptTokens: usage.PromptTokens != 0, CompletionTokens: usage.CompletionTokens != 0,
		TotalTokens: usage.TotalTokens != 0, CachedTokens: usage.CachedTokens != 0,
		CacheWriteTokens: usage.CacheWriteTokens != 0,
	}
}

func observedCount(value int, known bool) domain.ObservedTokenCount {
	if !known {
		return domain.ObservedTokenCount{}
	}
	return domain.ObservedTokenCount{Value: value, Known: true, Complete: true}
}

func addDuration(total, next time.Duration) time.Duration {
	if next > 0 && total > time.Duration(math.MaxInt64)-next {
		return time.Duration(math.MaxInt64)
	}
	return total + next
}