package metrics

import (
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

const (
	PromptSessionFreshnessFresh  = "fresh"
	PromptSessionFreshnessResume = "resume"
	MaxTokenEvidenceSamples      = 200_000
)

// PromptAttributionRecorder persists body-free Prompt/Context attribution.
type PromptAttributionRecorder interface {
	RecordPromptCompileMetrics(items []PromptCompileRecord)
}

// TokenAnalysisRecorder persists the detailed, high-cardinality evidence used
// by the opt-in AgentTeam Token diagnosis.
type TokenAnalysisRecorder interface {
	PromptAttributionRecorder
	RecordUsageObservation(item UsageObservationRecord, exposure *InvocationExposureRecord)
}

// TokenAnalysisQuerier exposes the one bounded evidence read required by the
// daemon-owned Token diagnosis.
type TokenAnalysisQuerier interface {
	QueryTokenAnalysisEvidence(query TokenEvidenceQuery) TokenAnalysisEvidence
}

// PromptCompileRecord is one block-level compile estimate stored locally.
type PromptCompileRecord struct {
	Timestamp          int64
	LaunchStartedAt    int64
	AgentID            string
	Runtime            string
	LaunchID           string
	SessionKind        string
	Freshness          string
	LeaderSessionKey   string
	EventID            string
	CompileEvent       string
	Stage              string
	Slot               string
	BlockID            string
	BlockLayer         string
	BlockSlot          string
	BlockType          string
	Durability         string
	TaskPayloadPresent bool
	TaskPayloadFP      string
	Records            uint64
	Tokens             uint64
}

// UsageObservationRecord is one normalized runtime usage delta correlated to
// a AgentTeam launch. It never contains prompt, task, tool, or model output text.
type UsageObservationRecord struct {
	Timestamp        int64  `json:"timestamp"`
	LaunchStartedAt  int64  `json:"launchStartedAt"`
	ObservationID    string `json:"-"`
	AgentID          string `json:"agentId"`
	Runtime          string `json:"runtime"`
	LaunchID         string `json:"launchId"`
	SessionKind      string `json:"sessionKind"`
	SessionID        string `json:"sessionId,omitempty"`
	Model            string `json:"model,omitempty"`
	InvocationID     string `json:"invocationId,omitempty"`
	InvocationCause  string `json:"invocationCause"`
	ExposureExcluded bool   `json:"-"`
	InputTokens      uint64 `json:"inputTokens"`
	CachedTokens     uint64 `json:"cachedTokens"`
	OutputTokens     uint64 `json:"outputTokens"`
}

var usageObservationSequence atomic.Uint64

// InvocationExposureRecord is the body-free compiler Context known to be
// active for one uniquely identified model invocation.
type InvocationExposureRecord struct {
	Complete bool
	Blocks   []InvocationExposureBlock
}

// InvocationExposureBlock is one semantic block active for an invocation.
type InvocationExposureBlock struct {
	BlockID    string
	Durability string
	Tokens     uint64
}

// ExposureCoverageRecord aggregates invocation-level correlation coverage.
type ExposureCoverageRecord struct {
	LaunchID             string
	Observations         uint64
	CompleteObservations uint64
}

// BlockExposureRecord aggregates one block's invocation exposure.
type BlockExposureRecord struct {
	LaunchID      string
	BlockID       string
	Durability    string
	ExposureCount uint64
	Tokens        uint64
}

// TokenEvidenceQuery selects one bounded detailed-evidence window.
type TokenEvidenceQuery struct {
	Window     TimeWindow
	AgentIDs   []string
	LaunchID   string
	SnapshotAt int64
	Limit      int
}

// TokenAnalysisEvidence is the one-pass, bounded evidence input consumed by
// the daemon-owned token analysis service.
type TokenAnalysisEvidence struct {
	PromptRecords    []PromptCompileRecord    `json:"promptRecords"`
	Usage            []UsageObservationRecord `json:"usage"`
	ExposureCoverage []ExposureCoverageRecord `json:"exposureCoverage"`
	BlockExposures   []BlockExposureRecord    `json:"blockExposures"`
	Scanned          int                      `json:"scanned"`
	Limit            int                      `json:"limit"`
	Truncated        bool                     `json:"truncated"`
}

// RecordPromptCompileMetrics persists block-level compile estimates.
func (tc *TSDBCollector) RecordPromptCompileMetrics(items []PromptCompileRecord) {
	if tc == nil || tc.db == nil || len(items) == 0 {
		return
	}
	for i := range items {
		if items[i].Timestamp == 0 {
			items[i].Timestamp = time.Now().UnixMilli()
		}
		items[i].AgentID = strings.TrimSpace(items[i].AgentID)
		items[i].Runtime = strings.TrimSpace(items[i].Runtime)
		items[i].LaunchID = strings.TrimSpace(items[i].LaunchID)
		items[i].SessionKind = strings.TrimSpace(items[i].SessionKind)
		items[i].Freshness = strings.TrimSpace(items[i].Freshness)
	}
	tc.BatchRecordPromptCompileMetrics(items)
}

// RecordUsageObservation persists one opt-in detailed runtime usage delta.
func (tc *TSDBCollector) RecordUsageObservation(
	item UsageObservationRecord,
	exposure *InvocationExposureRecord,
) {
	if tc == nil || tc.db == nil {
		return
	}
	if item.Timestamp == 0 {
		item.Timestamp = time.Now().UnixMilli()
	}
	if strings.TrimSpace(item.ObservationID) == "" {
		item.ObservationID = strconv.FormatInt(item.Timestamp, 36) + "-" +
			strconv.FormatUint(usageObservationSequence.Add(1), 36)
	}
	item.AgentID = strings.TrimSpace(item.AgentID)
	item.Runtime = strings.TrimSpace(item.Runtime)
	item.LaunchID = strings.TrimSpace(item.LaunchID)
	item.SessionKind = strings.TrimSpace(item.SessionKind)
	item.SessionID = strings.TrimSpace(item.SessionID)
	item.Model = strings.TrimSpace(item.Model)
	item.InvocationID = strings.TrimSpace(item.InvocationID)
	item.InvocationCause = strings.TrimSpace(item.InvocationCause)
	if item.AgentID == "" || item.Runtime == "" || item.LaunchID == "" || item.SessionKind == "" {
		return
	}
	if item.InvocationCause == "" {
		item.InvocationCause = "unknown"
	}
	if item.InvocationID == "" || item.ExposureExcluded {
		exposure = nil
	}
	tc.recordUsageObservation(item, exposure)
}

// QueryTokenAnalysisEvidence returns one bounded detailed evidence snapshot.
func (tc *TSDBCollector) QueryTokenAnalysisEvidence(query TokenEvidenceQuery) TokenAnalysisEvidence {
	if tc == nil || tc.promptQueryable() == nil {
		return TokenAnalysisEvidence{PromptRecords: []PromptCompileRecord{}, Usage: []UsageObservationRecord{}}
	}
	return tc.queryTokenAnalysisEvidence(query)
}