from __future__ import annotations

import asyncio
import json
import os
import time
import uuid
from typing import Any, AsyncIterator

import httpx
import pytest


_RUN_DEV_FT = os.getenv("INVESTIBOOK_RUN_DEV_FT") == "1"
_BASE_URL = os.getenv("INVESTIBOOK_DEV_BASE_URL", "").rstrip("/")
_TOKEN = os.getenv("INVESTIBOOK_DEV_TOKEN", "").strip()
_PROMPT = os.getenv(
    "INVESTIBOOK_DEV_FT_CHAT_PROMPT",
    "This is a functional test. Start answering, but do not finish quickly.",
).strip()

pytestmark = pytest.mark.skipif(
    not (_RUN_DEV_FT and _BASE_URL and _TOKEN),
    reason=(
        "Dev FT disabled. Set INVESTIBOOK_RUN_DEV_FT=1, "
        "INVESTIBOOK_DEV_BASE_URL, and INVESTIBOOK_DEV_TOKEN to run."
    ),
)


def _auth_headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {_TOKEN}"}


async def _iter_sse_events(response: httpx.Response) -> AsyncIterator[dict[str, Any]]:
    data_lines: list[str] = []
    async for line in response.aiter_lines():
        if not line:
            if data_lines:
                yield json.loads("\n".join(data_lines))
                data_lines = []
            continue
        if line.startswith(":"):
            continue
        if line.startswith("data: "):
            data_lines.append(line[6:])


async def _wait_for_assistant_message(
    client: httpx.AsyncClient,
    workspace_id: str,
    assistant_id: str,
    *,
    expected_status: str,
    timeout_s: float = 20.0,
) -> dict[str, Any]:
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        messages_res = await client.get(
            f"{_BASE_URL}/api/workspace/{workspace_id}/messages",
            headers=_auth_headers(),
        )
        messages_res.raise_for_status()
        messages = messages_res.json()["messages"]
        for message in messages:
            if message.get("id") != assistant_id:
                continue
            if message.get("status") == expected_status:
                return message
        await asyncio.sleep(0.5)
    raise AssertionError(
        f"Timed out waiting for assistant message {assistant_id} to reach status={expected_status}"
    )


@pytest.mark.asyncio
async def test_dev_chat_cancel_persists_cancelled_message() -> None:
    timeout = httpx.Timeout(60.0, connect=10.0)
    async with httpx.AsyncClient(timeout=timeout) as client:
        workspace_id: str | None = None
        try:
            create_res = await client.post(
                f"{_BASE_URL}/api/workspaces",
                headers=_auth_headers(),
                json={
                    "name": f"FT cancel persistence {uuid.uuid4().hex[:8]}",
                    "description": "Temporary workspace for dev FT",
                    "is_example": False,
                },
            )
            create_res.raise_for_status()
            workspace_id = create_res.json()["id"]

            status_before_res = await client.get(
                f"{_BASE_URL}/api/workspace/{workspace_id}/messages/status",
                headers=_auth_headers(),
            )
            status_before_res.raise_for_status()
            status_before = status_before_res.json()

            assistant_id: str | None = None
            saw_realtime_event = False

            async with client.stream(
                "GET",
                f"{_BASE_URL}/api/chat/stream",
                params={
                    "message": _PROMPT,
                    "workspace_id": workspace_id,
                    "token": _TOKEN,
                },
            ) as stream_res:
                stream_res.raise_for_status()
                async for event in _iter_sse_events(stream_res):
                    event_type = event.get("type")
                    if event_type == "message_id":
                        assistant_id = str(event["id"])
                        cancel_res = await client.post(
                            f"{_BASE_URL}/api/chat/cancel",
                            headers=_auth_headers(),
                            params={"workspace_id": workspace_id},
                        )
                        cancel_res.raise_for_status()
                        assert cancel_res.json()["status"] in {"cancelled", "no_session"}
                    elif event_type in {"text", "thought", "tool"}:
                        saw_realtime_event = True
                    elif event_type == "complete" and assistant_id:
                        break

            assert assistant_id, "Expected backend to emit assistant message_id before cancel"

            cancelled_message = await _wait_for_assistant_message(
                client,
                workspace_id,
                assistant_id,
                expected_status="cancelled",
            )

            assert cancelled_message["role"] == "assistant"
            assert cancelled_message["status"] == "cancelled"

            tool_parts = [
                part
                for part in cancelled_message.get("parts", [])
                if part.get("type") == "tool"
            ]
            assert all(
                (part.get("state") or {}).get("status") != "running"
                for part in tool_parts
            ), "Cancelled assistant message should not keep running tools"

            status_after_res = await client.get(
                f"{_BASE_URL}/api/workspace/{workspace_id}/messages/status",
                headers=_auth_headers(),
            )
            status_after_res.raise_for_status()
            status_after = status_after_res.json()

            assert status_after["count"] >= status_before["count"] + 2
            assert status_after["updated_at"] != status_before["updated_at"]

            if saw_realtime_event:
                assert cancelled_message.get("content", "") != ""
        finally:
            if workspace_id:
                delete_res = await client.delete(
                    f"{_BASE_URL}/api/workspaces/{workspace_id}",
                    headers=_auth_headers(),
                )
                delete_res.raise_for_status()