import json
import os
import asyncio
from pathlib import Path

import httpx
from fastapi.testclient import TestClient

import app
import config
import db
import routers.chat
import triage
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS


def _mock_triage_url(monkeypatch, url: str = config.LLAMA_SERVER_BASE):
    monkeypatch.setattr(config, "LLAMA_SERVER_BASE", url)


def make_client(tmp_path: Path) -> TestClient:
    os.environ["CAIC_ADMIN_PIN"] = "1234"
    db.DB_PATH = tmp_path / "caic-streaming.db"
    SESSIONS.clear()
    PIN_ATTEMPTS.clear()
    RATE_EVENTS.clear()
    db.init_db()
    return TestClient(app.app, raise_server_exceptions=False)


def parse_sse_payloads(body: str) -> list[dict]:
    payloads: list[dict] = []
    for chunk in body.split("\n\n"):
        chunk = chunk.strip()
        if not chunk.startswith("data: "):
            continue
        raw = chunk[len("data: ") :]
        payloads.append(json.loads(raw))
    return payloads


class _MockStreamResponse:
    def __init__(self, lines: list[str]):
        self._lines = lines

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return False

    async def aiter_lines(self):
        for line in self._lines:
            yield line


def _stream_json_lines(events: list[dict]) -> list[str]:
    return [json.dumps(event) for event in events]


def test_chat_stream_emits_tokens_and_done(tmp_path: Path, monkeypatch):
    _mock_triage_url(monkeypatch)
    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
            "session_id"
        ]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        events = _stream_json_lines(
            [
                {"message": {"content": "Hel"}, "logprobs": [{"logprob": -0.01}]},
                {"message": {"content": "lo"}, "logprobs": [{"logprob": -0.01}]},
                {"done": True, "eval_count": 2, "eval_duration": 1000000000},
            ]
        )

        def stream_stub(self, method, url, json=None, timeout=None):
            return _MockStreamResponse(events)

        monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)

        resp = client.post(
            "/api/chat",
            json={"message": "hello", "model": config.DEFAULT_MODEL},
            headers=headers,
        )
        assert resp.status_code == 200
        payloads = parse_sse_payloads(resp.text)

        token_text = "".join(p.get("token", "") for p in payloads if "token" in p)
        assert token_text == "Hello"
        done_events = [p for p in payloads if p.get("done")]
        assert done_events
        assert "searched" not in done_events[-1]


def test_chat_auto_search_trigger_emits_search_events(tmp_path: Path, monkeypatch):
    _mock_triage_url(monkeypatch)
    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
            "session_id"
        ]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        first_stream = _stream_json_lines(
            [
                {
                    "message": {"content": "I don't have current data on that question."},
                    "logprobs": [{"logprob": -5.0}],
                },
                {"done": True, "eval_count": 2, "eval_duration": 1000000000},
            ]
        )
        second_stream = _stream_json_lines(
            [
                {"message": {"content": "Based on current data: 42."}},
                {"done": True},
            ]
        )
        stream_batches = [first_stream, second_stream]

        def stream_stub(self, method, url, json=None, timeout=None):
            return _MockStreamResponse(stream_batches.pop(0))

        async def search_stub(query: str, max_results: int = 5):
            return [
                {
                    "title": "Answer",
                    "url": "https://example.com",
                    "content": "The value is 42.",
                }
            ]

        monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)
        monkeypatch.setattr(routers.chat, "query_searxng", search_stub)

        resp = client.post(
            "/api/chat",
            json={"message": "what is the latest value", "model": config.DEFAULT_MODEL},
            headers=headers,
        )
        assert resp.status_code == 200
        payloads = parse_sse_payloads(resp.text)

        assert any(p.get("searching") is True for p in payloads)
        assert any("search_results" in p for p in payloads)
        assert any(p.get("augmented") is True for p in payloads)
        done_events = [p for p in payloads if p.get("done")]
        assert done_events and done_events[-1].get("searched") is True


def test_chat_with_upload_context_id_injects_document(tmp_path: Path, monkeypatch):
    _mock_triage_url(monkeypatch)
    captured_payload = {}

    def stream_stub(self, method, url, json=None, timeout=None):
        nonlocal captured_payload
        captured_payload = json
        events = [{"message": {"content": "ok"}, "logprobs": [{"logprob": -0.01}]}, {"done": True, "eval_count": 1, "eval_duration": 1000000000}]
        return _MockStreamResponse([__import__('json').dumps(e) for e in events])

    monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)

    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        db_local = db.get_db()
        expires = "2099-12-31T23:59:59+00:00"
        cid = db.insert_upload_context(db_local, "conv-up", "report.txt", "Confidential document content here", expires, "text/plain")
        db_local.commit()
        db_local.close()

        resp = client.post(
            "/api/chat",
            json={"message": "summarize this", "upload_context_id": cid, "model": config.DEFAULT_MODEL},
            headers=headers,
        )
        assert resp.status_code == 200
        system_content = next((m["content"] for m in captured_payload.get("messages", []) if m["role"] == "system"), "")
        assert "Confidential document content here" in system_content


def test_chat_with_expired_upload_context_id_silent(tmp_path: Path, monkeypatch):
    _mock_triage_url(monkeypatch)
    captured_payload = {}

    def stream_stub(self, method, url, json=None, timeout=None):
        nonlocal captured_payload
        captured_payload = json
        events = [{"message": {"content": "ok"}, "logprobs": [{"logprob": -0.01}]}, {"done": True, "eval_count": 1, "eval_duration": 1000000000}]
        return _MockStreamResponse([__import__('json').dumps(e) for e in events])

    monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)

    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        import datetime
        db_local = db.get_db()
        expires = (datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2)).isoformat()
        cid = db.insert_upload_context(db_local, "conv-exp", "old.txt", "Stale data", expires, "text/plain")
        db_local.commit()
        db_local.close()

        resp = client.post(
            "/api/chat",
            json={"message": "hi", "upload_context_id": cid, "model": config.DEFAULT_MODEL},
            headers=headers,
        )
        assert resp.status_code == 200
        system_content = next((m["content"] for m in captured_payload.get("messages", []) if m["role"] == "system"), "")
        assert "Stale data" not in system_content


def test_memory_command_paths_remember_and_forget(tmp_path: Path, monkeypatch):
    _mock_triage_url(monkeypatch)
    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
            "session_id"
        ]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        base_stream = _stream_json_lines(
            [
                {"message": {"content": "ok"}, "logprobs": [{"logprob": -0.01}]},
                {"done": True, "eval_count": 1, "eval_duration": 1000000000},
            ]
        )

        def stream_stub(self, method, url, json=None, timeout=None):
            return _MockStreamResponse(base_stream)

        monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)

        remember_resp = client.post(
            "/api/chat",
            json={
                "message": "remember that my favorite language is rust",
                "model": config.DEFAULT_MODEL,
            },
            headers=headers,
        )
        assert remember_resp.status_code == 200
        remember_events = parse_sse_payloads(remember_resp.text)
        assert any("Remembered" in p.get("token", "") for p in remember_events)

        memories_after_add = client.get("/api/memories", headers={"X-Session-ID": sid, "Origin": "http://testserver"})
        assert memories_after_add.status_code == 200
        assert memories_after_add.json().get("count", 0) >= 1

        forget_resp = client.post(
            "/api/chat",
            json={
                "message": "forget about my favorite language",
                "model": config.DEFAULT_MODEL,
            },
            headers=headers,
        )
        assert forget_resp.status_code == 200
        forget_events = parse_sse_payloads(forget_resp.text)
        assert any("Forgot" in p.get("token", "") for p in forget_events)

def test_private_chat_does_not_persist(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(httpx.AsyncClient, "stream", lambda *a, **kw: _MockStreamResponse([
        'data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":null,"logprobs":{"content":[]}}]}',
        'data: {"choices":[{"delta":{"content":" world"},"finish_reason":null,"logprobs":{"content":[]}}]}',
        'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[]}}],"usage":{"completion_tokens":2,"prompt_tokens":10,"tokens_per_second":5.0}}',
        "data: [DONE]",
    ]))
    monkeypatch.setattr(triage, "classify_query", lambda q: "general")

    async def _mock_ensure(m): return True
    monkeypatch.setattr("model_pull.ensure_model", _mock_ensure)

    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        resp = client.post("/api/chat", json={
            "message": "test private message",
            "private_chat": True,
        }, headers=headers)
        assert resp.status_code == 200
        events = parse_sse_payloads(resp.text)
        tokens = [p for p in events if "token" in p]
        assert len(tokens) == 2
        done = [p for p in events if p.get("done")]
        assert len(done) == 1

        conv_resp = client.get("/api/conversations", headers=headers)
        assert conv_resp.status_code == 200
        assert len(conv_resp.json()) == 0


def test_private_chat_does_not_auto_search(tmp_path: Path, monkeypatch):
    """Private mode should skip auto-search even when search is enabled."""
    monkeypatch.setattr(httpx.AsyncClient, "stream", lambda *a, **kw: _MockStreamResponse([
        'data: {"choices":[{"delta":{"content":"I don\'t know"},"finish_reason":null,"logprobs":{"content":[{"logprob":-2.5}]}}]}',
        'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[{"logprob":-2.5}]}}],"usage":{"completion_tokens":1,"prompt_tokens":10,"tokens_per_second":5.0}}',
        "data: [DONE]",
    ]))
    monkeypatch.setattr(triage, "classify_query", lambda q: "general")
    monkeypatch.setattr(routers.chat, "query_searxng", lambda q: [{"title": "result"}])

    with make_client(tmp_path) as client:
        sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
        headers = {"X-Session-ID": sid, "Origin": "http://testserver"}

        resp = client.post("/api/chat", json={
            "message": "what is the weather",
            "private_chat": True,
        }, headers=headers)
        assert resp.status_code == 200
        events = parse_sse_payloads(resp.text)
        searching = [p for p in events if p.get("searching")]
        assert len(searching) == 0, "private chat should not auto-search"
