package metrics

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/tsdb"
)

// Metric names for TSDB storage
const (
	metricRequestsTotal          = "apr_requests_total"
	metricInputTokensTotal       = "apr_input_tokens_total"
	metricCachedInputTokensTotal = "apr_cached_input_tokens_total"
	metricOutputTokensTotal      = "apr_output_tokens_total"
	metricRequestDurationMs      = "apr_request_duration_ms_total"
	metricErrorsTotal            = "apr_errors_total"
	metricServerMessagesTotal    = "apr_server_messages_total"
	metricAgentMessagesTotal     = "apr_agent_messages_total"
	metricAgentTokenRequests     = "apr_agent_token_requests_total"
	metricAgentInputTokens       = "apr_agent_input_tokens_total"
	metricAgentCachedInputTokens = "apr_agent_cached_input_tokens_total"
	metricAgentOutputTokens      = "apr_agent_output_tokens_total"
	metricPromptCompileRecords   = "apr_prompt_compile_records_total"
	metricPromptCompileTokens    = "apr_prompt_compile_tokens_total"
	metricTokenUsageObservation  = "apr_token_analysis_usage_observation"
	metricTokenUsageInput        = "apr_token_analysis_input_tokens"
	metricTokenUsageCached       = "apr_token_analysis_cached_tokens"
	metricTokenUsageOutput       = "apr_token_analysis_output_tokens"
	metricInvocationExposure     = "apr_token_analysis_invocation_exposure"
	metricInvocationBlockTokens  = "apr_token_analysis_invocation_block_tokens"
)

// Label names
const (
	labelKeyID              = "key_id"
	labelModel              = "model"
	labelErrorType          = "error_type"
	labelServerID           = "server_id"
	labelMessageDirection   = "direction"
	labelAgentID            = "agent_id"
	labelRuntime            = "runtime"
	labelLaunchID           = "launch_id"
	labelLaunchStartedAt    = "launch_started_at"
	labelSessionKind        = "session_kind"
	labelSessionFreshness   = "session_freshness"
	labelLeaderSessionKey   = "leader_session_key"
	labelCompileEvent       = "compile_event"
	labelStage              = "stage"
	labelSlotKey            = "surface" // frozen persisted Prometheus label key; renaming cuts historical apr_prompt_compile_* series continuity
	labelBlockID            = "block_id"
	labelBlockLayer         = "block_layer"
	labelBlockSlot          = "block_slot"
	labelBlockType          = "block_type"
	labelSessionID          = "session_id"
	labelEventID            = "event_id"
	labelTaskPayloadPresent = "task_payload_present"
	labelTaskPayloadFP      = "task_payload_fingerprint"
	labelDurability         = "durability"
	labelInvocationID       = "invocation_id"
	labelInvocationCause    = "invocation_cause"
	labelExposureExcluded   = "exposure_excluded"
	labelObservationID      = "observation_id"
	labelExposureComplete   = "exposure_complete"
)

// TSDBCollector manages statistics collection using Prometheus TSDB.
type TSDBCollector struct {
	mu          sync.RWMutex
	db          *tsdb.DB
	promptDB    storage.Queryable
	stopCh      chan struct{}
	keyMaskMap  map[string]string // Maps keyID to masked API key
	maskMapPath string            // Path to persist key mask mapping
}

// loadKeyMaskMap loads the key mask mapping from disk
func (tc *TSDBCollector) loadKeyMaskMap() error {
	tc.mu.Lock()
	defer tc.mu.Unlock()

	// Check if file exists
	if _, err := os.Stat(tc.maskMapPath); os.IsNotExist(err) {
		// File doesn't exist, initialize empty map
		tc.keyMaskMap = make(map[string]string)
		return nil
	}

	// Read file content
	data, err := os.ReadFile(tc.maskMapPath)
	if err != nil {
		return fmt.Errorf("failed to read key mask map: %w", err)
	}

	// Unmarshal JSON
	var maskMap map[string]string
	if err := json.Unmarshal(data, &maskMap); err != nil {
		return fmt.Errorf("failed to parse key mask map: %w", err)
	}

	tc.keyMaskMap = maskMap
	slog.Info("Loaded key mask mapping", "count", len(tc.keyMaskMap))
	return nil
}

// saveKeyMaskMap saves the key mask mapping to disk
func (tc *TSDBCollector) saveKeyMaskMap() error {
	tc.mu.RLock()
	defer tc.mu.RUnlock()

	// Marshal to JSON
	data, err := json.MarshalIndent(tc.keyMaskMap, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal key mask map: %w", err)
	}

	// Write to file
	if err := os.WriteFile(tc.maskMapPath, data, 0600); err != nil {
		return fmt.Errorf("failed to write key mask map: %w", err)
	}

	return nil
}

// NewTSDBCollector creates a new TSDB-based statistics collector.
func NewTSDBCollector() (*TSDBCollector, error) {
	opts := tsdb.DefaultOptions()
	opts.RetentionDuration = int64(7 * 24 * time.Hour / time.Millisecond) // 7 days
	opts.MinBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.MaxBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.NoLockfile = false

	db, err := tsdb.Open(tsdbPath, nil, nil, opts, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to open TSDB: %w", err)
	}

	// Create mask map path in the same directory as TSDB
	maskMapPath := filepath.Join(filepath.Dir(tsdbPath), "key_mask_map.json")

	collector := &TSDBCollector{
		db:          db,
		promptDB:    db,
		stopCh:      make(chan struct{}),
		keyMaskMap:  make(map[string]string),
		maskMapPath: maskMapPath,
	}

	// Load existing mask map
	if err := collector.loadKeyMaskMap(); err != nil {
		slog.Warn("Failed to load existing key mask map, starting with empty map", "error", err)
	}

	return collector, nil
}

func (tc *TSDBCollector) promptQueryable() storage.Queryable {
	if tc == nil {
		return nil
	}
	if tc.promptDB != nil {
		return tc.promptDB
	}
	return tc.db
}

// Start begins any background maintenance (no-op for TSDB).
func (tc *TSDBCollector) Start(ctx context.Context) {
	// TSDB handles its own maintenance
}

// Stop gracefully stops the collector.
func (tc *TSDBCollector) Stop() {
	close(tc.stopCh)
}

// Close closes the TSDB database.
func (tc *TSDBCollector) Close() error {
	if tc.db != nil {
		return tc.db.Close()
	}
	return nil
}