"""Tests for engine-integrated observability (issues #66 #67 #68 #78 #80 #82).

Verifies:
- trace_id full-chain propagation (request → engine → MetricsCollector → Response)
- Structured JSON logging via TraceContextAdapter
- Prometheus metrics export from engine
- Engine-level metrics API (get_engine_metrics, enable_prometheus, reset)

No model loading is performed — tests use LLMEngine constructor only.
"""

from __future__ import annotations

import logging

import pytest

from sagellm_core.llm_engine import LLMEngine, LLMEngineConfig
from sagellm_core.observability import (
    EngineMetrics,
    MetricsCollector,
    PrometheusExporter,
    TraceContextAdapter,
    setup_logger,
)
from sagellm_core.observability.logger import StructuredFormatter


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
def engine_config() -> LLMEngineConfig:
    """Minimal engine config (no real model required — engine not started)."""
    return LLMEngineConfig(
        model_path="test-model",
        backend_type="cpu",
        engine_id="test-engine-obs-01",
    )


@pytest.fixture()
def engine(engine_config: LLMEngineConfig) -> LLMEngine:
    """LLMEngine instance (not started — observability API is available before start)."""
    return LLMEngine(engine_config)


# ---------------------------------------------------------------------------
# 1. Engine observability initialization
# ---------------------------------------------------------------------------


class TestEngineObservabilityInit:
    """Engine observability attributes are initialised in __init__."""

    def test_engine_has_metrics_collector(self, engine: LLMEngine) -> None:
        """Engine must have a MetricsCollector after construction."""
        assert hasattr(engine, "_engine_collector")
        assert isinstance(engine._engine_collector, MetricsCollector)

    def test_engine_collector_has_engine_id(self, engine: LLMEngine) -> None:
        """MetricsCollector must carry the engine_id for trace correlation."""
        assert engine._engine_collector.engine_id == "test-engine-obs-01"

    def test_engine_has_structured_logger(self, engine: LLMEngine) -> None:
        """Engine must have a structured logger instance."""
        assert hasattr(engine, "_structured_logger")
        assert isinstance(engine._structured_logger, logging.Logger)

    def test_structured_logger_has_json_formatter(self, engine: LLMEngine) -> None:
        """At least one handler must use StructuredFormatter (JSON output)."""
        assert any(
            isinstance(h.formatter, StructuredFormatter) for h in engine._structured_logger.handlers
        ), "Engine structured logger must have a StructuredFormatter handler"

    def test_prometheus_exporter_starts_none(self, engine: LLMEngine) -> None:
        """Prometheus exporter must be None until explicitly enabled."""
        assert engine._prometheus_exporter is None


# ---------------------------------------------------------------------------
# 2. get_engine_metrics()
# ---------------------------------------------------------------------------


class TestGetEngineMetrics:
    """get_engine_metrics() returns an EngineMetrics snapshot."""

    def test_returns_engine_metrics_type(self, engine: LLMEngine) -> None:
        """get_engine_metrics() must return EngineMetrics."""
        metrics = engine.get_engine_metrics()
        assert isinstance(metrics, EngineMetrics)

    def test_initial_metrics_are_zero(self, engine: LLMEngine) -> None:
        """Before any requests, all numeric metrics must be zero."""
        metrics = engine.get_engine_metrics()
        assert metrics.ttft_ms == 0.0
        assert metrics.throughput_tps == 0.0
        assert metrics.num_requests == 0
        assert metrics.error_rate == 0.0

    def test_metrics_after_simulated_request(self, engine: LLMEngine) -> None:
        """Metrics reflect data recorded into the engine-level collector."""
        coll = engine._engine_collector
        coll.record_ttft(42.0)
        coll.record_token_generation(5.0)
        coll.record_request_latency(200.0)
        coll.record_kv_usage(512)
        coll.record_memory_usage(1024.0)

        metrics = engine.get_engine_metrics()

        assert metrics.ttft_ms == 42.0
        assert metrics.num_requests == 1
        assert metrics.kv_used_tokens == 512
        assert metrics.peak_mem_mb == 1024.0

    def test_engine_id_in_metrics(self, engine: LLMEngine) -> None:
        """Returned EngineMetrics must carry engine_id."""
        # engine_collector.compute_metrics() propagates engine_id from the collector
        metrics = engine.get_engine_metrics()
        assert metrics.engine_id == "test-engine-obs-01"

    def test_reset_engine_metrics(self, engine: LLMEngine) -> None:
        """reset_engine_metrics() clears accumulated data."""
        engine._engine_collector.record_ttft(50.0)
        engine._engine_collector.record_request_latency(100.0)

        engine.reset_engine_metrics()

        metrics = engine.get_engine_metrics()
        assert metrics.ttft_ms == 0.0
        assert metrics.num_requests == 0

    def test_engine_id_preserved_after_reset(self, engine: LLMEngine) -> None:
        """reset_engine_metrics() must preserve engine_id."""
        engine.reset_engine_metrics()
        assert engine._engine_collector.engine_id == "test-engine-obs-01"


# ---------------------------------------------------------------------------
# 3. Prometheus export API (issues #68 #80)
# ---------------------------------------------------------------------------


class TestPrometheusAPI:
    """enable_prometheus() and get_prometheus_metrics() work correctly."""

    def test_enable_prometheus_returns_exporter(self, engine: LLMEngine) -> None:
        """enable_prometheus() must return a PrometheusExporter."""
        try:
            exporter = engine.enable_prometheus(namespace="sagellm_test")
            assert isinstance(exporter, PrometheusExporter)
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_enable_prometheus_is_idempotent(self, engine: LLMEngine) -> None:
        """Calling enable_prometheus() twice must return the same instance."""
        try:
            exp1 = engine.enable_prometheus()
            exp2 = engine.enable_prometheus()
            assert exp1 is exp2
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_prometheus_exporter_stored_on_engine(self, engine: LLMEngine) -> None:
        """After enable_prometheus(), _prometheus_exporter must be set."""
        try:
            engine.enable_prometheus()
            assert engine._prometheus_exporter is not None
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_get_prometheus_metrics_returns_string(self, engine: LLMEngine) -> None:
        """get_prometheus_metrics() must return a non-empty string."""
        try:
            text = engine.get_prometheus_metrics(namespace="sagellm_test")
            assert isinstance(text, str)
            assert len(text) > 0
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_prometheus_metrics_contain_namespace(self, engine: LLMEngine) -> None:
        """Prometheus output must include the specified namespace."""
        try:
            text = engine.get_prometheus_metrics(namespace="sagellm_test")
            assert "sagellm_test_" in text
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_prometheus_metrics_contain_kv_gauge(self, engine: LLMEngine) -> None:
        """Prometheus output must include the kv_used_tokens gauge by default."""
        try:
            engine._engine_collector.record_kv_usage(256)
            text = engine.get_prometheus_metrics(namespace="sagellm_test")
            assert "kv_used_tokens" in text
        except ImportError:
            pytest.skip("prometheus_client not installed")

    def test_prometheus_metrics_contain_throughput(self, engine: LLMEngine) -> None:
        """Prometheus output must include the throughput gauge."""
        try:
            text = engine.get_prometheus_metrics(namespace="sagellm_test")
            assert "throughput_tps" in text
        except ImportError:
            pytest.skip("prometheus_client not installed")


# ---------------------------------------------------------------------------
# 4. Structured logging with trace context (issues #66 #67)
# ---------------------------------------------------------------------------


class TestStructuredLogging:
    """TraceContextAdapter propagates trace_id / request_id / engine_id."""

    def test_trace_context_adapter_injects_trace_id(self, engine: LLMEngine) -> None:
        """TraceContextAdapter must inject trace_id into log extra fields."""
        adapter = TraceContextAdapter(
            engine._structured_logger,
            {
                "trace_id": "chain-test-trace",
                "request_id": "chain-req-001",
                "engine_id": engine.engine_id,
            },
        )
        _msg, kwargs = adapter.process("test message", {})
        extra = kwargs.get("extra", {})

        assert extra["trace_id"] == "chain-test-trace"
        assert extra["request_id"] == "chain-req-001"
        assert extra["engine_id"] == "test-engine-obs-01"

    def test_structured_formatter_outputs_json(self) -> None:
        """StructuredFormatter must produce valid JSON per log record."""
        import json

        formatter = StructuredFormatter()
        record = logging.LogRecord(
            name="test",
            level=logging.INFO,
            pathname="test.py",
            lineno=1,
            msg="test structured message",
            args=(),
            exc_info=None,
        )
        record.trace_id = "fmt-trace-001"
        record.request_id = "fmt-req-001"

        output = formatter.format(record)
        parsed = json.loads(output)

        assert parsed["message"] == "test structured message"
        assert parsed["trace_id"] == "fmt-trace-001"
        assert parsed["request_id"] == "fmt-req-001"
        assert parsed["level"] == "INFO"

    def test_setup_logger_structured_mode(self) -> None:
        """setup_logger(structured=True) must install StructuredFormatter."""
        logger = setup_logger("test.obs.engine.structured", structured=True)
        assert any(isinstance(h.formatter, StructuredFormatter) for h in logger.handlers)

    def test_engine_collector_mark_timestamps(self, engine: LLMEngine) -> None:
        """MetricsCollector.mark() must record timestamps that survive compute_metrics()."""
        engine._engine_collector.mark("req-001:queued")
        engine._engine_collector.mark("req-001:completed")

        metrics = engine.get_engine_metrics()
        assert "req-001:queued" in metrics.timestamps
        assert "req-001:completed" in metrics.timestamps


# ---------------------------------------------------------------------------
# 5. trace_id full-chain (request → MetricsCollector → EngineMetrics)
# ---------------------------------------------------------------------------


class TestTraceIdFullChain:
    """trace_id flows from request IDs through MetricsCollector to EngineMetrics."""

    def test_collector_trace_id_propagated_to_metrics(self) -> None:
        """MetricsCollector initialized with trace_id propagates it to EngineMetrics."""
        coll = MetricsCollector(
            trace_id="chain-trace-999",
            request_id="chain-req-999",
            engine_id="eng-chain-01",
        )
        coll.record_ttft(10.0)
        metrics = coll.compute_metrics()

        assert metrics.trace_id == "chain-trace-999"
        assert metrics.request_id == "chain-req-999"
        assert metrics.engine_id == "eng-chain-01"

    def test_engine_collector_engine_id_in_metrics(self, engine: LLMEngine) -> None:
        """Engine-level collector carries engine_id through to EngineMetrics."""
        metrics = engine.get_engine_metrics()
        assert metrics.engine_id == engine.engine_id

    def test_structured_log_contains_trace_fields(self) -> None:
        """StructuredFormatter includes trace_id, request_id, engine_id in JSON output."""
        import json

        formatter = StructuredFormatter()
        record = logging.LogRecord(
            name="sagellm.engine.eng-trace-test",
            level=logging.INFO,
            pathname="llm_engine.py",
            lineno=42,
            msg="generate: request received",
            args=(),
            exc_info=None,
        )
        record.trace_id = "trace-full-001"
        record.request_id = "req-full-001"
        record.engine_id = "eng-full-01"

        output = formatter.format(record)
        parsed = json.loads(output)

        assert parsed["trace_id"] == "trace-full-001"
        assert parsed["request_id"] == "req-full-001"
        assert parsed["engine_id"] == "eng-full-01"

    def test_metrics_to_json_contains_trace_id(self) -> None:
        """EngineMetrics.to_json() must include trace_id for downstream consumers."""
        import json

        coll = MetricsCollector(trace_id="json-trace-001", engine_id="json-eng-01")
        coll.record_ttft(25.0)
        metrics = coll.compute_metrics()

        parsed = json.loads(metrics.to_json())
        assert parsed["trace_id"] == "json-trace-001"
        assert parsed["engine_id"] == "json-eng-01"
        assert parsed["ttft_ms"] == 25.0
