package observability

import (
	"fmt"
	"sync"
	"time"
)

// Metrics holds all Prometheus metrics for the gateway.
type Metrics struct {
	mu sync.RWMutex

	// Request counters: request_total{route, upstream, status}
	requestTotal map[string]int64

	// Request latencies: request_duration_seconds (histogram)
	// Stored as cumulative buckets for Prometheus text format
	requestDuration map[string]int64 // stores duration samples in milliseconds
	requestDurationBuckets map[string]map[float64]int64 // histogram buckets

	// Bytes counters: gateway_bytes{direction, route, upstream}
	bytesIn  map[string]int64
	bytesOut map[string]int64

	// Upstream health: upstream_health{upstream} = 1 or 0
	upstreamHealth map[string]int

	// Streaming metrics
	streamingResponsesTotal map[string]int64
	streamingByteCount      map[string]int64

	// LLM inference metrics (TTFT and ITL)
	// ttftMs: Time-to-First-Token in milliseconds
	ttftMs map[string][]int64 // samples for histogram
	// itlMs: Inter-Token Latency in milliseconds
	itlMs map[string][]int64 // samples for histogram
	// Token counts
	tokenCount map[string]int64
}

// NewMetrics creates a new Metrics instance.
func NewMetrics() *Metrics {
	return &Metrics{
		requestTotal:           make(map[string]int64),
		requestDuration:        make(map[string]int64),
		requestDurationBuckets: make(map[string]map[float64]int64),
		bytesIn:                make(map[string]int64),
		bytesOut:               make(map[string]int64),
		upstreamHealth:         make(map[string]int),
		streamingResponsesTotal: make(map[string]int64),
		streamingByteCount:     make(map[string]int64),
		ttftMs:                 make(map[string][]int64),
		itlMs:                  make(map[string][]int64),
		tokenCount:             make(map[string]int64),
	}
}

// RecordRequest records a request with its route, upstream, status, and duration.
func (m *Metrics) RecordRequest(route, upstream string, statusCode int, duration time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("%s:%s:%d", route, upstream, statusCode)
	m.requestTotal[key]++

	// Record duration in milliseconds
	durationKey := fmt.Sprintf("%s:%s", route, upstream)
	m.requestDuration[durationKey] += int64(duration.Milliseconds())

	// Record in histogram buckets
	if _, ok := m.requestDurationBuckets[durationKey]; !ok {
		m.requestDurationBuckets[durationKey] = make(map[float64]int64)
	}

	// Prometheus histogram buckets: .005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10
	buckets := []float64{0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10}
	durationSeconds := duration.Seconds()

	for _, bucket := range buckets {
		if durationSeconds <= bucket {
			m.requestDurationBuckets[durationKey][bucket]++
		}
	}
}

// RecordBytesIn records incoming bytes.
func (m *Metrics) RecordBytesIn(route, upstream string, bytes int64) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("%s:%s", route, upstream)
	m.bytesIn[key] += bytes
}

// RecordBytesOut records outgoing bytes.
func (m *Metrics) RecordBytesOut(route, upstream string, bytes int64) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("%s:%s", route, upstream)
	m.bytesOut[key] += bytes
}

// SetUpstreamHealth sets the health status of an upstream (1 = healthy, 0 = unhealthy).
func (m *Metrics) SetUpstreamHealth(upstream string, healthy bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if healthy {
		m.upstreamHealth[upstream] = 1
	} else {
		m.upstreamHealth[upstream] = 0
	}
}

// RecordStreamingResponse records a streaming response with its total byte count and duration.
func (m *Metrics) RecordStreamingResponse(route, upstream string, totalBytes int64, duration time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("%s:%s", route, upstream)
	m.streamingResponsesTotal[key]++
	m.streamingByteCount[key] += totalBytes

	// Also record as request duration
	m.recordDuration(key, duration)
}

func (m *Metrics) recordDuration(key string, duration time.Duration) {
	m.requestDuration[key] += int64(duration.Milliseconds())
	
	if _, ok := m.requestDurationBuckets[key]; !ok {
		m.requestDurationBuckets[key] = make(map[float64]int64)
	}

	buckets := []float64{0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10}
	durationSeconds := duration.Seconds()

	for _, bucket := range buckets {
		if durationSeconds <= bucket {
			m.requestDurationBuckets[key][bucket]++
		}
	}
}

// GetMetrics returns a copy of current metrics (for testing/export).
func (m *Metrics) GetMetrics() map[string]interface{} {
	m.mu.RLock()
	defer m.mu.RUnlock()

	return map[string]interface{}{
		"request_total":             m.requestTotal,
		"request_duration":          m.requestDuration,
		"request_duration_buckets":  m.requestDurationBuckets,
		"bytes_in":                  m.bytesIn,
		"bytes_out":                 m.bytesOut,
		"upstream_health":           m.upstreamHealth,
		"streaming_responses_total": m.streamingResponsesTotal,
		"streaming_byte_count":      m.streamingByteCount,
		"llm_ttft_ms":               m.ttftMs,
		"llm_itl_ms":                m.itlMs,
		"llm_token_count":           m.tokenCount,
	}
}

// RecordTTFT records Time-to-First-Token in milliseconds
func (m *Metrics) RecordTTFT(model string, ttftMs int64) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("llm:ttft:%s", model)
	m.ttftMs[key] = append(m.ttftMs[key], ttftMs)
}

// RecordITL records Inter-Token Latency in milliseconds
func (m *Metrics) RecordITL(model string, itlMs int64) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("llm:itl:%s", model)
	m.itlMs[key] = append(m.itlMs[key], itlMs)
}

// RecordTokenCount records number of tokens in response
func (m *Metrics) RecordTokenCount(model string, count int64) {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := fmt.Sprintf("llm:tokens:%s", model)
	m.tokenCount[key] += count
}

// GetTTFTMetrics returns TTFT statistics for Prometheus export
func (m *Metrics) GetTTFTMetrics() map[string]interface{} {
	m.mu.RLock()
	defer m.mu.RUnlock()

	result := make(map[string]interface{})
	for key, samples := range m.ttftMs {
		if len(samples) > 0 {
			result[key] = map[string]interface{}{
				"count": len(samples),
				"sum":   sumInt64(samples),
				"avg":   sumInt64(samples) / int64(len(samples)),
				"min":   minInt64(samples),
				"max":   maxInt64(samples),
			}
		}
	}
	return result
}

// GetITLMetrics returns ITL statistics for Prometheus export
func (m *Metrics) GetITLMetrics() map[string]interface{} {
	m.mu.RLock()
	defer m.mu.RUnlock()

	result := make(map[string]interface{})
	for key, samples := range m.itlMs {
		if len(samples) > 0 {
			result[key] = map[string]interface{}{
				"count": len(samples),
				"sum":   sumInt64(samples),
				"avg":   sumInt64(samples) / int64(len(samples)),
				"min":   minInt64(samples),
				"max":   maxInt64(samples),
			}
		}
	}
	return result
}

func sumInt64(vals []int64) int64 {
	var s int64
	for _, v := range vals {
		s += v
	}
	return s
}

func minInt64(vals []int64) int64 {
	if len(vals) == 0 {
		return 0
	}
	min := vals[0]
	for _, v := range vals {
		if v < min {
			min = v
		}
	}
	return min
}

func maxInt64(vals []int64) int64 {
	if len(vals) == 0 {
		return 0
	}
	max := vals[0]
	for _, v := range vals {
		if v > max {
			max = v
		}
	}
	return max
}

// Reset clears all metrics (for testing).
func (m *Metrics) Reset() {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.requestTotal = make(map[string]int64)
	m.requestDuration = make(map[string]int64)
	m.requestDurationBuckets = make(map[string]map[float64]int64)
	m.bytesIn = make(map[string]int64)
	m.bytesOut = make(map[string]int64)
	m.upstreamHealth = make(map[string]int)
	m.streamingResponsesTotal = make(map[string]int64)
	m.streamingByteCount = make(map[string]int64)
	m.ttftMs = make(map[string][]int64)
	m.itlMs = make(map[string][]int64)
	m.tokenCount = make(map[string]int64)
}