package metrics

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

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

func (tc *TSDBCollector) queryTokenAnalysisEvidence(query TokenEvidenceQuery) TokenAnalysisEvidence {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	query = normalizeTokenEvidenceQuery(query)
	result := TokenAnalysisEvidence{
		PromptRecords:    []PromptCompileRecord{},
		Usage:            []UsageObservationRecord{},
		ExposureCoverage: []ExposureCoverageRecord{},
		BlockExposures:   []BlockExposureRecord{},
		Limit:            query.Limit,
	}
	if query.AgentIDs != nil && len(query.AgentIDs) == 0 {
		return result
	}
	promptLimit := query.Limit / 4
	usageLimit := query.Limit / 4
	exposureLimit := query.Limit - promptLimit - usageLimit
	promptQuery := promptCompileQuery{
		Window:     query.Window,
		AgentIDs:   query.AgentIDs,
		LaunchID:   query.LaunchID,
		SnapshotAt: query.SnapshotAt,
	}
	records, scanned, truncated := tc.queryPromptCompileRecordsBoundedLocked(promptQuery, promptLimit)
	result.PromptRecords = records
	result.Scanned += min(scanned, promptLimit)
	result.Truncated = truncated
	usage, scanned, truncated := tc.queryUsageObservationsBoundedLocked(query, usageLimit)
	result.Usage = usage
	result.Scanned += min(scanned, usageLimit)
	result.Truncated = result.Truncated || truncated
	coverage, blocks, scanned, truncated := tc.queryInvocationExposuresBoundedLocked(query, exposureLimit)
	result.ExposureCoverage = coverage
	result.BlockExposures = blocks
	result.Scanned += min(scanned, exposureLimit)
	result.Truncated = result.Truncated || truncated
	return result
}

func normalizeTokenEvidenceQuery(query TokenEvidenceQuery) TokenEvidenceQuery {
	if query.Window < TimeWindow6Hours || query.Window > TimeWindow1Week {
		query.Window = TimeWindow1Day
	}
	query.AgentIDs = normalizeTokenEvidenceValues(query.AgentIDs)
	query.LaunchID = strings.TrimSpace(query.LaunchID)
	if query.SnapshotAt <= 0 {
		query.SnapshotAt = time.Now().UnixMilli()
	}
	if query.Limit <= 0 || query.Limit > MaxTokenEvidenceSamples {
		query.Limit = MaxTokenEvidenceSamples
	}
	return query
}

func normalizeTokenEvidenceValues(values []string) []string {
	if values == nil {
		return nil
	}
	seen := make(map[string]struct{}, len(values))
	normalized := make([]string, 0, len(values))
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		normalized = append(normalized, value)
	}
	sort.Strings(normalized)
	return normalized
}

func (tc *TSDBCollector) queryUsageObservationsBoundedLocked(query TokenEvidenceQuery, sampleLimit int) ([]UsageObservationRecord, int, bool) {
	db := tc.promptQueryable()
	if db == nil || sampleLimit <= 0 {
		return nil, 0, false
	}
	querier, err := db.Querier(
		time.UnixMilli(query.SnapshotAt).Add(-query.Window.Duration()).UnixMilli(),
		query.SnapshotAt,
	)
	if err != nil {
		return nil, 0, false
	}
	defer querier.Close()

	type accumulator struct {
		record   UsageObservationRecord
		presence uint8
	}
	records := make(map[string]*accumulator)
	scanned := 0
	truncated := false
	merge := func(metric string, bit uint8, apply func(*UsageObservationRecord, uint64)) {
		matchers := tokenEvidenceMatchers(metric, query)
		set := querier.Select(context.Background(), false, nil, matchers...)
		for set.Next() {
			series := set.At()
			lbls := series.Labels()
			it := series.Iterator(nil)
			for it.Next() == chunkenc.ValFloat {
				ts, value := it.At()
				id := lbls.Get(labelObservationID)
				if id == "" {
					continue
				}
				item := records[id]
				if item == nil {
					if len(records) >= sampleLimit {
						truncated = true
						continue
					}
					item = &accumulator{record: UsageObservationRecord{
						Timestamp:        ts,
						ObservationID:    id,
						AgentID:          lbls.Get(labelAgentID),
						Runtime:          lbls.Get(labelRuntime),
						LaunchID:         lbls.Get(labelLaunchID),
						LaunchStartedAt:  parseTimestampLabel(lbls.Get(labelLaunchStartedAt)),
						SessionKind:      lbls.Get(labelSessionKind),
						SessionID:        lbls.Get(labelSessionID),
						Model:            lbls.Get(labelModel),
						InvocationID:     lbls.Get(labelInvocationID),
						InvocationCause:  lbls.Get(labelInvocationCause),
						ExposureExcluded: lbls.Get(labelExposureExcluded) == "true",
					}}
					records[id] = item
					scanned++
				}
				apply(&item.record, uint64(value))
				item.presence |= bit
			}
		}
		if set.Err() != nil {
			truncated = true
		}
	}
	merge(metricTokenUsageObservation, 1, func(*UsageObservationRecord, uint64) {})
	merge(metricTokenUsageInput, 2, func(record *UsageObservationRecord, value uint64) { record.InputTokens += value })
	merge(metricTokenUsageCached, 4, func(record *UsageObservationRecord, value uint64) { record.CachedTokens += value })
	merge(metricTokenUsageOutput, 8, func(record *UsageObservationRecord, value uint64) { record.OutputTokens += value })
	out := make([]UsageObservationRecord, 0, len(records))
	for _, item := range records {
		if item.presence == 15 {
			out = append(out, item.record)
		}
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Timestamp != out[j].Timestamp {
			return out[i].Timestamp < out[j].Timestamp
		}
		return out[i].ObservationID < out[j].ObservationID
	})
	return out, scanned, truncated
}

func (tc *TSDBCollector) queryInvocationExposuresBoundedLocked(
	query TokenEvidenceQuery,
	sampleLimit int,
) ([]ExposureCoverageRecord, []BlockExposureRecord, int, bool) {
	db := tc.promptQueryable()
	if db == nil || sampleLimit <= 0 {
		return nil, nil, 0, false
	}
	querier, err := db.Querier(
		time.UnixMilli(query.SnapshotAt).Add(-query.Window.Duration()).UnixMilli(),
		query.SnapshotAt,
	)
	if err != nil {
		return nil, nil, 0, false
	}
	defer querier.Close()

	type invocationCoverage struct {
		launchID string
		complete bool
	}
	invocations := make(map[string]invocationCoverage)
	blocks := make(map[string]*BlockExposureRecord)
	exposedBlocks := make(map[string]struct{})
	scanned := 0
	truncated := false
	selectMetric := func(metric string, apply func(labels.Labels, uint64)) {
		set := querier.Select(context.Background(), false, nil, tokenEvidenceMatchers(metric, query)...)
		for set.Next() {
			series := set.At()
			lbls := series.Labels()
			it := series.Iterator(nil)
			for it.Next() == chunkenc.ValFloat {
				if scanned >= sampleLimit {
					truncated = true
					break
				}
				_, value := it.At()
				scanned++
				apply(lbls, uint64(value))
			}
			if truncated {
				break
			}
		}
		if set.Err() != nil {
			truncated = true
		}
	}
	selectMetric(metricInvocationExposure, func(lbls labels.Labels, _ uint64) {
		launchID := lbls.Get(labelLaunchID)
		key := invocationExposureKey(lbls)
		complete := lbls.Get(labelExposureComplete) == "true"
		if current, ok := invocations[key]; ok {
			current.complete = current.complete && complete
			invocations[key] = current
			return
		}
		invocations[key] = invocationCoverage{launchID: launchID, complete: complete}
	})
	if !truncated {
		selectMetric(metricInvocationBlockTokens, func(lbls labels.Labels, value uint64) {
			launchID := lbls.Get(labelLaunchID)
			blockID := lbls.Get(labelBlockID)
			durability := lbls.Get(labelDurability)
			exposureKey := invocationExposureKey(lbls) + "\x00" + blockID + "\x00" + durability
			if _, ok := exposedBlocks[exposureKey]; ok {
				return
			}
			exposedBlocks[exposureKey] = struct{}{}
			key := launchID + "\x00" + blockID + "\x00" + durability
			record := blocks[key]
			if record == nil {
				record = &BlockExposureRecord{
					LaunchID:   launchID,
					BlockID:    blockID,
					Durability: durability,
				}
				blocks[key] = record
			}
			record.ExposureCount++
			record.Tokens += value
		})
	}
	coverage := make(map[string]*ExposureCoverageRecord)
	for _, invocation := range invocations {
		record := coverage[invocation.launchID]
		if record == nil {
			record = &ExposureCoverageRecord{LaunchID: invocation.launchID}
			coverage[invocation.launchID] = record
		}
		record.Observations++
		if invocation.complete {
			record.CompleteObservations++
		}
	}
	coverageOut := make([]ExposureCoverageRecord, 0, len(coverage))
	for _, record := range coverage {
		coverageOut = append(coverageOut, *record)
	}
	sort.Slice(coverageOut, func(i, j int) bool {
		return coverageOut[i].LaunchID < coverageOut[j].LaunchID
	})
	blockOut := make([]BlockExposureRecord, 0, len(blocks))
	for _, record := range blocks {
		blockOut = append(blockOut, *record)
	}
	sort.Slice(blockOut, func(i, j int) bool {
		if blockOut[i].LaunchID != blockOut[j].LaunchID {
			return blockOut[i].LaunchID < blockOut[j].LaunchID
		}
		if blockOut[i].BlockID != blockOut[j].BlockID {
			return blockOut[i].BlockID < blockOut[j].BlockID
		}
		return blockOut[i].Durability < blockOut[j].Durability
	})
	return coverageOut, blockOut, scanned, truncated
}

func invocationExposureKey(lbls labels.Labels) string {
	return lbls.Get(labelLaunchID) + "\x00" +
		lbls.Get(labelRuntime) + "\x00" +
		lbls.Get(labelSessionID) + "\x00" +
		lbls.Get(labelInvocationID)
}

func tokenEvidenceMatchers(metric string, query TokenEvidenceQuery) []*labels.Matcher {
	matchers := []*labels.Matcher{labels.MustNewMatcher(labels.MatchEqual, "__name__", metric)}
	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
}