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

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

	"go.maximumdirect.net/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
	parentRunID          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                UsageAccumulator
}

// Selection contains resolved identities only. Unresolved fields stay empty.
type Selection struct {
	PromptID  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}
}

// 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 }

// SetParentRunID records the immediate source identity before start delivery.
func (r *Recorder) SetParentRunID(id string) { r.parentRunID = id }

// 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.Observed()
	r.deliver(event)
}

func (r *Recorder) event(kind domain.ExecutionEventKind) domain.ExecutionEvent {
	return domain.ExecutionEvent{
		Kind: kind, RunID: r.runID, ParentRunID: r.parentRunID, Time: r.now(), PromptID: r.selection.PromptID,
		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
}

// UsageAccumulator combines presence-aware usage across dispatched attempts.
// Its zero value represents no dispatched attempts.
type UsageAccumulator struct {
	totals   usageTotals
	attempts uint64
}

// Add includes the usage reported by one dispatched attempt.
func (u *UsageAccumulator) Add(usage domain.TokenUsage, presence *domain.TokenUsagePresence) {
	if u.attempts == 0 {
		u.totals = completeUsageTotals()
	}
	u.attempts++
	u.totals.add(usage, presence)
}

// Observed returns aggregate counts, marking every field known and complete
// when no attempt has been added.
func (u UsageAccumulator) Observed() domain.ObservedTokenUsage {
	return u.totals.snapshot(u.attempts)
}

func completeUsageTotals() 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
}