package metrics

import (
	"context"
	"log/slog"
	"sort"
	"time"

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

// ServerMessageTSDBRecord represents a raw server message record from TSDB.
type ServerMessageTSDBRecord struct {
	Timestamp int64
	ServerID  string
	Direction string // "received" or "sent"; empty treated as "received" for legacy data
}

// AgentTSDBRecord represents a raw per-agent message record from TSDB.
type AgentTSDBRecord struct {
	Timestamp int64
	AgentID   string
	Direction string // "received" or "sent"; empty treated as "received" for legacy data
}

// QueryServerMessages queries server message records from TSDB within the given time window.
func (tc *TSDBCollector) QueryServerMessages(window TimeWindow) []ServerMessageTSDBRecord {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return nil
	}

	now := time.Now()
	mint := now.Add(-window.Duration()).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for server messages", "error", err)
		return nil
	}
	defer querier.Close()

	matcher := labels.MustNewMatcher(labels.MatchEqual, "__name__", metricServerMessagesTotal)
	seriesSet := querier.Select(ctx, false, nil, matcher)

	var records []ServerMessageTSDBRecord
	for seriesSet.Next() {
		series := seriesSet.At()
		lbls := series.Labels()
		serverID := lbls.Get(labelServerID)
		direction := lbls.Get(labelMessageDirection)
		if serverID == "" {
			continue
		}

		iterator := series.Iterator(nil)
		for iterator.Next() == chunkenc.ValFloat {
			t, _ := iterator.At()
			records = append(records, ServerMessageTSDBRecord{
				Timestamp: t,
				ServerID:  serverID,
				Direction: direction,
			})
		}

		if err := iterator.Err(); err != nil {
			slog.Warn("Error iterating server message samples", "error", err)
		}
	}

	if err := seriesSet.Err(); err != nil {
		slog.Warn("Error iterating server message series", "error", err)
	}

	return records
}

// QueryAgentMessages queries per-agent message records from TSDB within the given time window.
func (tc *TSDBCollector) QueryAgentMessages(window TimeWindow) []AgentTSDBRecord {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return nil
	}

	now := time.Now()
	mint := now.Add(-window.Duration()).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for agent messages", "error", err)
		return nil
	}
	defer querier.Close()

	matcher := labels.MustNewMatcher(labels.MatchEqual, "__name__", metricAgentMessagesTotal)
	seriesSet := querier.Select(ctx, false, nil, matcher)

	var records []AgentTSDBRecord
	for seriesSet.Next() {
		series := seriesSet.At()
		lbls := series.Labels()
		agentID := lbls.Get(labelAgentID)
		direction := lbls.Get(labelMessageDirection)
		if agentID == "" {
			continue
		}

		iterator := series.Iterator(nil)
		for iterator.Next() == chunkenc.ValFloat {
			t, _ := iterator.At()
			records = append(records, AgentTSDBRecord{
				Timestamp: t,
				AgentID:   agentID,
				Direction: direction,
			})
		}

		if err := iterator.Err(); err != nil {
			slog.Warn("Error iterating agent message samples", "error", err)
		}
	}

	if err := seriesSet.Err(); err != nil {
		slog.Warn("Error iterating agent message series", "error", err)
	}

	return records
}

// QueryAgentTokenUsage queries per-agent token usage records from TSDB within the given time window.
func (tc *TSDBCollector) QueryAgentTokenUsage(window TimeWindow) []AgentTokenItem {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	if tc.db == nil {
		return nil
	}

	now := time.Now()
	mint := now.Add(-window.Duration()).UnixMilli()
	maxt := now.UnixMilli()

	ctx := context.Background()
	querier, err := tc.db.Querier(mint, maxt)
	if err != nil {
		slog.Warn("Failed to create querier for agent token usage", "error", err)
		return nil
	}
	defer querier.Close()

	type recordKey struct {
		timestamp int64
		agentID   string
		model     string
	}
	recordsByKey := make(map[recordKey]*AgentTokenItem)

	mergeMetric := func(metricName string, apply func(*AgentTokenItem, uint64)) {
		matcher := labels.MustNewMatcher(labels.MatchEqual, "__name__", metricName)
		seriesSet := querier.Select(ctx, false, nil, matcher)
		for seriesSet.Next() {
			series := seriesSet.At()
			lbls := series.Labels()
			agentID := lbls.Get(labelAgentID)
			model := lbls.Get(labelModel)
			if agentID == "" || model == "" {
				continue
			}

			iterator := series.Iterator(nil)
			for iterator.Next() == chunkenc.ValFloat {
				t, v := iterator.At()
				key := recordKey{timestamp: t, agentID: agentID, model: model}
				record := recordsByKey[key]
				if record == nil {
					record = &AgentTokenItem{
						Timestamp: t,
						AgentID:   agentID,
						Model:     model,
					}
					recordsByKey[key] = record
				}
				apply(record, uint64(v))
			}

			if err := iterator.Err(); err != nil {
				slog.Warn("Error iterating agent token samples", "metric", metricName, "error", err)
			}
		}
		if err := seriesSet.Err(); err != nil {
			slog.Warn("Error iterating agent token series", "metric", metricName, "error", err)
		}
	}

	mergeMetric(metricAgentTokenRequests, func(record *AgentTokenItem, value uint64) {
		if value > 0 && record.InputToken == 0 && record.CachedInputToken == 0 && record.OutputToken == 0 {
			// Request count is represented by the presence of this record. The
			// value itself is intentionally not copied into token fields.
		}
	})
	mergeMetric(metricAgentInputTokens, func(record *AgentTokenItem, value uint64) {
		record.InputToken += value
	})
	mergeMetric(metricAgentCachedInputTokens, func(record *AgentTokenItem, value uint64) {
		record.CachedInputToken += value
	})
	mergeMetric(metricAgentOutputTokens, func(record *AgentTokenItem, value uint64) {
		record.OutputToken += value
	})

	records := make([]AgentTokenItem, 0, len(recordsByKey))
	for _, record := range recordsByKey {
		records = append(records, *record)
	}
	sort.Slice(records, func(i, j int) bool {
		if records[i].Timestamp == records[j].Timestamp {
			if records[i].AgentID == records[j].AgentID {
				return records[i].Model < records[j].Model
			}
			return records[i].AgentID < records[j].AgentID
		}
		return records[i].Timestamp < records[j].Timestamp
	})
	return records
}