"""Tests for observability utilities."""

from __future__ import annotations

import json

from sagellm_core.observability import (
    EngineMetrics,
    MetricsCollector,
    PrometheusExporter,
    StructuredFormatter,
    TraceContextAdapter,
    get_logger,
    setup_logger,
)


class TestEngineMetrics:
    """Test EngineMetrics."""

    def test_default_metrics(self):
        """Test default metrics."""
        metrics = EngineMetrics()
        assert metrics.ttft_ms == 0.0
        assert metrics.throughput_tps == 0.0
        assert metrics.num_requests == 0
        assert metrics.trace_id == ""
        assert metrics.kv_used_tokens == 0
        assert metrics.prefix_hit_rate == 0.0
        assert metrics.spec_accept_rate == 0.0

    def test_to_dict(self):
        """Test converting to dict."""
        metrics = EngineMetrics(
            trace_id="trace-123",
            request_id="req-456",
            engine_id="eng-0",
            ttft_ms=50.0,
            throughput_tps=100.0,
            kv_used_tokens=1024,
            prefix_hit_rate=0.85,
        )
        d = metrics.to_dict()

        assert d["trace_id"] == "trace-123"
        assert d["request_id"] == "req-456"
        assert d["engine_id"] == "eng-0"
        assert d["ttft_ms"] == 50.0
        assert d["throughput_tps"] == 100.0
        assert d["kv_used_tokens"] == 1024
        assert d["prefix_hit_rate"] == 0.85

    def test_to_json(self):
        """Test converting to JSON."""
        metrics = EngineMetrics(
            trace_id="trace-123",
            ttft_ms=50.0,
        )
        json_str = metrics.to_json()
        parsed = json.loads(json_str)

        assert parsed["trace_id"] == "trace-123"
        assert parsed["ttft_ms"] == 50.0


class TestMetricsCollector:
    """Test MetricsCollector."""

    def test_initialization_with_trace_context(self):
        """Test initialization with trace context."""
        collector = MetricsCollector(
            trace_id="trace-123",
            request_id="req-456",
            engine_id="eng-0",
        )
        assert collector.trace_id == "trace-123"
        assert collector.request_id == "req-456"
        assert collector.engine_id == "eng-0"

    def test_record_ttft(self):
        """Test recording TTFT."""
        collector = MetricsCollector()
        collector.record_ttft(50.0)
        collector.record_ttft(60.0)

        metrics = collector.compute_metrics()
        assert metrics.ttft_ms == 55.0  # Average

    def test_record_latency(self):
        """Test recording request latency."""
        collector = MetricsCollector()
        collector.record_request_latency(100.0)
        collector.record_request_latency(200.0)
        collector.record_request_latency(150.0)

        metrics = collector.compute_metrics()
        assert metrics.num_requests == 3
        assert metrics.latency_p50_ms == 150.0

    def test_record_token_generation(self):
        """Test recording token generation time."""
        collector = MetricsCollector()

        for _ in range(10):
            collector.record_token_generation(10.0)

        metrics = collector.compute_metrics()
        assert metrics.num_tokens_generated == 10
        assert metrics.tpot_ms == 10.0
        assert metrics.throughput_tps > 0

    def test_kv_cache_metrics(self):
        """Test KV cache metrics recording."""
        collector = MetricsCollector()
        collector.record_kv_usage(1024)
        collector.record_prefix_hit()
        collector.record_prefix_hit()
        collector.record_prefix_miss()

        metrics = collector.compute_metrics()
        assert metrics.kv_used_tokens == 1024
        assert metrics.prefix_hit_rate == 2.0 / 3.0  # 2 hits out of 3 total

    def test_speculative_decoding_metrics(self):
        """Test speculative decoding metrics."""
        collector = MetricsCollector()
        collector.record_spec_tokens(accepted=8, total=10)
        collector.record_spec_tokens(accepted=7, total=10)

        metrics = collector.compute_metrics()
        assert metrics.spec_accept_rate == 15.0 / 20.0  # 15 accepted out of 20 total

    def test_memory_metrics(self):
        """Test memory usage recording."""
        collector = MetricsCollector()
        collector.record_memory_usage(512.0)
        collector.record_memory_usage(768.0)
        collector.record_memory_usage(650.0)

        metrics = collector.compute_metrics()
        assert metrics.peak_mem_mb == 768.0  # Peak value

    def test_error_tracking(self):
        """Test error tracking."""
        collector = MetricsCollector()
        collector.record_request_latency(100.0)
        collector.record_error()
        collector.record_request_latency(200.0)

        metrics = collector.compute_metrics()
        assert metrics.error_rate == 1.0 / 3.0  # 1 error out of 3 total

    def test_timestamp_marking(self):
        """Test timestamp marking."""
        collector = MetricsCollector()
        collector.mark("queued")
        collector.mark("scheduled")
        collector.mark("executed")
        collector.mark("completed")

        metrics = collector.compute_metrics()
        assert "queued" in metrics.timestamps
        assert "scheduled" in metrics.timestamps
        assert "executed" in metrics.timestamps
        assert "completed" in metrics.timestamps
        assert metrics.timestamps["completed"] >= metrics.timestamps["queued"]

    def test_to_json(self):
        """Test JSON export."""
        collector = MetricsCollector(
            trace_id="trace-123",
            request_id="req-456",
        )
        collector.record_ttft(50.0)
        collector.mark("queued")

        json_str = collector.to_json()
        parsed = json.loads(json_str)

        assert parsed["trace_id"] == "trace-123"
        assert parsed["request_id"] == "req-456"
        assert "ttft_ms" in parsed

    def test_reset(self):
        """Test resetting metrics."""
        collector = MetricsCollector(trace_id="trace-123")
        collector.record_ttft(50.0)
        collector.record_kv_usage(1024)
        collector.mark("queued")
        collector.reset()

        assert collector.trace_id == "trace-123"  # Context preserved
        metrics = collector.get_metrics()
        assert metrics.ttft_ms == 0.0
        assert metrics.kv_used_tokens == 0
        assert len(metrics.timestamps) == 0


class TestLogger:
    """Test logging utilities."""

    def test_setup_logger(self):
        """Test setting up logger."""
        logger = setup_logger("test_logger")
        assert logger.name == "test_logger"

    def test_get_logger(self):
        """Test getting logger."""
        logger = get_logger("test_logger2")
        assert logger.name == "test_logger2"

    def test_structured_logging(self):
        """Test structured JSON logging."""
        logger = setup_logger("test_structured", structured=True)
        assert logger.name == "test_structured"
        assert any(isinstance(h.formatter, StructuredFormatter) for h in logger.handlers)

    def test_trace_context_adapter(self):
        """Test TraceContextAdapter."""
        logger = get_logger("test_trace")
        trace_logger = TraceContextAdapter(
            logger, {"trace_id": "trace-123", "request_id": "req-456"}
        )

        # Verify context is added
        msg, kwargs = trace_logger.process("test message", {})
        assert "extra" in kwargs
        assert kwargs["extra"]["trace_id"] == "trace-123"
        assert kwargs["extra"]["request_id"] == "req-456"


class TestPrometheusExporter:
    """Test PrometheusExporter."""

    def test_initialization(self):
        """Test exporter initialization."""
        try:
            exporter = PrometheusExporter(namespace="test")
            assert exporter.namespace == "test"
            assert exporter.registry is not None
        except ImportError:
            # Skip if prometheus_client not installed
            pass

    def test_update_from_collector(self):
        """Test updating Prometheus metrics from collector."""
        try:
            exporter = PrometheusExporter(namespace="test")
            collector = MetricsCollector(
                trace_id="trace-123",
                engine_id="eng-0",
            )
            collector.record_ttft(50.0)
            collector.record_kv_usage(1024)

            exporter.update_from_collector(collector)

            # Verify metrics can be exported
            metrics_text = exporter.export_text()
            assert "test_throughput_tps" in metrics_text
            assert "test_kv_used_tokens" in metrics_text
        except ImportError:
            # Skip if prometheus_client not installed
            pass

    def test_export(self):
        """Test exporting metrics."""
        try:
            exporter = PrometheusExporter(namespace="test")
            collector = MetricsCollector()
            collector.record_request_latency(100.0)

            exporter.update_from_collector(collector)

            metrics_bytes = exporter.export()
            assert isinstance(metrics_bytes, bytes)

            metrics_text = exporter.export_text()
            assert isinstance(metrics_text, str)
        except ImportError:
            # Skip if prometheus_client not installed
            pass
