"""Integration: short-lived DB sessions in the chat endpoint (SEC-14-04).

Verifies that the ``POST /api/chat`` endpoint uses short-lived sessions
for retrieval steps (steering notes, KB overview, retrieve) and for the
query_log write — no DB connection is held across the SSE stream.
"""
from __future__ import annotations

import json
import math
import re
import uuid
from collections.abc import AsyncIterator, Iterator
from typing import Any

import pytest
from fastapi.testclient import TestClient
from sqlalchemy import text
from sqlalchemy.orm import Session

from app.api import chat as chat_api
from app.db import SessionLocal
from app.main import app as fastapi_app
from app.models import Document
from app.rag.llm import StreamPiece, ToolCallPiece
from tests.conftest import ADMIN_PASSWORD

_DIM = 768
_TOKEN_RE = re.compile(r"[a-z0-9]+")


def _token_vec(text: str) -> list[float]:
    """Bag-of-words unit vector — same algorithm as the E2E mock."""
    import hashlib

    vec = [0.0] * _DIM
    for tok in _TOKEN_RE.findall(text.lower()):
        vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % _DIM] += 1.0
    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
    return [v / norm for v in vec]


class FakeChatLLM:
    """Minimal duck-typed LLM client for the chat endpoint."""

    def __init__(self, answer: str = "Test answer.") -> None:
        self.answer = answer

    async def embed_one(self, message: str) -> list[float]:
        return _token_vec(message)

    async def chat_stream(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        scaffolding: Any = None,
    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
        yield StreamPiece("thinking", "thinking")
        yield StreamPiece("content", self.answer)


@pytest.fixture()
def kb(db: Session) -> Iterator[None]:
    """Fresh documents + chunks tables for these tests."""
    db.execute(text("TRUNCATE chunks, documents, folder_summaries, query_log"))
    db.commit()
    yield
    db.execute(text("TRUNCATE chunks, documents, folder_summaries, query_log"))
    db.commit()


def _make_doc(
    source: str, path: str, title: str, content: str, db: Session | None = None
) -> Document:
    """Add a document row to the DB."""
    from app.models import Document

    doc = Document(
        id=uuid.uuid4(),
        source=source,
        path=path,
        full_path=f"/tmp/{path}",
        title=title,
        content=content,
        content_hash="0" * 64,
    )
    if db is not None:
        db.add(doc)
    return doc


class CountingSession:
    """A session wrapper that counts how many times it is created and
    closed, so tests can verify short-lived session usage."""

    _instances: list[CountingSession] = []
    _lock: Any = None

    def __init__(self, real: Session) -> None:
        self._real = real
        self._closed = False

    def __enter__(self) -> CountingSession:
        return self

    def __exit__(self, *args: Any) -> None:
        if not self._closed:
            self._closed = True
            self._real.close()

    def add(self, obj: Any) -> None:
        self._real.add(obj)

    def commit(self) -> None:
        self._real.commit()

    def scalars(self, stmt: Any) -> Any:
        return self._real.scalars(stmt)

    def get(self, model: Any, pk: Any) -> Any:
        return self._real.get(model, pk)

    def execute(self, stmt: Any, params: Any = None) -> Any:
        return self._real.execute(stmt, params)

    @property
    def closed(self) -> bool:
        return self._closed


# ---------- deflected turn uses short-lived sessions ----------


def test_deflected_turn_uses_short_lived_sessions(
    client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
) -> None:
    """A deflected turn (LOW mode) uses short-lived sessions for
    retrieval steps (steering notes, KB overview, retrieve) but does
    not call the session factory used by the agent loop (which does
    not run for deflected turns)."""
    # Seed a document so we have a KB
    _make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
    db.commit()

    # Mock the LLM
    fake_llm = FakeChatLLM(answer="I don't have info on that.")
    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm

    # Log in as admin
    login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
    assert login_resp.status_code == 204

    # Track SessionLocal calls
    original_session_local = SessionLocal
    session_calls: list[bool] = []

    def counting_factory() -> CountingSession:
        real = original_session_local()
        session_calls.append(True)
        return CountingSession(real)

    monkeypatch.setattr("app.api.chat.SessionLocal", counting_factory)

    # Ask a question that will be deflected (cosine below threshold)
    response = client.post(
        "/api/chat",
        json={"message": "completely unrelated question xyz123"},
    )

    assert response.status_code == 200
    frames = list(_parse_sse(response))

    # The turn should end with a done event
    done_frames = [f for f in frames if f["type"] == "done"]
    assert len(done_frames) == 1
    assert done_frames[0]["deflected"] is True

    # Short-lived sessions were used for retrieval
    assert len(session_calls) > 0


# ---------- grounded turn with tool calls uses short-lived sessions ----------


def test_grounded_turn_with_tools_uses_short_lived_sessions(
    client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
) -> None:
    """A grounded turn with tool calls uses a short-lived session for
    each tool round — the session is created, used, and closed per
    tool call."""
    # Seed a document
    _make_doc("Test", "doc.md", "Test Doc", "This is a test document about Kubernetes.", db)
    db.commit()

    # Mock the LLM
    fake_llm = FakeChatLLM(answer="The document is about Kubernetes.")
    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm

    # Log in as admin
    login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
    assert login_resp.status_code == 204

    # Track SessionLocal calls
    original_session_local = SessionLocal
    session_ids: list[int] = []

    def tracking_factory() -> Session:
        real = original_session_local()
        session_ids.append(id(real))
        return real

    monkeypatch.setattr("app.api.chat.SessionLocal", tracking_factory)

    # Ask a question that will be grounded
    response = client.post(
        "/api/chat",
        json={"message": "What is in the test document?"},
    )

    assert response.status_code == 200
    frames = list(_parse_sse(response))

    # The turn should end with a done event
    done_frames = [f for f in frames if f["type"] == "done"]
    assert len(done_frames) == 1

    # Multiple sessions were used (retrieval + tool calls + query_log)
    # Each tool call creates a new session
    assert len(session_ids) >= 1


# ---------- query_log write uses short-lived session ----------


def test_query_log_write_uses_short_lived_session(
    client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
) -> None:
    """The query_log write uses a short-lived session — if the write
    fails, the answer still goes out (the error is caught and logged)."""
    # Seed a document
    _make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
    db.commit()

    # Mock the LLM
    fake_llm = FakeChatLLM(answer="The document contains test content.")
    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm

    # Log in as admin
    login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
    assert login_resp.status_code == 204

    # Track SessionLocal calls
    original_session_local = SessionLocal
    query_log_sessions: list[int] = []

    def tracking_factory() -> Session:
        real = original_session_local()
        # The first few sessions are for retrieval; the last is for query_log
        query_log_sessions.append(id(real))
        return real

    monkeypatch.setattr("app.api.chat.SessionLocal", tracking_factory)

    response = client.post(
        "/api/chat",
        json={"message": "What is in the test document?"},
    )

    assert response.status_code == 200
    frames = list(_parse_sse(response))

    done_frames = [f for f in frames if f["type"] == "done"]
    assert len(done_frames) == 1

    # query_log was written (the short-lived session committed)
    query_log_rows = db.execute(text("SELECT count(*) FROM query_log")).scalar()
    assert query_log_rows == 1


# ---------- DB failure mid-stream works with short-lived sessions ----------


def test_db_failure_mid_stream_with_short_lived_sessions(
    client: TestClient, db, monkeypatch: pytest.MonkeyPatch, kb: None
) -> None:
    """If the DB fails mid-stream (during a tool call), the error
    path works correctly with short-lived sessions."""
    # Seed a document
    _make_doc("Test", "doc.md", "Test Doc", "This is a test document.", db)
    db.commit()

    # Mock the LLM
    fake_llm = FakeChatLLM(answer="The document contains test content.")
    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: fake_llm

    # Log in as admin
    login_resp = client.post("/api/login", json={"password": ADMIN_PASSWORD})
    assert login_resp.status_code == 204

    call_count = [0]

    def failing_factory() -> Session:
        call_count[0] += 1
        if call_count[0] > 2:
            # Fail on the third call (a tool round)
            raise RuntimeError("DB connection lost mid-stream")
        return SessionLocal()

    monkeypatch.setattr("app.api.chat.SessionLocal", failing_factory)

    response = client.post(
        "/api/chat",
        json={"message": "What is in the test document?"},
    )

    assert response.status_code == 200
    frames = list(_parse_sse(response))

    # Should get an error event (not a done event)
    error_frames = [f for f in frames if f["type"] == "error"]
    assert len(error_frames) == 1
    assert "offline" in error_frames[0]["detail"].lower()

    # No done event — the error is terminal
    done_frames = [f for f in frames if f["type"] == "done"]
    assert len(done_frames) == 0


# ---------- helpers ----------


def _parse_sse(response: Any) -> Iterator[dict[str, Any]]:
    """Parse SSE frames from a streaming response."""
    buf = ""
    for chunk in response.iter_text():
        buf += chunk
        while "\n\n" in buf:
            frame, buf = buf.split("\n\n", 1)
            frame = frame.strip()
            if frame.startswith("data:"):
                yield json.loads(frame.removeprefix("data:").strip())