package main

import (
	"sync"
	"time"

	pk "gitea.maximumdirect.net/eric/promptkit"
)

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"`
}
type Observation struct {
	Kind               string        `json:"kind"`
	RunID              string        `json:"run_id"`
	Role               string        `json:"role"`
	Backend            string        `json:"backend"`
	Profile            string        `json:"profile"`
	Generation         uint64        `json:"generation,omitempty"`
	Attempt            uint64        `json:"attempt,omitempty"`
	Dispatched         bool          `json:"dispatched,omitempty"`
	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"`
	Failure            string        `json:"failure,omitempty"`
	StatusCode         int           `json:"status_code,omitempty"`
	RetriesExhausted   bool          `json:"retries_exhausted,omitempty"`
	Attempts           uint64        `json:"attempts,omitempty"`
	Retries            uint64        `json:"retries,omitempty"`
	Repairs            uint64        `json:"repairs,omitempty"`
	Usage              Usage         `json:"usage"`
}

type Recorder struct {
	mu      sync.Mutex
	records []Observation
}

func (r *Recorder) observer(role string) pk.ExecutionObserver {
	return func(e pk.ExecutionEvent) {
		count := func(c pk.ObservedTokenCount) TokenCount { return TokenCount{c.Value, c.Known, c.Complete} }
		o := Observation{Kind: string(e.Kind), RunID: e.RunID, Role: role, Backend: e.BackendID, Profile: e.ProfileID,
			Generation: e.Generation, Attempt: e.Attempt, Dispatched: e.Dispatched, CapacityWait: e.CapacityWait, GenerationDuration: e.GenerationDuration, RetryDelay: e.RetryDelay, Duration: e.Duration,
			Outcome: string(e.Outcome), Failure: string(e.Failure), StatusCode: e.StatusCode, RetriesExhausted: e.RetriesExhausted, Attempts: e.Attempts, Retries: e.TransientRetries, Repairs: e.RepairCalls,
			Usage: Usage{count(e.Usage.PromptTokens), count(e.Usage.CompletionTokens), count(e.Usage.TotalTokens), count(e.Usage.CachedTokens), count(e.Usage.CacheWriteTokens)}}
		r.mu.Lock()
		r.records = append(r.records, o)
		r.mu.Unlock()
	}
}
func (r *Recorder) snapshot() []Observation {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]Observation(nil), r.records...)
}

// totals uses only finished records. Attempt usage is retained for diagnosis,
// but is never added to the already aggregated terminal usage.
func (r *Recorder) totals() (attempts, retries, repairs uint64) {
	for _, o := range r.snapshot() {
		if o.Kind == string(pk.ExecutionFinished) {
			attempts += o.Attempts
			retries += o.Retries
			repairs += o.Repairs
		}
	}
	return
}