package metrics

import (
	"context"
	"fmt"
	"log/slog"
	"math"
	"sync"
	"time"

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

const (
	metricTokenTrackingCalls  = "agentteam_token_tracking_calls_total"
	metricTokenTrackingInput  = "agentteam_token_tracking_input_tokens_total"
	metricTokenTrackingCached = "agentteam_token_tracking_cached_input_tokens_total"
	metricTokenTrackingOutput = "agentteam_token_tracking_output_tokens_total"

	labelTokenSourceKind = "source_kind"
	labelTokenSourceName = "source_name"

	tokenTrackingRetention = 7 * 24 * time.Hour
	tokenTrackingBlock     = 2 * time.Hour
)

// TokenTrackingStore owns the independent token-source TSDB.
type TokenTrackingStore struct {
	mu            sync.RWMutex
	db            *tsdb.DB
	path          string
	lastTimestamp int64
}

// NewTokenTrackingStore opens the standard local token-source database.
func NewTokenTrackingStore() (*TokenTrackingStore, error) {
	return openTokenTrackingStore(tokenTrackingPath, false)
}

func openTokenTrackingStore(path string, noLockfile bool) (*TokenTrackingStore, error) {
	opts := tsdb.DefaultOptions()
	opts.RetentionDuration = int64(tokenTrackingRetention / time.Millisecond)
	opts.MinBlockDuration = int64(tokenTrackingBlock / time.Millisecond)
	opts.MaxBlockDuration = int64(tokenTrackingBlock / time.Millisecond)
	opts.NoLockfile = noLockfile

	db, err := tsdb.Open(path, nil, nil, opts, nil)
	if err != nil {
		return nil, fmt.Errorf("open token tracking TSDB: %w", err)
	}
	lastTimestamp := db.Head().MaxTime()
	if lastTimestamp == math.MinInt64 {
		lastTimestamp = 0
	}
	return &TokenTrackingStore{
		db:            db,
		path:          path,
		lastTimestamp: lastTimestamp,
	}, nil
}

// RecordTokenTracking writes one sanitized, aligned four-family observation.
// Invalid observations fail closed and storage errors remain diagnostic only.
func (store *TokenTrackingStore) RecordTokenTracking(item TokenSourceObservation) {
	if store == nil {
		return
	}
	item, ok := sanitizeTokenSourceObservation(item)
	if !ok {
		return
	}

	store.mu.Lock()
	defer store.mu.Unlock()
	if store.db == nil {
		return
	}
	if item.Timestamp <= store.lastTimestamp {
		item.Timestamp = store.lastTimestamp + 1
	}

	base := labels.FromStrings(
		labelAgentID, item.AgentID,
		labelTokenSourceKind, item.SourceKind,
		labelTokenSourceName, item.SourceName,
	)
	appender := store.db.Appender(context.Background())
	for _, metric := range []struct {
		name  string
		value uint64
	}{
		{name: metricTokenTrackingCalls, value: item.CallCount},
		{name: metricTokenTrackingInput, value: item.InputToken},
		{name: metricTokenTrackingCached, value: item.CachedInputToken},
		{name: metricTokenTrackingOutput, value: item.OutputToken},
	} {
		if _, err := appender.Append(0, tokenTrackingLabels(base, metric.name), item.Timestamp, float64(metric.value)); err != nil {
			_ = appender.Rollback()
			slog.Warn("failed to append token tracking sample", "metric", metric.name, "error", err)
			return
		}
	}
	if err := appender.Commit(); err != nil {
		slog.Warn("failed to commit token tracking observation", "error", err)
		return
	}
	store.lastTimestamp = item.Timestamp
}

func tokenTrackingLabels(base labels.Labels, metric string) labels.Labels {
	return labels.FromStrings(
		"__name__", metric,
		labelAgentID, base.Get(labelAgentID),
		labelTokenSourceKind, base.Get(labelTokenSourceKind),
		labelTokenSourceName, base.Get(labelTokenSourceName),
	)
}

// Close closes the independent token-source database.
func (store *TokenTrackingStore) Close() error {
	if store == nil {
		return nil
	}
	store.mu.Lock()
	defer store.mu.Unlock()
	if store.db == nil {
		return nil
	}
	err := store.db.Close()
	store.db = nil
	return err
}