"""The wire protocol between a chat-completions client and a ChatGPT conversation.

The model has no tool-calling API here, so tool use is carried in the text: the
preamble instructs it to answer with exactly one fenced json block, and every
reply is parsed back into a tool call or a final answer.
"""

from __future__ import annotations

import json
from dataclasses import dataclass

FINAL_TOOL = "final"


@dataclass(frozen=True)
class ToolCall:
    name: str
    arguments: dict


@dataclass(frozen=True)
class Final:
    text: str


@dataclass(frozen=True)
class Malformed:
    reason: str


ReplyKind = ToolCall | Final | Malformed


CONTRACT = """\
You are the reasoning engine of a command-line coding agent. You do not talk to a
human. Every one of your replies is parsed by a program.

Answer with EXACTLY ONE fenced json block and nothing else — no prose before it,
no prose after it, never two blocks.

To use a tool:

```json
{"tool": "<tool name>", "arguments": {}}
```

To finish, when no tool is needed:

```json
{"tool": "final", "text": "<answer>"}
```

After each tool call I reply with its result under a `TOOL RESULT [id]` heading.
Never invent a tool result. Never call a tool that is not listed below.
"""


def render_preamble(system: str, tools: list[dict]) -> str:
    """The first message of a conversation: contract, tools, then the client's system prompt."""
    parts = [CONTRACT, "# TOOLS"]
    for tool in tools:
        spec = _tool_spec(tool)
        parts.append(
            f"## {spec['name']}\n{spec.get('description', '')}\n"
            f"arguments schema:\n{json.dumps(spec.get('schema') or {}, indent=1)}"
        )
    if not tools:
        parts.append("(no tools available — answer with the final block)")
    parts.append("# SYSTEM\n" + system)
    return "\n\n".join(parts)


def _tool_spec(tool: dict) -> dict:
    """Accept both the OpenAI (`function`) and Anthropic (`input_schema`) tool shapes."""
    if "function" in tool:
        fn = tool["function"]
        return {"name": fn.get("name", ""), "description": fn.get("description", ""),
                "schema": fn.get("parameters")}
    return {"name": tool.get("name", ""), "description": tool.get("description", ""),
            "schema": tool.get("input_schema") or tool.get("parameters")}


def tool_names(tools: list[dict]) -> set[str]:
    return {_tool_spec(t)["name"] for t in tools}


def parse_reply(blocks: list[str], tools: list[dict], text: str = "") -> ReplyKind:
    """Turn the reply's code blocks into a tool call or a final answer.

    `blocks` arrive already extracted from the reply's `pre code` elements — ChatGPT
    renders fences as DOM nodes, so fence markers are absent from the text and a
    fence regex would find nothing.
    """
    if not blocks:
        return Final(text)
    if len(blocks) > 1:
        return Malformed(f"expected exactly one json block, got {len(blocks)}")
    try:
        payload = json.loads(blocks[0])
    except ValueError as exc:
        return Malformed(f"the json block did not parse: {exc}")
    if not isinstance(payload, dict) or "tool" not in payload:
        return Malformed('the json block has no "tool" key')
    name = payload["tool"]
    if name == FINAL_TOOL:
        if not isinstance(payload.get("text"), str):
            return Malformed('a final block needs a "text" string')
        return Final(payload["text"])
    if name not in tool_names(tools):
        return Malformed(f"unknown tool {name!r}")
    arguments = payload.get("arguments")
    if not isinstance(arguments, dict):
        return Malformed(f'tool {name!r} needs an "arguments" object')
    return ToolCall(name, arguments)


def correction_prompt(reason: str) -> str:
    return (
        f"Your last reply broke the protocol: {reason}.\n"
        "Send the same answer again as EXACTLY ONE fenced json block and nothing "
        'else — either {"tool": "<name>", "arguments": {}} or '
        '{"tool": "final", "text": "<answer>"}.'
    )