"""Shared agent core: chat template, instructions, tool-call parsing, episode logic.

This exact module drives (1) RL training rollouts, (2) every evaluation, and
(3) the shipped standalone tool — zero train/inference mismatch.

Conventions (inherited from the reference client):
- NO system role: instructions are merged into the first user message.
- Strict user/assistant alternation; tool output comes back as the next user
  message, prefixed "Tool output: ".
- Plain-text tool calls, types "bash" and "final", lenient regex parsing, and
  parse errors are fed back to the model as the next user message.

Format: reasoning goes in <think>...</think>, the tool call in <answer>...</answer>
(the reference's "Thoughts:" field is dropped — it is redundant with <think>).
Assistant turns end with "</answer>" followed by EOS. Past assistant turns are
rendered with the <think> block stripped (saves context; template handles it).
"""
import re

# ----------------------------------------------------------------------------
# Chat template
# ----------------------------------------------------------------------------

PREAMBLE = (
    "A conversation between User and Assistant. The user asks a question, and the "
    "Assistant solves it. The assistant first thinks about the reasoning process in "
    "the mind and then provides the user with the answer. The reasoning process and "
    "answer are enclosed within <think>...</think> and <answer>...</answer> tags, "
    "respectively, i.e., <think> reasoning process here </think> "
    "<answer>answer here </answer>."
)

# Past assistant messages are rendered with everything up to and including
# </think> removed. The generation prompt ends with "Assistant: <think>" + newline.
CHAT_TEMPLATE = (
    "{{- '" + PREAMBLE.replace("'", "\\'") + "\\n\\n' -}}"
    "{%- for message in messages -%}"
    "{%- if message['role'] == 'user' -%}"
    "{{- 'User: ' + message['content'] + '\\n\\n' -}}"
    "{%- elif message['role'] == 'assistant' -%}"
    "{%- set content = message['content'] -%}"
    "{%- if '</think>' in content -%}"
    "{%- set content = content.split('</think>')[-1] -%}"
    "{%- endif -%}"
    "{{- 'Assistant: ' + content.strip() + '\\n\\n' -}}"
    "{%- endif -%}"
    "{%- endfor -%}"
    "{%- if add_generation_prompt -%}"
    "{{- 'Assistant: <think>\\n' -}}"
    "{%- endif -%}"
)

THINK_PREFIX = "<think>\n"  # generation prompt ends with this; completions start after it
STOP_STRING = "</answer>"

# Sampling defaults — used for RL rollouts, every eval, and shipped generation_config.
GEN_TEMPERATURE = 1.0
GEN_TOP_P = 0.95
MAX_TURN_TOKENS = 1024  # shared training/eval/shipped-tool per-turn output budget


def render_messages(messages, add_generation_prompt: bool = True) -> str:
    """Render messages exactly like CHAT_TEMPLATE, without depending on Jinja.

    Offline vLLM/Transformers calls use this function; server-mode vLLM uses the
    tokenizer chat template below. Keep both byte-identical.
    """
    parts = [PREAMBLE + "\n\n"]
    for message in messages:
        role = message["role"]
        content = message["content"]
        if role == "user":
            parts.append("User: " + content + "\n\n")
        elif role == "assistant":
            if "</think>" in content:
                content = content.split("</think>")[-1]
            parts.append("Assistant: " + content.strip() + "\n\n")
    if add_generation_prompt:
        parts.append("Assistant: <think>\n")
    return "".join(parts)

# ----------------------------------------------------------------------------
# Instructions (merged into the first user message)
# ----------------------------------------------------------------------------

TOOL_INSTRUCTIONS = """You are an AI agentic coding assistant. You complete the user's task by calling tools step by step, one tool call per turn.

First reason inside <think>...</think>, then put exactly one tool call inside <answer>...</answer> in exactly this format:

<answer>Tool type: [bash or final]
Tool query: [the shell command to run, or the final answer]</answer>

Available tools:
1. bash: runs a shell command in an isolated Linux sandbox and returns its output (e.g. write files with heredocs, run python, test code). The sandbox is stateless: the working directory and variables reset between calls, so chain commands with && where needed.
2. final: ends the task; put the complete final answer in Tool query.

The next user message will contain the tool output as "Tool output: ...". Keep outputs small (use head/tail). After your final tool call the episode ends."""


def coding_task_message(statement: str, starter_code: str = "", max_turns: int = 2) -> str:
    """First user message for a competitive-programming task (training AND eval)."""
    if starter_code and starter_code.strip():
        req = (
            "Write a Python 3 solution that completes the following starter code, "
            "keeping the exact same class name and method signature:\n"
            "```python\n" + starter_code.strip() + "\n```"
        )
    else:
        req = ("Write a complete Python 3 program that reads the input from stdin "
               "and prints the required output to stdout.")
    n_bash = max_turns - 1
    one_turn_rule = (
        "- This is a one-turn episode: NEVER call `bash`; submit the complete Python solution "
        "with your single `final` tool call immediately.\n"
        if max_turns == 1 else ""
    )
    return f"""{TOOL_INSTRUCTIONS}

Task: solve this competitive programming problem in Python 3.

{statement.strip()}

{req}

Rules:
- You have at most {max_turns} tool calls in total, so at most {n_bash} bash call{"s" if n_bash != 1 else ""}; you may use bash to test your solution on the example input before submitting.
{one_turn_rule}- Your last tool call must be "final" and its Tool query must contain ONLY your complete solution inside one ```python ... ``` code block.
- The solution is judged against hidden tests; handle all edge cases and stay within a few seconds of runtime."""


PARSE_ERROR_FEEDBACK = """Parsing Error: {err}
Please repeat, strictly matching this format:

<answer>Tool type: [bash or final]
Tool query: [the shell command, or the final answer]</answer>"""

# ----------------------------------------------------------------------------
# Parsing
# ----------------------------------------------------------------------------


def parse_action(text: str):
    """Parse an assistant completion into (tool_type, tool_query, error).

    `text` is the completion text (reasoning + <answer> block, with or without
    the closing tag). Lenient: falls back to scanning the whole text.
    """
    m = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
    if m:
        seg = m.group(1)
    elif "<answer>" in text:
        seg = text.split("<answer>", 1)[1]
    else:
        seg = text.split("</think>")[-1]

    tt = re.search(r"Tool type:\s*(.*?)(?=\n\s*Tool query:|$)", seg, re.DOTALL | re.IGNORECASE)
    tq = re.search(r"Tool query:\s*(.*)", seg, re.DOTALL | re.IGNORECASE)
    tool_type = tt.group(1).strip().lower() if tt else None
    # The reference protocol presents its choices as ``[bash or final]`` and
    # base completions commonly echo the brackets (``Tool type: [final]``).
    # Treat that as the same lenient plain-text call rather than wasting a turn.
    if tool_type:
        tool_type = tool_type.strip(" \t[](){}<>`*_\"'")
    tool_query = tq.group(1).strip() if tq else None

    if not tool_type or tool_query is None or tool_query == "":
        return None, None, ("Could not parse 'Tool type' or 'Tool query' from the response. "
                            "Ensure you output an <answer> block with 'Tool type:' and 'Tool query:'.")
    if tool_type not in ("bash", "final"):
        return None, None, f"Unknown tool type '{tool_type}'; must be 'bash' or 'final'."
    if tool_type == "bash":
        # strip a markdown fence around the command, reference-style
        tool_query = re.sub(r"^```[a-zA-Z]*\n", "", tool_query)
        tool_query = re.sub(r"\n```$", "", tool_query)
        tool_query = tool_query.strip()
    return tool_type, tool_query, None


def extract_code(final_answer: str) -> str:
    """Extract the python solution from a final answer (last ```python block, else raw)."""
    blocks = re.findall(r"```(?:python|py)?\n(.*?)```", final_answer, re.DOTALL)
    if blocks:
        return blocks[-1].strip()
    return final_answer.strip()


def normalize_completion(content: str, finish_reason: str = "stop") -> str:
    """Make server-mode completions byte-identical to training rollouts: the stop
    string is part of the turn. If the server stripped it (or stopped at EOS right
    after it), re-append."""
    if STOP_STRING not in content and finish_reason == "stop" and "<answer>" in content:
        content = content.rstrip() + STOP_STRING
    return content


# ----------------------------------------------------------------------------
# Episode state machine (transport-agnostic)
# ----------------------------------------------------------------------------


class Episode:
    """One agentic episode with strict user/assistant alternation.

    Use: messages -> (render+generate outside) -> step(completion) which returns
    ("bash", cmd) | ("final", answer) | ("continue", None) after appending messages.
    After executing a bash command, call add_tool_output(out).
    """

    def __init__(self, first_user_msg: str, max_turns: int = 2, meta=None):
        self.messages = [{"role": "user", "content": first_user_msg}]
        self.max_turns = max_turns
        self.turns = 0
        self.done = False
        self.final_answer = None
        self.end_reason = None
        self.meta = meta or {}

    def step(self, completion: str):
        """completion: model output text (starts after 'Assistant: <think>\\n')."""
        assert not self.done
        content = THINK_PREFIX + completion
        self.messages.append({"role": "assistant", "content": content})
        self.turns += 1
        tool_type, tool_query, err = parse_action(completion)

        if err is not None:
            if self.turns >= self.max_turns:
                self.done, self.end_reason = True, "parse_error"
                return ("continue", None)
            self.add_tool_output(PARSE_ERROR_FEEDBACK.format(err=err))
            return ("continue", None)

        if tool_type == "final":
            self.done, self.end_reason = True, "final"
            self.final_answer = tool_query
            return ("final", tool_query)

        # bash
        if self.turns >= self.max_turns:
            self.done, self.end_reason = True, "out_of_turns"
            return ("continue", None)
        return ("bash", tool_query)

    def add_tool_output(self, output: str):
        self.messages.append({"role": "user", "content": f"Tool output: {output}"})


def run_episode(generate_fn, bash_fn, first_user_msg: str, max_turns: int = 2, meta=None):
    """Sequential episode driver (eval + shipped tool). generate_fn(messages)->completion."""
    ep = Episode(first_user_msg, max_turns=max_turns, meta=meta)
    while not ep.done:
        completion = generate_fn(ep.messages)
        kind, arg = ep.step(completion)
        if kind == "bash":
            ep.add_tool_output(bash_fn(arg))
    return ep


def openai_generate_fn(client, model: str, temperature: float = GEN_TEMPERATURE,
                       top_p: float = GEN_TOP_P, max_tokens: int = MAX_TURN_TOKENS):
    """generate_fn over an OpenAI-compatible chat endpoint (vLLM server). The server
    applies this repo's chat template (shipped in the tokenizer)."""
    def fn(messages):
        kwargs = dict(model=model, messages=messages, temperature=temperature,
                      top_p=top_p, max_tokens=max_tokens, stop=[STOP_STRING])
        try:
            resp = client.chat.completions.create(
                **kwargs, extra_body={"include_stop_str_in_output": True})
        except Exception:
            resp = client.chat.completions.create(**kwargs)
        choice = resp.choices[0]
        return normalize_completion(choice.message.content or "", choice.finish_reason)
    return fn