package metrics

import (
	"context"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

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

type promptCompileQuery struct {
	Window     TimeWindow
	AgentIDs   []string
	LaunchID   string
	SnapshotAt int64
}

type promptCompileSampleKey struct {
	Timestamp          int64
	AgentID            string
	Runtime            string
	LaunchID           string
	LaunchStartedAt    int64
	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
}

type promptCompileAccumulator struct {
	record   PromptCompileRecord
	presence uint8
}

func (tc *TSDBCollector) queryPromptCompileRecordsBounded(query promptCompileQuery, sampleLimit int) ([]PromptCompileRecord, int, bool) {
	tc.mu.RLock()
	defer tc.mu.RUnlock()
	return tc.queryPromptCompileRecordsBoundedLocked(query, sampleLimit)
}

func (tc *TSDBCollector) queryPromptCompileRecordsBoundedLocked(query promptCompileQuery, sampleLimit int) ([]PromptCompileRecord, int, bool) {
	promptDB := tc.promptQueryable()
	if promptDB == nil {
		return nil, 0, false
	}
	maxt := query.SnapshotAt
	if maxt <= 0 {
		maxt = time.Now().UnixMilli()
	}
	querier, err := promptDB.Querier(time.UnixMilli(maxt).Add(-query.Window.Duration()).UnixMilli(), maxt)
	if err != nil {
		return nil, 0, false
	}
	defer querier.Close()
	ctx := context.Background()
	recordsByKey := make(map[promptCompileSampleKey]*promptCompileAccumulator)
	scanned := 0
	truncated := false
	perMetricLimit := sampleLimit / 2
	if perMetricLimit == 0 && sampleLimit > 0 {
		perMetricLimit = 1
	}
	mergeMetric := func(metricName string, presence uint8, apply func(*PromptCompileRecord, uint64)) {
		metricSamples := 0
		seriesSet := querier.Select(ctx, false, nil, promptMetricMatchers(metricName, query)...)
		for seriesSet.Next() {
			series := seriesSet.At()
			lbls := series.Labels()
			it := series.Iterator(nil)
			for it.Next() == chunkenc.ValFloat {
				if perMetricLimit > 0 && metricSamples >= perMetricLimit {
					truncated = true
					break
				}
				metricSamples++
				scanned++
				ts, v := it.At()
				key := promptCompileKey(ts, lbls)
				accumulator := recordsByKey[key]
				if accumulator == nil {
					accumulator = &promptCompileAccumulator{record: promptCompileRecord(key)}
					recordsByKey[key] = accumulator
				}
				apply(&accumulator.record, uint64(v))
				accumulator.presence |= presence
			}
			if truncated && perMetricLimit > 0 && metricSamples >= perMetricLimit {
				break
			}
		}
	}
	mergeMetric(metricPromptCompileRecords, 1, func(record *PromptCompileRecord, value uint64) {
		record.Records += value
	})
	mergeMetric(metricPromptCompileTokens, 2, func(record *PromptCompileRecord, value uint64) {
		record.Tokens += value
	})
	out := make([]PromptCompileRecord, 0, len(recordsByKey))
	for _, accumulator := range recordsByKey {
		if accumulator.presence == 3 {
			out = append(out, accumulator.record)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Timestamp < out[j].Timestamp })
	return out, scanned, truncated
}

func promptMetricMatchers(metricName string, query promptCompileQuery) []*labels.Matcher {
	matchers := []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "__name__", metricName)}
	if query.LaunchID != "" {
		matchers = append(matchers, labels.MustNewMatcher(labels.MatchEqual, labelLaunchID, query.LaunchID))
	}
	switch len(query.AgentIDs) {
	case 0:
	case 1:
		matchers = append(matchers, labels.MustNewMatcher(labels.MatchEqual, labelAgentID, query.AgentIDs[0]))
	default:
		values := make([]string, len(query.AgentIDs))
		for i, value := range query.AgentIDs {
			values[i] = regexp.QuoteMeta(value)
		}
		matchers = append(matchers, labels.MustNewMatcher(labels.MatchRegexp, labelAgentID, strings.Join(values, "|")))
	}
	return matchers
}

func promptCompileKey(timestamp int64, lbls labels.Labels) promptCompileSampleKey {
	return promptCompileSampleKey{
		Timestamp:          timestamp,
		AgentID:            lbls.Get(labelAgentID),
		Runtime:            lbls.Get(labelRuntime),
		LaunchID:           lbls.Get(labelLaunchID),
		LaunchStartedAt:    parseTimestampLabel(lbls.Get(labelLaunchStartedAt)),
		SessionKind:        lbls.Get(labelSessionKind),
		Freshness:          lbls.Get(labelSessionFreshness),
		LeaderSessionKey:   lbls.Get(labelLeaderSessionKey),
		EventID:            lbls.Get(labelEventID),
		CompileEvent:       lbls.Get(labelCompileEvent),
		Stage:              lbls.Get(labelStage),
		Slot:               lbls.Get(labelSlotKey),
		BlockID:            lbls.Get(labelBlockID),
		BlockLayer:         lbls.Get(labelBlockLayer),
		BlockSlot:          lbls.Get(labelBlockSlot),
		BlockType:          lbls.Get(labelBlockType),
		Durability:         lbls.Get(labelDurability),
		TaskPayloadPresent: lbls.Get(labelTaskPayloadPresent) == "true",
		TaskPayloadFP:      lbls.Get(labelTaskPayloadFP),
	}
}

func promptCompileRecord(key promptCompileSampleKey) PromptCompileRecord {
	return PromptCompileRecord{
		Timestamp:          key.Timestamp,
		AgentID:            key.AgentID,
		Runtime:            key.Runtime,
		LaunchID:           key.LaunchID,
		LaunchStartedAt:    key.LaunchStartedAt,
		SessionKind:        key.SessionKind,
		Freshness:          key.Freshness,
		LeaderSessionKey:   key.LeaderSessionKey,
		EventID:            key.EventID,
		CompileEvent:       key.CompileEvent,
		Stage:              key.Stage,
		Slot:               key.Slot,
		BlockID:            key.BlockID,
		BlockLayer:         key.BlockLayer,
		BlockSlot:          key.BlockSlot,
		BlockType:          key.BlockType,
		Durability:         key.Durability,
		TaskPayloadPresent: key.TaskPayloadPresent,
		TaskPayloadFP:      key.TaskPayloadFP,
	}
}

func parseTimestampLabel(value string) int64 {
	timestamp, _ := strconv.ParseInt(strings.TrimSpace(value), 10, 64)
	return timestamp
}