"""Canonical trajectory format and the one compact dialect we train on first ("ta-v1").

A trajectory is a list of messages:
  {"role": "system"|"user", "content": str}
  {"role": "assistant", "think": str|None, "content": str, "tool_calls": [{"name", "arguments"}]}
  {"role": "tool", "results": [str, ...]}          # one result per call, same order
Rendering is a pure function so other harness dialects (Hermes, Pi, Claude Code, ...) can be
added later as extra renderers over the same records.

ta-v1 text:
  <|im_start|>system\n...<|im_end|>\n<|im_start|>user\n...<|im_end|>\n
  <|im_start|>assistant\n<think>\n...\n</think>\n<tool_call>\n{"name": ..., "arguments": {...}}\n</tool_call><|im_end|>\n
  <|im_start|>tool\n<tool_response>\nRAW RESULT TEXT\n</tool_response><|im_end|>\n
Tool results are raw text, not JSON-escaped, so copying a span from them stays a plain copy.
"""
from __future__ import annotations

import json
import re

from tiny_agent.tools import TOOL_SPECS

SYSTEM_TA_V1 = (
    "You are an agent working in /work. Find information with tools; do not answer from memory. "
    "If the information is not in the workspace, submit NOT_FOUND.\nTools:\n"
    + "\n".join(f"{n}({a}): {d}" for n, a, d in TOOL_SPECS)
)


def render(messages: list[dict], add_generation_prompt: bool = False) -> str:
    out = []
    for m in messages:
        r = m["role"]
        if r in ("system", "user"):
            out.append(f"<|im_start|>{r}\n{m['content']}<|im_end|>\n")
        elif r == "assistant":
            s = "<|im_start|>assistant\n"
            if m.get("think"):
                s += f"<think>\n{m['think']}\n</think>\n"
            if m.get("content"):
                s += m["content"]
            for c in m.get("tool_calls") or []:
                s += "<tool_call>\n" + json.dumps({"name": c["name"], "arguments": c["arguments"]}, ensure_ascii=False) + "\n</tool_call>"
            out.append(s + "<|im_end|>\n")
        elif r == "tool":
            body = "".join(f"<tool_response>\n{x}\n</tool_response>" for x in m["results"])
            out.append(f"<|im_start|>tool\n{body}<|im_end|>\n")
    if add_generation_prompt:
        out.append("<|im_start|>assistant\n")
    return "".join(out)


_CALL = re.compile(r"<tool_call>(.*?)</tool_call>", re.S)
_THINK = re.compile(r"<think>(.*?)</think>", re.S)


def parse_assistant(text: str) -> dict:
    """Parse one generated assistant turn (text up to <|im_end|>). Malformed calls are reported
    as {"error": ...} entries so the environment can return the error as a tool result."""
    text = text.split("<|im_end|>")[0]
    think = _THINK.search(text)
    calls = []
    for raw in _CALL.findall(text):
        try:
            obj = json.loads(raw.strip())
            if not isinstance(obj, dict) or "name" not in obj:
                raise ValueError("missing name")
            args = obj.get("arguments") or {}
            if not isinstance(args, dict):
                raise ValueError("arguments must be an object")
            calls.append({"name": str(obj["name"]), "arguments": args})
        except (json.JSONDecodeError, ValueError) as e:
            calls.append({"error": f"Error: could not parse tool call ({e}). Use JSON: "
                                   '{"name": "...", "arguments": {...}}'})
    content = _CALL.sub("", _THINK.sub("", text)).strip()
    return {"role": "assistant", "think": think.group(1).strip() if think else None,
            "content": content, "tool_calls": calls}


def repeated_calls(messages: list[dict]) -> list[int]:
    """Per assistant turn, how many of its tool calls repeat an earlier call (same name and
    arguments, canonical JSON) with no edit/write since, i.e. calls that cannot see anything new.
    Re-running tests after an edit is not a repeat. (MiMo-V2.6 found mild repetition that goes
    unpenalized gets amplified by RL into tool-call flooding.)"""
    seen, out = set(), []
    for m in messages:
        if m["role"] != "assistant":
            continue
        n = 0
        for c in m.get("tool_calls") or []:
            if "error" in c:
                continue
            if c["name"] in ("edit", "write"):
                seen.clear()
                continue
            key = (c["name"], json.dumps(c["arguments"], sort_keys=True, ensure_ascii=False))
            n += key in seen
            seen.add(key)
        out.append(n)
    return out