from __future__ import annotations

import asyncio
import json
import os
import uuid
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any, Callable

_JSONRPC_VERSION = "2.0"
_NANOBOT_API_STREAM_LIMIT = 2 * 1024 * 1024
NANOBOT_API_SOCKET = Path(
    os.getenv("NANOBOT_API_SOCKET", str(Path.home() / ".nanobot" / "api.sock"))
).expanduser()


class _NanobotApiError(RuntimeError):
    def __init__(self, code: int, message: str) -> None:
        super().__init__(message)
        self.code = code


def _jsonrpc_request(request_id: str, method: str, params: dict[str, Any]) -> dict[str, Any]:
    return {
        "jsonrpc": _JSONRPC_VERSION,
        "id": request_id,
        "method": method,
        "params": params,
    }


def _jsonrpc_notification(method: str, params: dict[str, Any]) -> dict[str, Any]:
    return {
        "jsonrpc": _JSONRPC_VERSION,
        "method": method,
        "params": params,
    }


async def _open_nanobot_api_socket(
    socket_path: Path = NANOBOT_API_SOCKET,
    *,
    stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
    if not socket_path.exists():
        raise RuntimeError(
            f"Nanobot API socket not found at {socket_path}. "
            "Enable channels.api and start `nanobot gateway`."
        )
    try:
        return await asyncio.open_unix_connection(
            path=str(socket_path),
            limit=stream_limit,
        )
    except OSError as exc:
        raise RuntimeError(f"failed to connect to Nanobot API socket: {exc}") from exc


async def _send_nanobot_api_request(
    method: str,
    params: dict[str, Any],
    *,
    timeout_seconds: float,
    socket_path: Path = NANOBOT_API_SOCKET,
    stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
) -> Any:
    request_id = str(uuid.uuid4())
    reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit)
    try:
        writer.write(
            (
                json.dumps(_jsonrpc_request(request_id, method, params), ensure_ascii=False) + "\n"
            ).encode("utf-8")
        )
        await writer.drain()

        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout_seconds

        while True:
            remaining = deadline - loop.time()
            if remaining <= 0:
                raise RuntimeError(f"timed out waiting for Nanobot API response to {method}")

            try:
                line = await asyncio.wait_for(reader.readline(), timeout=remaining)
            except ValueError as exc:
                raise RuntimeError(
                    "Nanobot API response exceeded the configured stream limit"
                ) from exc
            if not line:
                raise RuntimeError("Nanobot API socket closed before responding")

            try:
                message = json.loads(line.decode("utf-8", errors="replace"))
            except json.JSONDecodeError:
                continue
            if not isinstance(message, dict):
                continue
            if message.get("jsonrpc") != _JSONRPC_VERSION:
                continue
            if "method" in message:
                continue
            if str(message.get("id", "")).strip() != request_id:
                continue
            if "error" in message:
                error = message.get("error", {})
                if isinstance(error, dict):
                    raise _NanobotApiError(
                        int(error.get("code", -32000)),
                        str(error.get("message", "unknown Nanobot API error")),
                    )
                raise _NanobotApiError(-32000, str(error))
            return message.get("result")
    finally:
        writer.close()
        await writer.wait_closed()


async def _run_nanobot_agent_turn(
    content: str,
    *,
    chat_id: str,
    metadata: dict[str, Any] | None = None,
    timeout_seconds: float = 90.0,
    notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None],
    socket_path: Path = NANOBOT_API_SOCKET,
    stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
) -> dict[str, Any]:
    final_message = ""
    saw_activity = False
    async for typed_event in _stream_nanobot_agent_turn(
        content,
        chat_id=chat_id,
        metadata=metadata,
        timeout_seconds=timeout_seconds,
        notification_to_event=notification_to_event,
        socket_path=socket_path,
        stream_limit=stream_limit,
    ):
        event_type = typed_event.get("type")
        if event_type == "message":
            saw_activity = True
            if (
                not bool(typed_event.get("is_progress", False))
                and typed_event.get("role") == "nanobot"
            ):
                final_message = str(typed_event.get("content", ""))
            continue

        if event_type == "agent_state":
            saw_activity = True

    return {
        "status": "ok",
        "message": final_message,
        "completed": saw_activity,
    }


async def _stream_nanobot_agent_turn(
    content: str,
    *,
    chat_id: str,
    metadata: dict[str, Any] | None = None,
    timeout_seconds: float = 60.0,
    idle_grace_seconds: float = 1.0,
    notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None],
    socket_path: Path = NANOBOT_API_SOCKET,
    stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
) -> AsyncIterator[dict[str, Any]]:
    reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit)
    saw_final_message = False
    saw_activity = False
    try:
        writer.write(
            (
                json.dumps(
                    _jsonrpc_notification(
                        "message.send",
                        {
                            "content": content,
                            "chat_id": chat_id,
                            "metadata": metadata or {},
                        },
                    ),
                    ensure_ascii=False,
                )
                + "\n"
            ).encode("utf-8")
        )
        await writer.drain()

        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout_seconds

        while True:
            remaining = deadline - loop.time()
            if remaining <= 0:
                raise RuntimeError("Timed out waiting for a Nanobot response.")

            timeout = min(remaining, idle_grace_seconds) if saw_final_message else remaining
            try:
                line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            except asyncio.TimeoutError:
                if saw_final_message:
                    break
                raise RuntimeError("Timed out waiting for Nanobot to start responding")
            except ValueError as exc:
                raise RuntimeError(
                    "Nanobot API response exceeded the configured stream limit"
                ) from exc

            if not line:
                break

            try:
                obj = json.loads(line.decode("utf-8", errors="replace").strip())
            except json.JSONDecodeError:
                continue
            if not isinstance(obj, dict):
                continue

            typed_event = notification_to_event(obj)
            if typed_event is None:
                continue

            event_type = typed_event.get("type")
            if event_type == "message":
                saw_activity = True
                if (
                    not bool(typed_event.get("is_progress", False))
                    and typed_event.get("role") == "nanobot"
                ):
                    saw_final_message = True

            yield typed_event

            if (
                saw_final_message
                and event_type == "agent_state"
                and str(typed_event.get("state", "")).strip().lower() == "idle"
            ):
                break

        if not saw_activity:
            raise RuntimeError("Timed out waiting for Nanobot to start responding")
    finally:
        writer.close()
        await writer.wait_closed()


NanobotApiError = _NanobotApiError
send_nanobot_api_request = _send_nanobot_api_request
run_nanobot_agent_turn = _run_nanobot_agent_turn
stream_nanobot_agent_turn = _stream_nanobot_agent_turn

__all__ = [
    "NANOBOT_API_SOCKET",
    "NanobotApiError",
    "run_nanobot_agent_turn",
    "send_nanobot_api_request",
    "stream_nanobot_agent_turn",
]
