package dispatcher

import (
	"strings"
	"time"

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

func promptCompileRecords(agentID, runtime, launchID, sessionKind, freshness, stage string, compiled promptcompiler.RuntimeInputEnvelope, rendered promptcompiler.RenderedPromptPayload) []metrics.PromptCompileRecord {
	agentID = strings.TrimSpace(agentID)
	runtime = strings.TrimSpace(runtime)
	launchID = strings.TrimSpace(launchID)
	sessionKind = strings.TrimSpace(sessionKind)
	freshness = strings.TrimSpace(freshness)
	stage = strings.TrimSpace(stage)
	if agentID == "" || runtime == "" || launchID == "" || sessionKind == "" || freshness == "" || stage == "" {
		return nil
	}
	timestamp := time.Now().UnixMilli()
	blocks := promptAttributionBlocks(compiled, rendered)
	items := make([]metrics.PromptCompileRecord, 0, len(blocks))
	if len(blocks) > 0 {
		taskPayloadFP := token.Default().FingerprintTokenAnalysisPayload([]byte(compiled.TaskPayloadDigest))
		for _, block := range blocks {
			items = append(items, metrics.PromptCompileRecord{
				Timestamp:          timestamp,
				AgentID:            agentID,
				Runtime:            runtime,
				LaunchID:           launchID,
				SessionKind:        sessionKind,
				Freshness:          freshness,
				EventID:            strings.TrimSpace(compiled.EventID),
				CompileEvent:       string(compiled.EventType),
				Stage:              stage,
				Slot:               string(block.Slot),
				BlockID:            promptMetricsBlockID(block),
				BlockLayer:         string(block.Layer),
				BlockSlot:          string(block.Slot),
				BlockType:          string(block.Type),
				Durability:         string(block.Durability),
				TaskPayloadPresent: compiled.TaskPayloadPresent,
				TaskPayloadFP:      taskPayloadFP,
				Tokens:             uint64(block.TokenEstimate),
			})
		}
	}
	if len(items) == 0 {
		for _, surface := range promptcompiler.RenderedSurfaceAttributions(rendered) {
			items = append(items, metrics.PromptCompileRecord{
				Timestamp:    timestamp,
				AgentID:      agentID,
				Runtime:      runtime,
				LaunchID:     launchID,
				SessionKind:  sessionKind,
				Freshness:    freshness,
				EventID:      strings.TrimSpace(compiled.EventID),
				CompileEvent: string(compiled.EventType),
				Stage:        stage,
				Slot:         surface.Slot,
				Tokens:       uint64(surface.Tokens),
			})
		}
	}
	return items
}

func promptMetricsBlockID(block promptcompiler.PromptBlock) string {
	return strings.TrimSpace(block.ID)
}

func promptAttributionBlocks(compiled promptcompiler.RuntimeInputEnvelope, rendered promptcompiler.RenderedPromptPayload) []promptcompiler.PromptBlock {
	if len(rendered.AttributionBlocks) > 0 {
		return rendered.AttributionBlocks
	}
	if compiled.Bundle != nil {
		return compiled.Bundle.Blocks
	}
	return nil
}