package agent

import (
	"context"
	"sort"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/metrics"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
)

var promptAttributionRecorder = metrics.DefaultPromptAttributionRecorder
var tokenAnalysisRecorder = metrics.DefaultTokenAnalysisRecorder

func defaultTokenSourceRecord(observation TokenSourceObservation) {
	recorder := metrics.DefaultTokenTrackingRecorder()
	if recorder == nil {
		return
	}
	recorder.RecordTokenTracking(metrics.TokenSourceObservation{
		Timestamp:        observation.Timestamp,
		AgentID:          observation.AgentID,
		Runtime:          observation.Runtime,
		Model:            observation.Model,
		SourceKind:       observation.Source.Kind,
		SourceName:       observation.Source.Name,
		InputToken:       observation.InputTokens,
		CachedInputToken: observation.CachedInputTokens,
		OutputToken:      observation.OutputTokens,
		CallCount:        observation.CallCount,
	})
}

func defaultTokenSourceAuditRecord(observation TokenSourceAuditObservation) {
	recorder := metrics.DefaultTokenTrackingAuditRecorder()
	if recorder == nil {
		return
	}
	candidates := make([]metrics.TokenTrackingAuditSource, 0, len(observation.Candidates))
	for _, candidate := range observation.Candidates {
		candidates = append(candidates, metrics.TokenTrackingAuditSource{
			SourceKind:     candidate.source.Kind,
			SourceName:     candidate.source.Name,
			InputEstimate:  candidate.inputTokens,
			OutputEstimate: candidate.outputTokens,
		})
	}
	allocations := make([]metrics.TokenTrackingAuditAllocation, 0, len(observation.Allocations))
	for _, allocation := range observation.Allocations {
		allocations = append(allocations, metrics.TokenTrackingAuditAllocation{
			SourceKind: allocation.Source.Kind,
			SourceName: allocation.Source.Name,
			Input:      allocation.InputTokens,
			Cached:     allocation.CachedInputTokens,
			Output:     allocation.OutputTokens,
		})
	}
	recorder.RecordTokenTrackingAudit(metrics.TokenTrackingAuditEntry{
		Timestamp:   observation.Timestamp,
		Event:       observation.Event,
		Runtime:     observation.Runtime,
		ToolName:    observation.ToolName,
		RawCommand:  observation.RawCommand,
		SourceKind:  observation.Source.Kind,
		SourceName:  observation.Source.Name,
		Decision:    observation.Decision,
		Candidates:  candidates,
		Allocations: allocations,
	})
}

// recordUsage records token usage from a runtime usage event to the global
// metrics collector. The model is resolved from the event (if set). The model
// configured at process start is only a safe fallback for the primary runtime
// session; a foreign Codex subagent may use a different model.
func (ap *Process) recordUsage(ev drivercontract.ParsedEvent) {
	if ev.InputTokens == 0 && ev.CachedInputTokens == 0 && ev.OutputTokens == 0 {
		return
	}

	agentID, apiKey, model, driverName := ap.usageMetricContext(ev)
	// Machine Overview uses ServerMessageCollector, not RequestRecorder. Keep
	// source recording beside that Machine metric and independent of provider API
	// keys; both stores enforce the same non-empty Agent/model dimensions.
	metricModel := model
	if metricModel == "" {
		metricModel = unknownUsageModel(driverName)
	}
	if agentID != "" {
		ap.recordTokenSource(ev, agentID, metricModel, driverName)
		if serverMessages := metrics.DefaultServerMessageCollector(); serverMessages != nil {
			serverMessages.RecordAgentTokenUsage(agentID, metricModel, ev.InputTokens, ev.CachedInputTokens, ev.OutputTokens)
		}
	} else {
		// Even when the observation is outside the persisted Machine metric
		// dimensions, consume the completed source exactly once so it cannot be
		// charged to a later model call.
		consumption := ap.consumeTokenSource(ev.SessionID, ev.ToolCallID)
		ap.recordTokenSourceConsumptionAudit(driverName, consumption, nil)
	}
	if recorder := tokenAnalysisRecorder(); recorder != nil {
		ap.mu.Lock()
		observation := metrics.UsageObservationRecord{
			LaunchStartedAt:  timestampMillis(ap.launchStartedAt),
			AgentID:          strings.TrimSpace(ap.AgentID),
			Runtime:          strings.TrimSpace(ap.runtime),
			LaunchID:         strings.TrimSpace(ap.launchID),
			SessionKind:      strings.TrimSpace(ap.sessionKind),
			SessionID:        strings.TrimSpace(ev.SessionID),
			Model:            strings.TrimSpace(model),
			InvocationID:     strings.TrimSpace(ev.ModelInvocationID),
			InvocationCause:  "unknown",
			ExposureExcluded: ev.ContextExposureExcluded,
			InputTokens:      ev.InputTokens,
			CachedTokens:     ev.CachedInputTokens,
			OutputTokens:     ev.OutputTokens,
		}
		if observation.SessionID == "" {
			observation.SessionID = strings.TrimSpace(ap.SessionID)
		}
		exposure := ap.invocationExposureLocked(ev)
		ap.mu.Unlock()
		recorder.RecordUsageObservation(observation, exposure)
	} else if ev.ContextCompacted {
		ap.mu.Lock()
		ap.removeCompactablePromptMetricsLocked()
		ap.mu.Unlock()
	}

	if apiKey == "" && strings.TrimSpace(driverName) == "codex" && strings.TrimSpace(agentID) != "" {
		apiKey = "codex:" + strings.TrimSpace(agentID)
	}
	if model == "" || apiKey == "" {
		return
	}

	if collector := metrics.DefaultRequestRecorder(); collector != nil {
		collector.RecordRequest(apiKey, model, ev.InputTokens, ev.CachedInputTokens, ev.OutputTokens, time.Duration(0))
	}
}

func unknownUsageModel(runtime string) string {
	runtime = strings.TrimSpace(runtime)
	if runtime == "" {
		return "unknown"
	}
	return "unknown/" + runtime
}

// SetLaunchPromptCompileMetrics retains only body-free compiler attribution
// metadata for this live Process. If detailed collection is enabled at launch
// it is written immediately; otherwise the daemon can write it once when a
// human enables diagnosis, so future usage from an already-running Session
// still has its persistent Context ledger.
func (ap *Process) SetLaunchPromptCompileMetrics(items []metrics.PromptCompileRecord) {
	if ap == nil || len(items) == 0 {
		return
	}
	ap.mu.Lock()
	ap.launchPromptMetrics = append(ap.launchPromptMetrics[:0], items...)
	ap.resetActivePromptMetricsLocked(items)
	if !metrics.DetailedTokenCollectionEnabled() {
		ap.contextExposureComplete = false
	}
	ap.launchMetricsWritten = false
	ap.mu.Unlock()
	ap.RecordRetainedPromptCompileMetrics()
}

// LaunchPromptCompileMetrics returns the retained body-free launch ledger.
// It exists so dispatcher tests using constructor fakes can assert the same
// Process-owned evidence without depending on the global TSDB recorder.
func (ap *Process) LaunchPromptCompileMetrics() []metrics.PromptCompileRecord {
	if ap == nil {
		return nil
	}
	ap.mu.Lock()
	defer ap.mu.Unlock()
	return append([]metrics.PromptCompileRecord(nil), ap.launchPromptMetrics...)
}

// RecordRetainedPromptCompileMetrics writes the live Process launch ledger at
// most once. It never reconstructs historical deliveries or runtime usage.
func (ap *Process) RecordRetainedPromptCompileMetrics() bool {
	if ap == nil || !metrics.DetailedTokenCollectionEnabled() {
		return false
	}
	recorder := promptAttributionRecorder()
	if recorder == nil {
		return false
	}
	ap.mu.Lock()
	if ap.launchMetricsWritten || len(ap.launchPromptMetrics) == 0 {
		ap.mu.Unlock()
		return false
	}
	items := append([]metrics.PromptCompileRecord(nil), ap.launchPromptMetrics...)
	startedAt := timestampMillis(ap.launchStartedAt)
	for index := range items {
		if items[index].LaunchStartedAt == 0 {
			items[index].LaunchStartedAt = startedAt
		}
	}
	ap.launchMetricsWritten = true
	ap.mu.Unlock()
	recorder.RecordPromptCompileMetrics(items)
	return true
}

func (ap *Process) recordDeliveryPromptCompileMetrics(
	compiled promptcompiler.RuntimeInputEnvelope,
	rendered promptcompiler.RenderedPromptPayload,
) {
	recorder := promptAttributionRecorder()
	if recorder == nil {
		return
	}
	ap.mu.Lock()
	agentID := strings.TrimSpace(ap.AgentID)
	runtime := strings.TrimSpace(ap.runtime)
	launchID := strings.TrimSpace(ap.launchID)
	sessionKind := strings.TrimSpace(ap.sessionKind)
	freshness := strings.TrimSpace(ap.sessionFreshness)
	launchStartedAt := timestampMillis(ap.launchStartedAt)
	ap.mu.Unlock()
	if agentID == "" || runtime == "" || launchID == "" || sessionKind == "" || freshness == "" {
		return
	}
	blocks := rendered.AttributionBlocks
	if len(blocks) == 0 && compiled.Bundle != nil {
		blocks = compiled.Bundle.Blocks
	}
	if len(blocks) == 0 {
		return
	}
	timestamp := time.Now().UnixMilli()
	taskPayloadFP := fingerprintTaskPayload(compiled.TaskPayloadDigest)
	items := make([]metrics.PromptCompileRecord, 0, len(blocks))
	for _, block := range blocks {
		blockID := strings.TrimSpace(block.ID)
		items = append(items, metrics.PromptCompileRecord{
			Timestamp:          timestamp,
			LaunchStartedAt:    launchStartedAt,
			AgentID:            agentID,
			Runtime:            runtime,
			LaunchID:           launchID,
			SessionKind:        sessionKind,
			Freshness:          freshness,
			EventID:            strings.TrimSpace(compiled.EventID),
			CompileEvent:       string(compiled.EventType),
			Stage:              string(compiled.EventKind),
			Slot:               string(block.Slot),
			BlockID:            blockID,
			BlockLayer:         string(block.Layer),
			BlockSlot:          string(block.Slot),
			BlockType:          string(block.Type),
			Durability:         string(block.Durability),
			TaskPayloadPresent: compiled.TaskPayloadPresent,
			TaskPayloadFP:      taskPayloadFP,
			Tokens:             uint64(max(block.TokenEstimate, 0)),
		})
	}
	recorder.RecordPromptCompileMetrics(items)
}

func timestampMillis(value time.Time) int64 {
	if value.IsZero() {
		return 0
	}
	return value.UnixMilli()
}

func (ap *Process) recordGeneratedPromptCompileMetrics(
	eventType promptcompiler.CompileEventType,
	kind promptcompiler.GeneratedPromptKind,
) {
	if !metrics.DetailedTokenCollectionEnabled() || ap == nil || ap.driver == nil {
		return
	}
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())
	carrier, err := compiler.CompileGeneratedPrompt(
		context.Background(), eventType, promptcompiler.NewEventID(eventType),
		promptcompiler.GeneratedPromptInput{Kind: kind}, ap.driver.Name(),
	)
	if err != nil {
		return
	}
	ap.recordDeliveryPromptCompileMetrics(carrier.Envelope(), carrier.Rendered())
	ap.activatePromptBoundary(promptMetricsBoundary{compiled: carrier.Envelope(), rendered: carrier.Rendered()})
}

type promptMetricsBoundary struct {
	compiled promptcompiler.RuntimeInputEnvelope
	rendered promptcompiler.RenderedPromptPayload
}

func (ap *Process) stageDeliveryPromptCompileMetrics(
	inputID string,
	compiled promptcompiler.RuntimeInputEnvelope,
	rendered promptcompiler.RenderedPromptPayload,
) {
	inputID = strings.TrimSpace(inputID)
	if inputID == "" || !metrics.DetailedTokenCollectionEnabled() {
		return
	}
	ap.mu.Lock()
	if ap.pendingPromptMetrics == nil {
		ap.pendingPromptMetrics = make(map[string]promptMetricsBoundary)
	}
	ap.pendingPromptMetrics[inputID] = promptMetricsBoundary{compiled: compiled, rendered: rendered}
	ap.mu.Unlock()
}

func (ap *Process) appendDeliveryPromptCompileBlock(inputID string, block promptcompiler.PromptBlock) {
	inputID = strings.TrimSpace(inputID)
	if inputID == "" || !metrics.DetailedTokenCollectionEnabled() || block.TokenEstimate <= 0 {
		return
	}
	ap.mu.Lock()
	boundary, ok := ap.pendingPromptMetrics[inputID]
	if !ok {
		ap.mu.Unlock()
		return
	}
	for index := range boundary.rendered.AttributionBlocks {
		if boundary.rendered.AttributionBlocks[index].ID == block.ID {
			boundary.rendered.AttributionBlocks[index] = block
			ap.pendingPromptMetrics[inputID] = boundary
			ap.mu.Unlock()
			return
		}
	}
	boundary.rendered.AttributionBlocks = append(boundary.rendered.AttributionBlocks, block)
	ap.pendingPromptMetrics[inputID] = boundary
	ap.mu.Unlock()
}

func (ap *Process) commitDeliveryPromptCompileMetrics(inputID string) {
	inputID = strings.TrimSpace(inputID)
	if inputID == "" {
		return
	}
	ap.mu.Lock()
	boundary, ok := ap.pendingPromptMetrics[inputID]
	delete(ap.pendingPromptMetrics, inputID)
	if len(ap.pendingPromptMetrics) == 0 {
		ap.pendingPromptMetrics = nil
	}
	ap.mu.Unlock()
	if ok {
		ap.recordDeliveryPromptCompileMetrics(boundary.compiled, boundary.rendered)
		ap.activatePromptBoundary(boundary)
	}
}

func (ap *Process) discardDeliveryPromptCompileMetrics(inputID string) {
	inputID = strings.TrimSpace(inputID)
	ap.mu.Lock()
	delete(ap.pendingPromptMetrics, inputID)
	if len(ap.pendingPromptMetrics) == 0 {
		ap.pendingPromptMetrics = nil
	}
	ap.mu.Unlock()
}

func (ap *Process) resetActivePromptMetricsLocked(items []metrics.PromptCompileRecord) {
	ap.activePromptMetrics = make(map[string]metrics.InvocationExposureBlock)
	ap.contextExposureComplete = true
	blocks := make([]metrics.InvocationExposureBlock, 0, len(items))
	for _, item := range items {
		blocks = append(blocks, metrics.InvocationExposureBlock{
			BlockID:    strings.TrimSpace(item.BlockID),
			Durability: strings.TrimSpace(item.Durability),
			Tokens:     item.Tokens,
		})
	}
	ap.activatePromptMetricsLocked(blocks)
	ap.promptMetricsBaselineOK = ap.contextExposureComplete
}

func (ap *Process) activatePromptBoundary(boundary promptMetricsBoundary) {
	blocks := boundary.rendered.AttributionBlocks
	if len(blocks) == 0 && boundary.compiled.Bundle != nil {
		blocks = boundary.compiled.Bundle.Blocks
	}
	exposures := make([]metrics.InvocationExposureBlock, 0, len(blocks))
	for _, block := range blocks {
		exposures = append(exposures, metrics.InvocationExposureBlock{
			BlockID:    promptExposureBlockID(block.ID),
			Durability: string(block.Durability),
			Tokens:     uint64(max(block.TokenEstimate, 0)),
		})
	}
	ap.mu.Lock()
	ap.contextExposureComplete = ap.promptMetricsBaselineOK
	ap.removePromptMetricsByDurabilityLocked(string(promptcompiler.DurabilityTurnOnly))
	ap.activatePromptMetricsLocked(exposures)
	ap.mu.Unlock()
}

func (ap *Process) activatePromptMetricsLocked(blocks []metrics.InvocationExposureBlock) {
	if ap.activePromptMetrics == nil {
		ap.activePromptMetrics = make(map[string]metrics.InvocationExposureBlock)
	}
	for _, block := range blocks {
		block.BlockID = strings.TrimSpace(block.BlockID)
		block.Durability = strings.TrimSpace(block.Durability)
		if block.BlockID == "" || block.Durability == "" {
			ap.contextExposureComplete = false
			continue
		}
		switch block.Durability {
		case string(promptcompiler.DurabilityPersistent),
			string(promptcompiler.DurabilityTurnOnly),
			string(promptcompiler.DurabilityUntilAcknowledged),
			string(promptcompiler.DurabilityCompactable):
		default:
			ap.contextExposureComplete = false
			continue
		}
		ap.activePromptMetrics[block.BlockID+"\x00"+block.Durability] = block
	}
}

func (ap *Process) invocationExposureLocked(ev drivercontract.ParsedEvent) *metrics.InvocationExposureRecord {
	if ev.ContextCompacted {
		ap.removeCompactablePromptMetricsLocked()
	}
	invocationID := strings.TrimSpace(ev.ModelInvocationID)
	if invocationID == "" || ev.ContextExposureExcluded {
		return nil
	}
	complete := ev.ContextExposureKnown && ap.contextExposureComplete
	blocks := make([]metrics.InvocationExposureBlock, 0, len(ap.activePromptMetrics))
	if ev.ContextExposureKnown {
		for _, block := range ap.activePromptMetrics {
			blocks = append(blocks, block)
		}
	}
	sort.Slice(blocks, func(i, j int) bool {
		if blocks[i].BlockID != blocks[j].BlockID {
			return blocks[i].BlockID < blocks[j].BlockID
		}
		return blocks[i].Durability < blocks[j].Durability
	})
	return &metrics.InvocationExposureRecord{
		Complete: complete,
		Blocks:   blocks,
	}
}

func (ap *Process) removeCompactablePromptMetricsLocked() {
	ap.removePromptMetricsByDurabilityLocked(string(promptcompiler.DurabilityCompactable))
}

func (ap *Process) removePromptMetricsByDurabilityLocked(durabilities ...string) {
	for key, block := range ap.activePromptMetrics {
		for _, durability := range durabilities {
			if block.Durability == durability {
				delete(ap.activePromptMetrics, key)
				break
			}
		}
	}
}

// InvalidateTokenAnalysisExposure prevents a newly enabled collection window
// from claiming Context completeness until the next fully observed boundary.
func (ap *Process) InvalidateTokenAnalysisExposure() {
	if ap == nil {
		return
	}
	ap.mu.Lock()
	ap.contextExposureComplete = false
	ap.mu.Unlock()
}

func promptExposureBlockID(blockID string) string {
	return strings.TrimSpace(blockID)
}

func (ap *Process) recordTokenSource(
	ev drivercontract.ParsedEvent,
	agentID string,
	model string,
	driverName string,
) {
	consumption := ap.consumeTokenSource(ev.SessionID, ev.ToolCallID)
	ap.mu.Lock()
	recorder := ap.tokenSourceRecorder
	ap.mu.Unlock()
	base := TokenSourceObservation{
		Timestamp:         time.Now().UnixMilli(),
		AgentID:           strings.TrimSpace(agentID),
		Runtime:           strings.TrimSpace(driverName),
		Model:             strings.TrimSpace(model),
		InputTokens:       ev.InputTokens,
		CachedInputTokens: ev.CachedInputTokens,
		OutputTokens:      ev.OutputTokens,
	}
	allocations := allocateTokenSourceObservations(base, consumption)
	ap.recordTokenSourceConsumptionAudit(driverName, consumption, allocations)
	if recorder != nil {
		for _, allocation := range allocations {
			recorder(allocation)
		}
	}
}

func (ap *Process) usageMetricDimensions(ev drivercontract.ParsedEvent) (string, string, bool) {
	agentID, apiKey, model, driverName := ap.usageMetricContext(ev)
	if apiKey == "" && strings.TrimSpace(driverName) == "codex" && strings.TrimSpace(agentID) != "" {
		apiKey = "codex:" + strings.TrimSpace(agentID)
	}
	if model == "" || apiKey == "" {
		return "", "", false
	}
	return apiKey, model, true
}

func (ap *Process) usageMetricContext(ev drivercontract.ParsedEvent) (agentID, apiKey, model, driverName string) {
	ap.mu.Lock()
	model = strings.TrimSpace(ev.Model)
	eventSessionID := strings.TrimSpace(ev.SessionID)
	primarySessionID := strings.TrimSpace(ap.SessionID)
	if model == "" && (eventSessionID == "" || (primarySessionID != "" && eventSessionID == primarySessionID)) {
		model = ap.model
	}
	apiKey = ap.providerAPIKey
	agentID = ap.AgentID
	if ap.driver != nil {
		driverName = ap.driver.Name()
	}
	ap.mu.Unlock()

	model = strings.TrimSpace(model)
	apiKey = strings.TrimSpace(apiKey)
	agentID = strings.TrimSpace(agentID)
	driverName = strings.TrimSpace(driverName)
	return agentID, apiKey, model, driverName
}