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
}

// 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),
	}
}

// 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,
	}
}

// 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)
}