package metrics

import (
	"context"
	"log/slog"
	"strconv"
	"strings"

	"github.com/prometheus/prometheus/model/labels"
)

// BatchRequestItem represents a single request for batch TSDB writes.
type BatchRequestItem struct {
	KeyID            string
	Model            string
	Timestamp        int64
	InputToken       uint64
	CachedInputToken uint64
	OutputToken      uint64
	DurationMs       uint64
}

// BatchErrorItem represents a single error for batch TSDB writes.
type BatchErrorItem struct {
	Model     string
	ErrorType ErrorType
	Timestamp int64
}

// BatchRecordRequests writes multiple request records to TSDB in a single commit.
// It ensures unique timestamps per series (key_id + model) by incrementing
// duplicate timestamps by 1ms, since Prometheus TSDB requires strictly increasing
// timestamps within a single series.
func (tc *TSDBCollector) BatchRecordRequests(items []BatchRequestItem) {
	if tc.db == nil || len(items) == 0 {
		return
	}

	// Track the last timestamp used per series to avoid duplicates.
	type seriesKey struct{ keyID, model string }
	lastTS := make(map[seriesKey]int64)

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	for _, item := range items {
		key := seriesKey{item.KeyID, item.Model}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts

		baseLabels := labels.FromStrings(labelKeyID, item.KeyID, labelModel, item.Model)

		if _, err := app.Append(0, buildLabels(baseLabels, metricRequestsTotal), ts, 1.0); err != nil {
			slog.Warn("Failed to append batch request count", "error", err)
		}
		if _, err := app.Append(0, buildLabels(baseLabels, metricInputTokensTotal), ts, float64(item.InputToken)); err != nil {
			slog.Warn("Failed to append batch input tokens", "error", err)
		}
		if _, err := app.Append(0, buildLabels(baseLabels, metricCachedInputTokensTotal), ts, float64(item.CachedInputToken)); err != nil {
			slog.Warn("Failed to append batch cached input tokens", "error", err)
		}
		if item.OutputToken > 0 {
			if _, err := app.Append(0, buildLabels(baseLabels, metricOutputTokensTotal), ts, float64(item.OutputToken)); err != nil {
				slog.Warn("Failed to append batch output tokens", "error", err)
			}
		}
		if item.DurationMs > 0 {
			if _, err := app.Append(0, buildLabels(baseLabels, metricRequestDurationMs), ts, float64(item.DurationMs)); err != nil {
				slog.Warn("Failed to append batch duration", "error", err)
			}
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit batch stats to TSDB", "error", err, "count", len(items))
	}
}

// BatchRecordErrors writes multiple error records to TSDB in a single commit.
// It ensures unique timestamps per series (model + error_type) by incrementing
// duplicate timestamps by 1ms.
func (tc *TSDBCollector) BatchRecordErrors(items []BatchErrorItem) {
	if tc.db == nil || len(items) == 0 {
		return
	}

	type seriesKey struct {
		model     string
		errorType ErrorType
	}
	lastTS := make(map[seriesKey]int64)

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	for _, item := range items {
		key := seriesKey{item.Model, item.ErrorType}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts

		errorLabels := labels.FromStrings(
			"__name__", metricErrorsTotal,
			labelModel, item.Model,
			labelErrorType, string(item.ErrorType),
		)
		if _, err := app.Append(0, errorLabels, ts, 1.0); err != nil {
			slog.Warn("Failed to append batch error count", "error", err)
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit batch errors to TSDB", "error", err, "count", len(items))
	}
}

// ServerMessageItem represents a single server message for batch TSDB writes.
type ServerMessageItem struct {
	ServerID  string
	Direction string // "received" or "sent"; defaults to "received" if empty
	Timestamp int64
}

// AgentMessageItem represents a single per-agent message for batch TSDB writes.
type AgentMessageItem struct {
	AgentID   string
	Direction string // "received" or "sent"; defaults to "received" if empty
	Timestamp int64
}

// AgentTokenItem represents one per-agent token usage event for batch TSDB writes.
type AgentTokenItem struct {
	AgentID          string
	Model            string
	Timestamp        int64
	InputToken       uint64
	CachedInputToken uint64
	OutputToken      uint64
}

// BatchRecordServerMessages writes multiple server message records to TSDB in a single commit.
// It ensures unique timestamps per series (server_id + direction) by incrementing
// duplicate timestamps by 1ms, since Prometheus TSDB requires strictly increasing
// timestamps within a single series.
func (tc *TSDBCollector) BatchRecordServerMessages(items []ServerMessageItem) {
	if tc.db == nil || len(items) == 0 {
		return
	}

	// Track the last timestamp used per series to avoid duplicates.
	// TSDB deduplicates samples with the same timestamp within a series.
	type seriesKey struct{ id, dir string }
	lastTS := make(map[seriesKey]int64)

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	for _, item := range items {
		dir := item.Direction
		if dir == "" {
			dir = "received"
		}
		key := seriesKey{item.ServerID, dir}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts

		messageLabels := labels.FromStrings(
			"__name__", metricServerMessagesTotal,
			labelServerID, item.ServerID,
			labelMessageDirection, dir,
		)
		if _, err := app.Append(0, messageLabels, ts, 1.0); err != nil {
			slog.Warn("Failed to append batch server message", "error", err)
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit batch server messages to TSDB", "error", err, "count", len(items))
	}
}

// BatchRecordAgentMessages writes multiple per-agent message records to TSDB in a single commit.
func (tc *TSDBCollector) BatchRecordAgentMessages(items []AgentMessageItem) {
	if tc.db == nil || len(items) == 0 {
		return
	}

	type seriesKey struct{ id, dir string }
	lastTS := make(map[seriesKey]int64)

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	for _, item := range items {
		if item.AgentID == "" {
			continue
		}
		dir := item.Direction
		if dir == "" {
			dir = DirectionReceived
		}
		key := seriesKey{item.AgentID, dir}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts

		agentLabels := labels.FromStrings(
			"__name__", metricAgentMessagesTotal,
			labelAgentID, item.AgentID,
			labelMessageDirection, dir,
		)
		if _, err := app.Append(0, agentLabels, ts, 1.0); err != nil {
			slog.Warn("Failed to append batch agent message", "error", err)
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit batch agent messages to TSDB", "error", err, "count", len(items))
	}
}

// BatchRecordAgentTokenUsage writes multiple per-agent token usage records to TSDB in a single commit.
func (tc *TSDBCollector) BatchRecordAgentTokenUsage(items []AgentTokenItem) {
	if tc.db == nil || len(items) == 0 {
		return
	}

	type seriesKey struct{ agentID, model string }
	lastTS := make(map[seriesKey]int64)

	ctx := context.Background()
	app := tc.db.Appender(ctx)

	for _, item := range items {
		if item.AgentID == "" || item.Model == "" {
			continue
		}
		key := seriesKey{item.AgentID, item.Model}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts

		baseLabels := labels.FromStrings(labelAgentID, item.AgentID, labelModel, item.Model)
		if _, err := app.Append(0, buildAgentTokenLabels(baseLabels, metricAgentTokenRequests), ts, 1.0); err != nil {
			slog.Warn("Failed to append batch agent token request", "error", err)
		}
		if _, err := app.Append(0, buildAgentTokenLabels(baseLabels, metricAgentInputTokens), ts, float64(item.InputToken)); err != nil {
			slog.Warn("Failed to append batch agent input tokens", "error", err)
		}
		if _, err := app.Append(0, buildAgentTokenLabels(baseLabels, metricAgentCachedInputTokens), ts, float64(item.CachedInputToken)); err != nil {
			slog.Warn("Failed to append batch agent cached input tokens", "error", err)
		}
		if item.OutputToken > 0 {
			if _, err := app.Append(0, buildAgentTokenLabels(baseLabels, metricAgentOutputTokens), ts, float64(item.OutputToken)); err != nil {
				slog.Warn("Failed to append batch agent output tokens", "error", err)
			}
		}
	}

	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit batch agent token usage to TSDB", "error", err, "count", len(items))
	}
}

func (tc *TSDBCollector) BatchRecordPromptCompileMetrics(items []PromptCompileRecord) {
	if tc.db == nil || len(items) == 0 {
		return
	}
	tc.mu.Lock()
	defer tc.mu.Unlock()
	ctx := context.Background()
	app := tc.db.Appender(ctx)
	type seriesKey struct {
		agentID, runtime, launchID, launchStartedAt, sessionKind, freshness, leaderSessionKey     string
		eventID, compileEvent, stage, slot, blockID, blockLayer, blockSlot, blockType, durability string
		taskPayloadPresent                                                                        bool
		taskPayloadFP                                                                             string
	}
	lastTS := make(map[seriesKey]int64)
	for _, item := range items {
		if item.AgentID == "" || item.Runtime == "" || item.LaunchID == "" || item.SessionKind == "" || item.Freshness == "" || item.CompileEvent == "" || item.Stage == "" || item.Slot == "" {
			continue
		}
		key := seriesKey{item.AgentID, item.Runtime, item.LaunchID, strconv.FormatInt(item.LaunchStartedAt, 10), item.SessionKind, item.Freshness, item.LeaderSessionKey, item.EventID, item.CompileEvent, item.Stage, item.Slot, item.BlockID, item.BlockLayer, item.BlockSlot, item.BlockType, item.Durability, item.TaskPayloadPresent, item.TaskPayloadFP}
		ts := item.Timestamp
		if prev, ok := lastTS[key]; ok && ts <= prev {
			ts = prev + 1
		}
		lastTS[key] = ts
		lbls := labels.FromStrings(
			labelAgentID, item.AgentID,
			labelRuntime, item.Runtime,
			labelLaunchID, item.LaunchID,
			labelLaunchStartedAt, strconv.FormatInt(item.LaunchStartedAt, 10),
			labelSessionKind, item.SessionKind,
			labelSessionFreshness, item.Freshness,
			labelLeaderSessionKey, item.LeaderSessionKey,
			labelEventID, item.EventID,
			labelCompileEvent, item.CompileEvent,
			labelStage, item.Stage,
			labelSlotKey, item.Slot,
			labelBlockID, item.BlockID,
			labelBlockLayer, item.BlockLayer,
			labelBlockSlot, item.BlockSlot,
			labelBlockType, item.BlockType,
			labelDurability, item.Durability,
			labelTaskPayloadPresent, strconv.FormatBool(item.TaskPayloadPresent),
			labelTaskPayloadFP, item.TaskPayloadFP,
		)
		records := item.Records
		if records == 0 {
			records = 1
		}
		if _, err := app.Append(0, buildPromptCompileLabels(lbls, metricPromptCompileRecords), ts, float64(records)); err != nil {
			slog.Warn("Failed to append prompt compile record", "error", err)
		}
		if _, err := app.Append(0, buildPromptCompileLabels(lbls, metricPromptCompileTokens), ts, float64(item.Tokens)); err != nil {
			slog.Warn("Failed to append prompt compile tokens", "error", err)
		}
	}
	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit prompt compile metrics to TSDB", "error", err, "count", len(items))
	}
}

func (tc *TSDBCollector) recordUsageObservation(
	item UsageObservationRecord,
	exposure *InvocationExposureRecord,
) {
	if tc.db == nil {
		return
	}
	tc.mu.Lock()
	defer tc.mu.Unlock()
	base := labels.FromStrings(
		labelObservationID, item.ObservationID,
		labelAgentID, item.AgentID,
		labelRuntime, item.Runtime,
		labelLaunchID, item.LaunchID,
		labelLaunchStartedAt, strconv.FormatInt(item.LaunchStartedAt, 10),
		labelSessionKind, item.SessionKind,
		labelSessionID, item.SessionID,
		labelModel, item.Model,
		labelInvocationID, item.InvocationID,
		labelInvocationCause, item.InvocationCause,
		labelExposureExcluded, strconv.FormatBool(item.ExposureExcluded),
	)
	app := tc.db.Appender(context.Background())
	for _, metric := range []struct {
		name  string
		value uint64
	}{
		{metricTokenUsageObservation, 1},
		{metricTokenUsageInput, item.InputTokens},
		{metricTokenUsageCached, item.CachedTokens},
		{metricTokenUsageOutput, item.OutputTokens},
	} {
		if _, err := app.Append(0, buildUsageObservationLabels(base, metric.name), item.Timestamp, float64(metric.value)); err != nil {
			slog.Warn("Failed to append token analysis usage observation", "error", err)
			_ = app.Rollback()
			return
		}
	}
	if exposure != nil {
		exposureLabels := labels.FromStrings(
			labelAgentID, item.AgentID,
			labelRuntime, item.Runtime,
			labelLaunchID, item.LaunchID,
			labelLaunchStartedAt, strconv.FormatInt(item.LaunchStartedAt, 10),
			labelSessionKind, item.SessionKind,
			labelSessionID, item.SessionID,
			labelInvocationID, item.InvocationID,
			labelExposureComplete, strconv.FormatBool(exposure.Complete),
		)
		if _, err := app.Append(
			0,
			buildInvocationExposureLabels(exposureLabels, metricInvocationExposure, "", ""),
			item.Timestamp,
			1,
		); err != nil {
			slog.Warn("Failed to append token analysis invocation exposure", "error", err)
			_ = app.Rollback()
			return
		}
		for _, block := range exposure.Blocks {
			blockID := strings.TrimSpace(block.BlockID)
			durability := strings.TrimSpace(block.Durability)
			if blockID == "" || durability == "" {
				continue
			}
			if _, err := app.Append(
				0,
				buildInvocationExposureLabels(exposureLabels, metricInvocationBlockTokens, blockID, durability),
				item.Timestamp,
				float64(block.Tokens),
			); err != nil {
				slog.Warn("Failed to append token analysis block exposure", "error", err)
				_ = app.Rollback()
				return
			}
		}
	}
	if err := app.Commit(); err != nil {
		slog.Warn("Failed to commit token analysis usage observation", "error", err)
	}
}

func buildPromptCompileLabels(base labels.Labels, metricName string) labels.Labels {
	return labels.FromStrings(
		"__name__", metricName,
		labelAgentID, base.Get(labelAgentID),
		labelRuntime, base.Get(labelRuntime),
		labelLaunchID, base.Get(labelLaunchID),
		labelLaunchStartedAt, base.Get(labelLaunchStartedAt),
		labelSessionKind, base.Get(labelSessionKind),
		labelSessionFreshness, base.Get(labelSessionFreshness),
		labelLeaderSessionKey, base.Get(labelLeaderSessionKey),
		labelEventID, base.Get(labelEventID),
		labelCompileEvent, base.Get(labelCompileEvent),
		labelStage, base.Get(labelStage),
		labelSlotKey, base.Get(labelSlotKey),
		labelBlockID, base.Get(labelBlockID),
		labelBlockLayer, base.Get(labelBlockLayer),
		labelBlockSlot, base.Get(labelBlockSlot),
		labelBlockType, base.Get(labelBlockType),
		labelDurability, base.Get(labelDurability),
		labelTaskPayloadPresent, base.Get(labelTaskPayloadPresent),
		labelTaskPayloadFP, base.Get(labelTaskPayloadFP),
	)
}

func buildUsageObservationLabels(base labels.Labels, metricName string) labels.Labels {
	return labels.FromStrings(
		"__name__", metricName,
		labelObservationID, base.Get(labelObservationID),
		labelAgentID, base.Get(labelAgentID),
		labelRuntime, base.Get(labelRuntime),
		labelLaunchID, base.Get(labelLaunchID),
		labelLaunchStartedAt, base.Get(labelLaunchStartedAt),
		labelSessionKind, base.Get(labelSessionKind),
		labelSessionID, base.Get(labelSessionID),
		labelModel, base.Get(labelModel),
		labelInvocationID, base.Get(labelInvocationID),
		labelInvocationCause, base.Get(labelInvocationCause),
		labelExposureExcluded, base.Get(labelExposureExcluded),
	)
}

func buildInvocationExposureLabels(base labels.Labels, metricName, blockID, durability string) labels.Labels {
	return labels.FromStrings(
		"__name__", metricName,
		labelAgentID, base.Get(labelAgentID),
		labelRuntime, base.Get(labelRuntime),
		labelLaunchID, base.Get(labelLaunchID),
		labelLaunchStartedAt, base.Get(labelLaunchStartedAt),
		labelSessionKind, base.Get(labelSessionKind),
		labelSessionID, base.Get(labelSessionID),
		labelInvocationID, base.Get(labelInvocationID),
		labelExposureComplete, base.Get(labelExposureComplete),
		labelBlockID, blockID,
		labelDurability, durability,
	)
}