"""FastAI remote-runner wire-protocol + кодек SDK <-> JSON.

Используется с обеих сторон: сервер FastAI и runner у друга.
Все фреймы — JSON-объекты, отправляются как WS text frames.

Версия протокола: V = 1. При несовместимых изменениях бампать V.

ФРЕЙМЫ
======

Контрольные (используются на handshake/heartbeat):

    {"v":1, "type":"hello", "runner_id":"abc", "server_time":1717689296}
        — server → runner после успешного auth.

    {"v":1, "type":"ping", "ts":1717689296.123}
    {"v":1, "type":"pong", "ts":1717689296.123}
        — heartbeat; ts эхо-копируется в pong.

Query-фреймы (один query = одна сессия запрос+ответ):

    {"v":1, "type":"query_start", "qid":"<uuid>", "prompt":"...", "options":{...}}
        — server → runner. options сериализуется через encode_options().

    {"v":1, "type":"query_message", "qid":"<uuid>", "msg":{...encode_message()...}}
        — runner → server. Один WS-фрейм = одно SDK-сообщение из async for.

    {"v":1, "type":"query_end", "qid":"<uuid>"}
        — runner → server. query() закончился штатно (ResultMessage уже отправлен
        в последнем query_message; этот фрейм — сигнал конца async-итерации).

    {"v":1, "type":"query_error", "qid":"<uuid>", "error":"...", "kind":"<class>"}
        — runner → server. Любая ошибка во время query(). На стороне сервера
        её надо поднять в async-генераторе как исключение.

    {"v":1, "type":"query_cancel", "qid":"<uuid>"}
        — server → runner. Запрос отмены (юзер нажал stop).

ОГРАНИЧЕНИЯ (5a.1)
==================

`encode_options` НЕ передаёт hooks/can_use_tool (это callable) и mcp_servers
со сложным state. В runner-mode hooks игнорируются — это known-limitation,
будет решено отдельным roundtrip-протоколом позже.

`env` передаётся как есть — но runner ВОЛЕН подменить ANTHROPIC_* ключи
своими, чтобы трафик к Anthropic шёл с его IP/ключа.

`cwd` передаётся как hint; runner может использовать свой workdir.

`resume` передаётся, но имеет смысл только если runner шарит CLI-state
с предыдущим query на этом же runner_id — что не гарантировано.
"""
from __future__ import annotations

import dataclasses
import json
import time
import uuid
from typing import Any, Iterable

PROTOCOL_VERSION = 1


# ---------- сериализация сообщений (SDK -> JSON-safe dict) -----------------

def encode_block(block: Any) -> dict:
    """Сериализует один ContentBlock в dict с тегом '_t'."""
    cls = type(block).__name__
    if cls == "TextBlock":
        return {"_t": "TextBlock", "text": block.text}
    if cls == "ThinkingBlock":
        return {"_t": "ThinkingBlock", "thinking": block.thinking, "signature": block.signature}
    if cls == "ToolUseBlock":
        return {"_t": "ToolUseBlock", "id": block.id, "name": block.name, "input": block.input}
    if cls == "ToolResultBlock":
        return {
            "_t": "ToolResultBlock",
            "tool_use_id": block.tool_use_id,
            "content": block.content,
            "is_error": block.is_error,
        }
    raise ValueError(f"unknown content block type: {cls}")


def encode_content(content: Any) -> Any:
    """Сериализует content (str или list[ContentBlock])."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return [encode_block(b) for b in content]
    if content is None:
        return None
    raise ValueError(f"unsupported content type: {type(content).__name__}")


def encode_message(msg: Any) -> dict:
    """Сериализует SDK-сообщение (AssistantMessage / ResultMessage / UserMessage /
    SystemMessage) в JSON-safe dict с тегом '_t'."""
    cls = type(msg).__name__
    if cls == "AssistantMessage":
        return {
            "_t": "AssistantMessage",
            "content": [encode_block(b) for b in msg.content],
            "model": msg.model,
            "parent_tool_use_id": msg.parent_tool_use_id,
            "error": msg.error,
        }
    if cls == "ResultMessage":
        return {
            "_t": "ResultMessage",
            "subtype": msg.subtype,
            "duration_ms": msg.duration_ms,
            "duration_api_ms": msg.duration_api_ms,
            "is_error": msg.is_error,
            "num_turns": msg.num_turns,
            "session_id": msg.session_id,
            "stop_reason": msg.stop_reason,
            "total_cost_usd": msg.total_cost_usd,
            "usage": msg.usage,
            "result": msg.result,
            "structured_output": msg.structured_output,
        }
    if cls == "UserMessage":
        return {
            "_t": "UserMessage",
            "content": encode_content(msg.content),
            "uuid": getattr(msg, "uuid", None),
            "parent_tool_use_id": getattr(msg, "parent_tool_use_id", None),
            "tool_use_result": getattr(msg, "tool_use_result", None),
        }
    if cls == "SystemMessage":
        return {"_t": "SystemMessage", "subtype": msg.subtype, "data": msg.data}
    # StreamEvent или что-то другое — пытаемся через asdict как fallback
    if dataclasses.is_dataclass(msg):
        d = dataclasses.asdict(msg)
        d["_t"] = cls
        return d
    raise ValueError(f"unknown message type: {cls}")


# ---------- десериализация (JSON-safe dict -> SDK-объект) ------------------

def decode_block(data: dict) -> Any:
    """Восстанавливает ContentBlock из dict. Импорты ленивые — protocol.py
    можно запускать без установленного claude_agent_sdk (например для тестов
    кодирования словарей)."""
    from claude_agent_sdk.types import TextBlock, ThinkingBlock, ToolUseBlock, ToolResultBlock
    t = data.get("_t")
    if t == "TextBlock":
        return TextBlock(text=data["text"])
    if t == "ThinkingBlock":
        return ThinkingBlock(thinking=data["thinking"], signature=data["signature"])
    if t == "ToolUseBlock":
        return ToolUseBlock(id=data["id"], name=data["name"], input=data["input"])
    if t == "ToolResultBlock":
        return ToolResultBlock(
            tool_use_id=data["tool_use_id"],
            content=data.get("content"),
            is_error=data.get("is_error"),
        )
    raise ValueError(f"unknown block tag: {t}")


def decode_content(content: Any) -> Any:
    if isinstance(content, str) or content is None:
        return content
    if isinstance(content, list):
        return [decode_block(b) for b in content]
    raise ValueError(f"unsupported content payload: {type(content).__name__}")


def decode_message(data: dict) -> Any:
    """Восстанавливает SDK-сообщение из dict."""
    from claude_agent_sdk import AssistantMessage, ResultMessage
    from claude_agent_sdk.types import UserMessage, SystemMessage
    t = data.get("_t")
    if t == "AssistantMessage":
        return AssistantMessage(
            content=[decode_block(b) for b in data["content"]],
            model=data["model"],
            parent_tool_use_id=data.get("parent_tool_use_id"),
            error=data.get("error"),
        )
    if t == "ResultMessage":
        return ResultMessage(
            subtype=data["subtype"],
            duration_ms=data["duration_ms"],
            duration_api_ms=data["duration_api_ms"],
            is_error=data["is_error"],
            num_turns=data["num_turns"],
            session_id=data["session_id"],
            stop_reason=data.get("stop_reason"),
            total_cost_usd=data.get("total_cost_usd"),
            usage=data.get("usage"),
            result=data.get("result"),
            structured_output=data.get("structured_output"),
        )
    if t == "UserMessage":
        return UserMessage(
            content=decode_content(data["content"]),
            uuid=data.get("uuid"),
            parent_tool_use_id=data.get("parent_tool_use_id"),
            tool_use_result=data.get("tool_use_result"),
        )
    if t == "SystemMessage":
        return SystemMessage(subtype=data["subtype"], data=data["data"])
    # Generic fallback: любой новый dataclass-message из SDK (TaskStartedMessage,
    # TaskProgressMessage, TaskNotificationMessage, SessionMessage и т.п.).
    # Симметрично encode_message → dataclasses.asdict + _t.
    if isinstance(t, str):
        try:
            from claude_agent_sdk import types as _sdk_types  # type: ignore
            cls = getattr(_sdk_types, t, None)
            if cls is not None and dataclasses.is_dataclass(cls):
                field_names = {f.name for f in dataclasses.fields(cls)}
                kwargs = {k: v for k, v in data.items() if k in field_names}
                return cls(**kwargs)
        except Exception:
            pass
    raise ValueError(f"unknown message tag: {t}")


# ---------- сериализация options ------------------------------------------
# Передаём только JSON-safe поля. hooks / can_use_tool — callable, опускаем.

_OPTIONS_SCALAR_FIELDS = (
    "allowed_tools",
    "disallowed_tools",
    "permission_mode",
    "cwd",
    "model",
    "env",
    "system_prompt",
    "resume",
    "max_turns",
    "max_thinking_tokens",
    "add_dirs",
    "extra_args",
    "include_partial_messages",
    "fork_session",
    "agents",
    "setting_sources",
    "user",
)


def encode_options(options: Any) -> dict:
    """Сериализует ClaudeAgentOptions в JSON-safe dict (только скалярные поля).

    Пропускает:
      - hooks (callable)
      - can_use_tool (callable)
      - mcp_servers (могут содержать сложные объекты)
    Это known-limitation шага 5a.1 — расширим позже roundtrip-протоколом.
    """
    out: dict[str, Any] = {}
    for name in _OPTIONS_SCALAR_FIELDS:
        if hasattr(options, name):
            val = getattr(options, name)
            if val is None:
                continue
            try:
                json.dumps(val)  # быстрый тест сериализуемости
            except (TypeError, ValueError):
                continue  # пропускаем непереводимое
            out[name] = val
    return out


def decode_options(data: dict) -> Any:
    """Конструирует ClaudeAgentOptions из dict. Поля, которых нет в data,
    остаются дефолтными. Runner может ДО вызова query() патчить env/cwd
    своими значениями."""
    from claude_agent_sdk import ClaudeAgentOptions
    safe = {k: v for k, v in data.items() if k in _OPTIONS_SCALAR_FIELDS}
    return ClaudeAgentOptions(**safe)


# ---------- frame builders -------------------------------------------------

def new_qid() -> str:
    return uuid.uuid4().hex


def frame_hello(runner_id: str) -> dict:
    return {"v": PROTOCOL_VERSION, "type": "hello", "runner_id": runner_id, "server_time": int(time.time())}


def frame_ping(ts: float | None = None) -> dict:
    return {"v": PROTOCOL_VERSION, "type": "ping", "ts": ts if ts is not None else time.time()}


def frame_pong(ts: float) -> dict:
    return {"v": PROTOCOL_VERSION, "type": "pong", "ts": ts}


def frame_query_start(qid: str, prompt: str, options: Any, session_id: str = "") -> dict:
    f = {
        "v": PROTOCOL_VERSION,
        "type": "query_start",
        "qid": qid,
        "prompt": prompt,
        "options": encode_options(options),
    }
    if session_id:
        f["session_id"] = session_id
    return f


def frame_query_message(qid: str, msg: Any) -> dict:
    return {
        "v": PROTOCOL_VERSION,
        "type": "query_message",
        "qid": qid,
        "msg": encode_message(msg),
    }


def frame_query_end(qid: str) -> dict:
    return {"v": PROTOCOL_VERSION, "type": "query_end", "qid": qid}


def frame_query_error(qid: str, error: str, kind: str = "Exception") -> dict:
    return {"v": PROTOCOL_VERSION, "type": "query_error", "qid": qid, "error": error, "kind": kind}


def frame_query_cancel(qid: str) -> dict:
    return {"v": PROTOCOL_VERSION, "type": "query_cancel", "qid": qid}


def dump(frame: dict) -> str:
    """Сериализация WS-фрейма в текст. ensure_ascii=False — экономия места на
    русских/китайских строках; separators — компактнее."""
    return json.dumps(frame, ensure_ascii=False, separators=(",", ":"))


def parse(text: str) -> dict:
    """Парсит WS-фрейм. Валидирует версию."""
    data = json.loads(text)
    if not isinstance(data, dict):
        raise ValueError("frame must be a JSON object")
    v = data.get("v")
    if v != PROTOCOL_VERSION:
        raise ValueError(f"protocol version mismatch: got {v!r}, expected {PROTOCOL_VERSION}")
    if not isinstance(data.get("type"), str):
        raise ValueError("frame missing 'type'")
    return data