from __future__ import annotations

from pathlib import Path
import sys
import time

from bootstrap.app_server import AppServerStdioClient


def _write_stub_server(path: Path) -> None:
    path.write_text(
        r"""
import json
import sys


def send(msg: dict) -> None:
    sys.stdout.write(json.dumps(msg, separators=(",", ":")) + "\n")
    sys.stdout.flush()


initialized = False
THREAD_ID = "thr-123"
TURN_ID = "turn-1"


for line in sys.stdin:
    raw = line.strip()
    if not raw:
        continue
    msg = json.loads(raw)
    method = msg.get("method")
    req_id = msg.get("id")
    params = msg.get("params") or {}

    if method == "initialize":
        send({"id": req_id, "result": {"userAgent": "stub/1"}})
        continue

    if method == "initialized":
        initialized = True
        continue

    if method == "thread/start":
        if not initialized:
            send({"id": req_id, "error": {"code": -32002, "message": "Not initialized"}})
            continue
        send({"id": req_id, "result": {"thread": {"id": THREAD_ID}}})
        send(
            {
                "method": "sessionConfigured",
                "params": {
                    "sessionId": THREAD_ID,
                    "model": "gpt-test",
                    "reasoningEffort": None,
                    "historyLogId": 0,
                    "historyEntryCount": 0,
                    "initialMessages": None,
                    "rolloutPath": "rollout.jsonl",
                },
            }
        )
        send({"method": "thread/started", "params": {"thread": {"id": THREAD_ID}}})
        continue

    if method == "turn/start":
        if params.get("threadId") != THREAD_ID:
            send({"id": req_id, "error": {"code": -32602, "message": "bad threadId"}})
            continue

        # Ask the client to approve command execution.
        approval_id = 99
        send(
            {
                "id": approval_id,
                "method": "item/commandExecution/requestApproval",
                "params": {
                    "threadId": THREAD_ID,
                    "turnId": TURN_ID,
                    "itemId": "item-1",
                    "reason": None,
                    "command": "echo hi",
                    "cwd": params.get("cwd") or "",
                },
            }
        )
        resp = json.loads(sys.stdin.readline())
        if resp.get("id") != approval_id:
            raise SystemExit("expected approval response")
        decision = (resp.get("result") or {}).get("decision")
        if decision not in ("accept", "acceptForSession", "decline", "cancel"):
            raise SystemExit(f"unexpected approval decision: {decision!r}")

        # Respond to turn/start.
        send(
            {
                "id": req_id,
                "result": {
                    "turn": {"id": TURN_ID, "items": [], "status": "inProgress", "error": None}
                },
            }
        )

        # Stream a few notifications.
        send(
            {
                "method": "turn/started",
                "params": {
                    "threadId": THREAD_ID,
                    "turn": {"id": TURN_ID, "items": [], "status": "inProgress", "error": None},
                },
            }
        )
        send(
            {
                "method": "item/agentMessage/delta",
                "params": {"threadId": THREAD_ID, "turnId": TURN_ID, "itemId": "msg-1", "delta": "hello"},
            }
        )
        send(
            {
                "method": "thread/tokenUsage/updated",
                "params": {
                    "threadId": THREAD_ID,
                    "turnId": TURN_ID,
                    "tokenUsage": {
                        "total": {
                            "totalTokens": 10,
                            "inputTokens": 7,
                            "cachedInputTokens": 2,
                            "outputTokens": 1,
                            "reasoningOutputTokens": 0,
                        },
                        "last": {
                            "totalTokens": 10,
                            "inputTokens": 7,
                            "cachedInputTokens": 2,
                            "outputTokens": 1,
                            "reasoningOutputTokens": 0,
                        },
                        "modelContextWindow": None,
                    },
                },
            }
        )
        send(
            {
                "method": "turn/completed",
                "params": {
                    "threadId": THREAD_ID,
                    "turn": {"id": TURN_ID, "items": [], "status": "completed", "error": None},
                },
            }
        )
        continue

    send({"id": req_id, "error": {"code": -32601, "message": "not found"}})
""".lstrip(),
        encoding="utf-8",
    )


def test_app_server_stdio_client_roundtrip(tmp_path: Path) -> None:
    stub = tmp_path / "stub_server.py"
    _write_stub_server(stub)

    notifications: list[str] = []

    def on_notification(msg: dict[str, object]) -> None:
        method = msg.get("method")
        if isinstance(method, str):
            notifications.append(method)

    def request_handler(msg: dict[str, object]) -> dict[str, object]:
        if msg.get("method") == "item/commandExecution/requestApproval":
            return {"decision": "acceptForSession"}
        return {}

    client = AppServerStdioClient(
        argv=[sys.executable, "-u", str(stub)],
        env=None,
        cwd=tmp_path,
        on_notification=on_notification,
        request_handler=request_handler,
    )

    init = client.initialize(
        client_info={"name": "flywheel-bootstrap-test", "version": "0"}
    )
    assert init.user_agent == "stub/1"

    thread = client.request(
        method="thread/start",
        params={"cwd": str(tmp_path), "experimentalRawEvents": False},
        timeout_seconds=2,
    )
    assert thread.get("thread", {}).get("id") == "thr-123"

    _turn = client.request(
        method="turn/start",
        params={"threadId": "thr-123", "input": [{"type": "text", "text": "hi"}]},
        timeout_seconds=2,
    )

    deadline = time.time() + 2
    while time.time() < deadline and "turn/completed" not in notifications:
        time.sleep(0.01)

    assert "sessionConfigured" in notifications
    assert "thread/started" in notifications
    assert "turn/started" in notifications
    assert "item/agentMessage/delta" in notifications
    assert "thread/tokenUsage/updated" in notifications
    assert "turn/completed" in notifications

    client.close()
