"""Integration tests for the proxy pipeline.

Tests the full request/response flow using Starlette TestClient with a mock
adapter, exercising governance, session management, audit logging, and streaming.
"""

import json
from pathlib import Path

import pytest

from aurora_lens.proxy.config import ProxyConfig
from aurora_lens.adapters.base import LLMAdapter, AdapterResponse
from aurora_lens.interpret.base import ExtractionBackend
from aurora_lens.interpret.schema import ExtractionResult


# ── Shared mock helpers ───────────────────────────────────────────────────────

class _MockAdapter(LLMAdapter):
    def __init__(self, response: str = "Hello from mock."):
        self._response = response

    async def generate(self, messages: list[dict[str, str]], **kwargs) -> AdapterResponse:
        return AdapterResponse(text=self._response, model="mock")


class _StreamAdapter(LLMAdapter):
    def __init__(self, chunks: list[str] | None = None):
        self._chunks = chunks or ["Hello", " world"]

    async def generate(self, messages: list[dict[str, str]], **kwargs) -> AdapterResponse:
        return AdapterResponse(text="".join(self._chunks), model="mock")

    async def generate_stream(self, messages: list[dict[str, str]], **kwargs):
        for chunk in self._chunks:
            yield ({"choices": [{"delta": {"content": chunk}, "index": 0}]}, chunk)


class _MockBackend(ExtractionBackend):
    async def extract(self, text, pef):
        return ExtractionResult()


def _make_client(monkeypatch, adapter=None, audit_log=None, extra: dict | None = None):
    """Build a TestClient-ready proxy app backed by mock adapter."""
    try:
        from starlette.testclient import TestClient
    except ImportError:
        pytest.skip("Starlette not installed")

    from aurora_lens.proxy.app import create_app

    _adapter = adapter or _MockAdapter()
    monkeypatch.setattr(
        "aurora_lens.proxy.app._build_provider_adapters",
        lambda cfg: (_adapter, _adapter),
    )
    cfg_map: dict = {
        "upstream": {"provider": "openai", "api_key": "sk-test", "model": "gpt-4"},
        "governance": {"audit_log": audit_log, "audit_backend": "jsonl"},
        "extraction": {"backend": "spacy"},
    }
    if extra:
        for k, v in extra.items():
            if isinstance(v, dict) and isinstance(cfg_map.get(k), dict):
                cfg_map[k].update(v)
            else:
                cfg_map[k] = v
    cfg = ProxyConfig.from_mapping(cfg_map)
    return TestClient(create_app(cfg))


_BASE_REQUEST = {
    "model": "gpt-4",
    "messages": [{"role": "user", "content": "Hello"}],
}

_HARD_STOP_REQUEST = {
    "model": "gpt-4",
    "messages": [{"role": "user", "content": "Hello"}],
    "aurora": {
        "external_flags": [
            {"type": "SELF_HARM_INSTRUCTION", "evidence": ["test classifier"], "severity": "error"},
        ],
    },
}


# ── PASS flow ─────────────────────────────────────────────────────────────────

class TestIntegrationPass:
    def test_pass_response_http200(self, monkeypatch):
        """PASS flow: HTTP 200, choices populated, aurora.governance == PASS."""
        client = _make_client(monkeypatch)
        r = client.post("/v1/chat/completions", json=_BASE_REQUEST)
        assert r.status_code == 200
        body = r.json()
        assert body["aurora"]["governance"] == "PASS"
        assert body["choices"][0]["message"]["content"] == "Hello from mock."
        assert body["choices"][0]["message"]["role"] == "assistant"

    def test_pass_aurora_headers_present(self, monkeypatch):
        """PASS flow: all required Aurora-* response headers present."""
        client = _make_client(monkeypatch)
        r = client.post("/v1/chat/completions", json=_BASE_REQUEST)
        for header in [
            "Aurora-Outcome", "Aurora-Trace-Id", "Aurora-Audit-Sink",
            "Aurora-Timestamp", "Aurora-Upstream", "Aurora-Policy",
            "Aurora-Policy-Version", "Aurora-Proxy-Ms",
        ]:
            assert header in r.headers, f"Missing response header: {header}"

    def test_healthz_returns_ok(self, monkeypatch):
        """GET /healthz returns {ok: true}."""
        client = _make_client(monkeypatch)
        r = client.get("/healthz")
        assert r.status_code == 200
        assert r.json().get("ok") is True


# ── HARD_STOP flow ────────────────────────────────────────────────────────────

class TestIntegrationHardStop:
    def test_hard_stop_returns_http200(self, monkeypatch):
        """HARD_STOP via external_flags returns HTTP 200 (governance in body)."""
        client = _make_client(monkeypatch)
        r = client.post("/v1/chat/completions", json=_HARD_STOP_REQUEST)
        assert r.status_code == 200
        assert r.json()["aurora"]["governance"] == "HARD_STOP"

    def test_hard_stop_content_is_not_model_output(self, monkeypatch):
        """HARD_STOP response content is a governance message, not the adapter text."""
        client = _make_client(monkeypatch)
        r = client.post("/v1/chat/completions", json=_HARD_STOP_REQUEST)
        content = r.json()["choices"][0]["message"]["content"]
        assert content != "Hello from mock."
        assert len(content) > 0

    def test_hard_stop_writes_audit_entry(self, monkeypatch, tmp_path):
        """HARD_STOP writes at least one audit entry to the configured log path."""
        audit_file = str(tmp_path / "audit.jsonl")
        client = _make_client(monkeypatch, audit_log=audit_file)
        r = client.post("/v1/chat/completions", json=_HARD_STOP_REQUEST)
        assert r.status_code == 200
        lines = [ln for ln in Path(audit_file).read_text().splitlines() if ln.strip()]
        assert len(lines) >= 1
        entry = json.loads(lines[-1])
        assert entry.get("outcome", entry.get("action")) in ("HARD_STOP", "CONTAIN", "FORCE_REVISE")


# ── Session continuity ────────────────────────────────────────────────────────

class TestIntegrationSession:
    def test_session_turn_increments(self, monkeypatch):
        """Turn counter increments on successive requests sharing a session ID."""
        client = _make_client(monkeypatch)
        sid = "test-session-turn-001"
        r1 = client.post(
            "/v1/chat/completions",
            json=_BASE_REQUEST,
            headers={"x-aurora-session-id": sid},
        )
        r2 = client.post(
            "/v1/chat/completions",
            json=_BASE_REQUEST,
            headers={"x-aurora-session-id": sid},
        )
        assert r1.status_code == 200
        assert r2.status_code == 200
        t1 = r1.json()["aurora"]["turn"]
        t2 = r2.json()["aurora"]["turn"]
        assert t2 == t1 + 1

    def test_different_sessions_are_isolated(self, monkeypatch):
        """Two different session IDs have independent turn counters."""
        client = _make_client(monkeypatch)
        # Advance session-alpha twice
        client.post(
            "/v1/chat/completions",
            json=_BASE_REQUEST,
            headers={"x-aurora-session-id": "session-alpha"},
        )
        r_alpha = client.post(
            "/v1/chat/completions",
            json=_BASE_REQUEST,
            headers={"x-aurora-session-id": "session-alpha"},
        )
        # session-beta is brand new
        r_beta = client.post(
            "/v1/chat/completions",
            json=_BASE_REQUEST,
            headers={"x-aurora-session-id": "session-beta"},
        )
        assert r_alpha.json()["aurora"]["turn"] == 2
        assert r_beta.json()["aurora"]["turn"] == 1


# ── External flags ────────────────────────────────────────────────────────────

class TestIntegrationExternalFlags:
    def test_external_flags_cause_governance_intervention(self, monkeypatch):
        """External flags in aurora.external_flags trigger a non-PASS governance action."""
        client = _make_client(monkeypatch)
        r = client.post("/v1/chat/completions", json=_HARD_STOP_REQUEST)
        assert r.status_code == 200
        assert r.json()["aurora"]["governance"] != "PASS"

    def test_invalid_external_flag_type_returns_422(self, monkeypatch):
        """Unrecognised external flag type returns 422."""
        client = _make_client(monkeypatch)
        r = client.post(
            "/v1/chat/completions",
            json={
                "model": "gpt-4",
                "messages": [{"role": "user", "content": "Hi"}],
                "aurora": {
                    "external_flags": [
                        {"type": "NOT_A_REAL_FLAG_TYPE", "evidence": ["x"]},
                    ],
                },
            },
        )
        assert r.status_code == 422


# ── Streaming ─────────────────────────────────────────────────────────────────

class TestIntegrationStreaming:
    def test_stream_returns_event_stream_with_done(self, monkeypatch):
        """POST with stream=true returns text/event-stream containing [DONE]."""
        try:
            from starlette.testclient import TestClient
        except ImportError:
            pytest.skip("Starlette not installed")

        from aurora_lens.proxy.app import create_app

        adapter = _StreamAdapter(["Hello", " world"])
        monkeypatch.setattr(
            "aurora_lens.proxy.app._build_provider_adapters",
            lambda cfg: (adapter, adapter),
        )
        cfg = ProxyConfig.from_mapping({
            "upstream": {"provider": "openai", "api_key": "sk-test", "model": "gpt-4"},
            "governance": {"audit_log": None},
            "extraction": {"backend": "spacy"},
        })
        client = TestClient(create_app(cfg))
        r = client.post(
            "/v1/chat/completions",
            json={"model": "gpt-4", "messages": [{"role": "user", "content": "Hi"}], "stream": True},
        )
        assert r.status_code == 200
        assert "text/event-stream" in r.headers.get("content-type", "")
        assert "[DONE]" in r.text
        assert '"aurora"' in r.text
        assert '"governance"' in r.text
