"""
Utility functions for agent-related operations.

This module provides functions for handling tool descriptions,
building system prompts, and rendering prompt templates.
"""
import json
from typing import Any, Dict, Iterable, List, Optional, Tuple
import re
import yaml
from jinja2 import Environment
from mcp.types import Tool


SUMMARIZE_PROMPT = (
    "***Important: You have reached the maximum number of interaction turns. "
    "Therefore, you have to output the final answer with the specified format "
    "without calling any tools in the final response.***"
)


def get_fc_cnt_list(scheduler_mode: dict) -> List[int]:
    """
    Given a scheduler_mode dict where keys are iteration indices and values are counts,
    return a list whose length matches the highest iteration,
    with each segment filled according to the provided count.

    Example:
        {25: 3, 50: 2, 100: 1}
        -> [3]*25 + [2]*25 + [1]*50  (total length = 100)
    """
    # Sort scheduler_mode keys in ascending order
    # filter the keys that can be converted to int
    sorted_iters = sorted(k for k in scheduler_mode.keys() if isinstance(k, int))
    fc_cnt_list = []
    prev_iter = 0
    for _idx, iter_num in enumerate(sorted_iters):
        count = scheduler_mode[iter_num]
        span = iter_num - prev_iter
        fc_cnt_list.extend([count] * span)
        prev_iter = iter_num
    return fc_cnt_list


def get_iteration_prompt_content(
    iter_num: int,
    max_iterations: int,
    scheduler_mode: Optional[Dict],
) -> Optional[str]:
    """
    Return the user prompt for this iteration: summarize on last step,
    or countdown + optional scheduler hint when not last.
    Returns None when nothing should be appended.
    """
    is_last_iter = iter_num == max_iterations - 1
    steps_remaining = max_iterations - iter_num
    countdown = f"You have {steps_remaining} steps remaining. Please continue."

    result: Optional[str] = None
    if is_last_iter:
        result = SUMMARIZE_PROMPT
    elif scheduler_mode:
        if "dynamic" in scheduler_mode:
            dynamic_prompt = scheduler_mode["dynamic"]
            result = f"{countdown} {dynamic_prompt}"
        else:
            fc_cnt_list = get_fc_cnt_list(scheduler_mode)
            fc_cnt_this_iter = fc_cnt_list[iter_num]
            if fc_cnt_this_iter == 1:
                result = (
                    f"{countdown} At next step, if you need to make function calls, "
                    f"you MUST make exactly {fc_cnt_this_iter} function calls in a single response."
                )
            else:
                result = (
                    f"{countdown} At next step, if you need to make function calls, "
                    f"you MUST make at least {fc_cnt_this_iter} but not more than {fc_cnt_this_iter + 1} "
                    "function calls in a single response to gather information extensively."
                )
    elif iter_num > 0:
        result = countdown
    return result


def get_tools_description(tools: Dict[str, List[Tool]]) -> str:
    """
    Generate a formatted description of the specified tools.

    This function creates a detailed description of each tool, including
    the server name, tool name, description, and arguments.

    Args:
        tools (Dict[str, List[Tool]]): A dictionary of tools, where keys are
            server names and values are lists of Tool objects.

    Returns:
        str: A formatted string containing descriptions of all tools.
    """
    descriptions = []
    for server_name, tool_list in tools.items():
        for tool in tool_list:
            args = []
            if "properties" in tool.inputSchema:
                for param_name, param_info in tool.inputSchema["properties"].items():
                    info = "\n".join(
                        [
                            "    " + line
                            for line in yaml.dump(
                                param_info, sort_keys=False, indent=2
                            ).split("\n")
                        ]
                    )
                    arg = f"- {param_name}:\n{info}".strip()
                    if param_name in tool.inputSchema.get("required", []):
                        arg += "\n    required: true"
                    args.append(arg.strip())
            lines = [line for line in tool.description.split("\n") if line.strip()]
            arguments = f"\n{chr(10).join(args)}" if args else " No arguments"
            description = (
                f"Server: {server_name}\n"
                f"Tool: {tool.name}\n"
                f"Description:\n{chr(10).join(lines)}\n"
                f"Arguments:{arguments}"
            )
            descriptions.append(description)
    return "\n\n".join(descriptions).strip()


def build_system_prompt(
    system_prompt_template: str,
    tool_prompt_template: str = "",
    tools: Optional[Dict[str, List[Tool]]] = None,
    include_tool_description: Optional[bool] = True,
    **kwargs,
) -> str:
    """
    Build an agent system prompt using provided templates and tools.

    This function combines system and tool prompt templates with tool descriptions
    to create a comprehensive system prompt for an agent.

    Args:
        system_prompt_template (str): The template for the system prompt. If it
            ends with ".j2", it's treated as a path to a Jinja2 template file.
        tool_prompt_template (str, optional): The template for the tool prompt. If it
            ends with ".j2", it's treated as a path to a Jinja2 template file.
        tools (Dict[str, List[Tool]], optional): A dictionary of tools, where keys
            are server names and values are lists of Tool objects.
        include_tool_description (bool, optional): Whether to include tool descriptions
            in the prompt if tools exist.
        **kwargs: Additional keyword arguments to be passed to the template rendering.

    Returns:
        str: The rendered system prompt.

    Note:
        If both tool_prompt_template and tools are provided, a tools prompt will be
        generated and included in the final system prompt.
    """
    if system_prompt_template.endswith(".j2"):
        with open(system_prompt_template, "r", encoding="utf-8") as f:
            system_prompt_template = f.read()
    if tool_prompt_template.endswith(".j2"):
        with open(tool_prompt_template, "r", encoding="utf-8") as f:
            tool_prompt_template = f.read()

    tools_prompt = ""
    tools_description = get_tools_description(tools) if tools else ""

    if include_tool_description and tool_prompt_template and tools_description:
        env = Environment(trim_blocks=True, lstrip_blocks=True)
        template = env.from_string(tool_prompt_template)
        kwargs.update({"TOOLS_DESCRIPTION": tools_description})
        tools_prompt = template.render(**kwargs)

    env = Environment(trim_blocks=True, lstrip_blocks=True)
    template = env.from_string(system_prompt_template)
    if tools_prompt:
        kwargs.update({"TOOLS_PROMPT": tools_prompt})
    return template.render(**kwargs).strip()


def render_prompt_template(prompt_template: str, **kwargs):
    """
    Render a prompt using a given template and variables.

    This function takes a prompt template (either as a string or a file path)
    and renders it using the provided variables.

    Args:
        prompt_template (str): The prompt template string or path to a .j2 template file.
        **kwargs: Variables to be used in template rendering.

    Returns:
        str: The rendered prompt.

    Note:
        If prompt_template ends with ".j2", it's treated as a path to a Jinja2 template file.
    """
    if prompt_template.endswith(".j2"):
        with open(prompt_template, "r", encoding="utf-8") as f:
            prompt_template = f.read()
    env = Environment(trim_blocks=True, lstrip_blocks=True)
    template = env.from_string(prompt_template)
    return template.render(**kwargs).strip()


def _sanitize_ident(name: str) -> str:
    """Relaxed Harmony-style identifier: keep letters/digits/_/-; replace others with '_'."""
    out = []
    for ch in name:
        if ch.isalnum() or ch in ("_", "-"):
            out.append(ch)
        else:
            out.append("_")
    s = "".join(out)
    return s or "_"  # avoid empty


_TS_IDENT_RE = re.compile(r"^[A-Za-z_]\w*$")


def _format_ts_prop_key(key: str) -> str:
    """
    Keep hyphens for object property keys by quoting when not a bare identifier.
    """
    if _TS_IDENT_RE.match(key):
        return key
    return json.dumps(key, ensure_ascii=False)  # e.g. "API-post-page"


def _jsonschema_to_ts(schema: Optional[Dict[str, Any]]) -> str:  # pylint: disable=too-many-return-statements
    """Tiny JSON-Schema → TS-ish mapper for Harmony tool arg blocks."""
    if not schema:
        return "any"

    def to_ts(s: Any) -> str:  # pylint: disable=too-many-return-statements
        if s is None:
            return "any"
        if isinstance(s, bool):
            return "any" if s else "never"
        if not isinstance(s, dict):
            return "any"

        t = s.get("type")
        if isinstance(t, list):
            return " | ".join(to_ts({"type": x}) for x in t)

        if "enum" in s and isinstance(s["enum"], list):
            vals = s["enum"]
            if all(isinstance(v, str) for v in vals):
                return " | ".join(json.dumps(v, ensure_ascii=False) for v in vals)

        if t == "string":
            return "string"
        if t in ("number", "integer"):
            return "number"
        if t == "boolean":
            return "boolean"
        if t == "array":
            items = s.get("items") or {}
            return f"{to_ts(items)}[]"

        if t == "object" or "properties" in s or "additionalProperties" in s:
            props = s.get("properties") or {}
            required = set(s.get("required") or [])
            lines: List[str] = []
            for k, v in props.items():
                k_id = _format_ts_prop_key(k)  # <— preserve hyphens by quoting if needed
                opt = "" if k in required else "?"
                desc = v.get("description")
                if isinstance(desc, str) and desc.strip():
                    lines.append(f"// {desc.strip()}")
                lines.append(f"{k_id}{opt}: {to_ts(v)},")
            if "additionalProperties" in s:
                ap = s["additionalProperties"]
                if ap not in (False, None):
                    ap_ts = to_ts(ap)
                    lines.append(f"[key: string]: {ap_ts},")
            return "{\n" + ("\n".join(lines)) + "\n}"

        for key in ("oneOf", "anyOf"):
            if isinstance(s.get(key), list) and s[key]:
                return " | ".join(to_ts(sub) for sub in s[key])

        if isinstance(s.get("allOf"), list) and s["allOf"]:
            return " & ".join(to_ts(sub) for sub in s["allOf"])

        return "any"

    return to_ts(schema)


def render_tools_namespace(
    tools_by_server: Dict[str, Iterable[Any]],
    *,
    title: str = "functions",
) -> str:
    """
    Render a Harmony 'namespace functions { ... }' that aggregates tools from multiple MCP servers.
    - Each tool is emitted as:  type {server}__{tool} = (_: <TS>) => any;
    - Tool objects may be dicts (OpenAI-style) or objects with attributes: name, description, inputSchema.
    - JSON-Schema is taken from: tool['function']['parameters'] OR tool.inputSchema
    """
    lines: List[str] = ["# Tools", f"## {title}", "", f"namespace {title} {{", ""]

    def extract_one(tool_obj: Any) -> Optional[Dict[str, Any]]:
        """Extract name/description/schema from a tool object or function dict."""
        if isinstance(tool_obj, dict) and tool_obj.get("type") == "function":
            f = tool_obj.get("function") or {}
            return {
                "name": f.get("name"),
                "description": f.get("description"),
                "schema": f.get("parameters"),
            }
        name = getattr(tool_obj, "name", None)
        if name:
            return {
                "name": name,
                "description": getattr(tool_obj, "description", None),
                "schema": getattr(tool_obj, "inputSchema", None),
            }
        return None

    for server in sorted(tools_by_server.keys()):
        tools = list(tools_by_server[server] or [])
        if not tools:
            continue
        lines.append(f"// --- server: {server} ---")
        for t in tools:
            meta = extract_one(t)
            if not meta or not meta.get("name"):
                continue
            tool_name = meta["name"]
            # Hyphens are preserved here
            fq_name = f"{_sanitize_ident(server)}__{_sanitize_ident(tool_name)}"
            desc = (meta.get("description") or "").strip()
            if desc:
                lines.append(f"// {desc}")
            else:
                lines.append(f"// Tool from '{server}': {tool_name}")
            ts_args = _jsonschema_to_ts(meta.get("schema"))
            lines.append(f"type {fq_name} = (_: {ts_args}) => any;")
            lines.append("")
        lines.append("")

    lines.append(f"}} // namespace {title}")
    return "\n".join(lines)


def render_harmony_chain(
    developer_instructions: str,
    user_first_message: str,
    # Each round MUST have analysis; tool_call/tool_result are optional
    rounds: List[Dict[str, Any]],
    system_identity: str = "You are ChatGPT, a large language model trained by OpenAI.",
    knowledge_cutoff: str = "2024-06",
    reasoning: str = "low",
    tools_namespace_ts: Optional[str] = "",  # string from render_tools_namespace(...)
) -> str:
    """
    Format a multi-round tool-calling transcript in Harmony

    Parameters:
    - developer_instructions: your instructions for the assistant/developer message.
    - tools_namespace_ts: TypeScript namespace block defining tools.
    - user_first_message: the user’s initial input.
    - rounds: a list of steps; each includes at least:
         {
           "analysis": str,        # required
           "tool_call": {"name": str, "arguments": Any},  # optional
           "tool_result": Any                # optional
         }
    """
    parts: List[str] = []

    # --- system message
    sys_lines = [
        "<|start|>system<|message|>" + system_identity,
        f"Knowledge cutoff: {knowledge_cutoff}",
        "",
        f"Reasoning: {reasoning}",
        "",
        "# Valid channels: analysis, commentary, final. Channel must be included for every message.",
    ]
    if tools_namespace_ts:
        sys_lines.append("Calls to these tools must go to the commentary channel: 'functions'.")
    sys_lines.append("<|end|>")
    parts.append("\n".join(sys_lines) + "\n")

    # --- developer + tools
    dev = ["\n<|start|>developer<|message|># Instructions\n", developer_instructions.strip(), ""]
    if tools_namespace_ts:
        dev += [tools_namespace_ts, ""]
    dev.append("<|end|>")
    parts.append("\n".join(dev))

    # --- user message
    parts.append(f"<|start|>user<|message|>{user_first_message}<|end|>")

    # --- rounds of tool calls with required analysis
    for idx, step in enumerate(rounds):
        # enforce required fields
        if "analysis" not in step:
            raise ValueError(f"Round {idx} missing 'analysis'.")

        # 1. Assistant analysis
        parts.append("<|start|>assistant<|channel|>analysis<|message|>")
        parts.append(step["analysis"].strip())
        parts.append("<|end|>\n")

        # 2. Optional Tool call (assistant commentary)
        tool_call = step.get("tool_call") if isinstance(step.get("tool_call"), dict) else None
        name = tool_call.get("name") if tool_call else None
        if name:
            args = tool_call.get("arguments", {})
            parts.append(
                f"<|start|>assistant<|channel|>commentary to=functions.{name} "
                f"<|constrain|>json<|message|>{json.dumps(args, ensure_ascii=False)}<|call|>\n"
            )

            # 3. Optional Tool result (tool → assistant commentary)
            if "tool_result" in step and step["tool_result"] is not None:
                parts.append(
                    (
                        f"<|start|>functions.{name} to=assistant<|channel|>commentary<|message|>"
                        f"{step['tool_result']}<|end|>\n"
                    )
                )
    parts.append("<|start|>assistant<|channel|>analysis<|message|>")
    # Join everything
    return "".join(parts)


# ---------- Regex ----------
COMMENTARY_HEADER_RE = re.compile(
    r"<\|start\|>\s*assistant.*?<\|channel\|>\s*commentary\b.*?<\|message\|>",
    re.DOTALL,
)
FINAL_HEADER_RE = re.compile(
    r"<\|start\|>\s*assistant.*?<\|channel\|>\s*final\b.*?<\|message\|>",
    re.DOTALL,
)
TO_NAME_RE = re.compile(r"\bto=functions\.([^\s<{]+)")
# Stop at any structural tag that can follow args, including optional <|call|>
NEXT_TAG_RE = re.compile(r"(<\|start\|>|<\|end\|>|<\|call\|>)", re.DOTALL)

# Compact forms, e.g.:
# assistantcommentary to=functions.server__tool json{...}
# assistantfinal {...}
COMPACT_COMMENTARY_RE = re.compile(
    r"\bassistant(?:<\|channel\|>)?commentary\b"
    r"(?P<after>.*?)"
    r"(?=("
    r"\bassistant(?:<\|channel\|>)?commentary\b|"
    r"\bassistant(?:<\|channel\|>)?final\b|"
    r"<\|start\|>|"
    r"<\|end\|>|"
    r"\Z"
    r"))",
    re.DOTALL,
)

COMPACT_FINAL_RE = re.compile(
    r"\bassistant(?:<\|channel\|>)?final\b"
    r"(?P<after>.*?)"
    r"(?=("
    r"\bassistant(?:<\|channel\|>)?commentary\b|"
    r"\bassistant(?:<\|channel\|>)?final\b|"
    r"<\|start\|>|"
    r"<\|end\|>|"
    r"\Z"
    r"))",
    re.DOTALL,
)

# ---------- Balanced JSON/Array scanner ----------
def _scan_balanced_json_like(s: str, start_idx: int) -> Optional[Dict[str, Any]]:
    """
    Scan starting at either '{' or '[' and return a parsed JSON value (or raw text) plus end index.
    Handles strings/escapes and nested braces/brackets.
    Returns:
        {"raw": <str>, "end": <int>, "value": <Any>} or None
    """
    if start_idx >= len(s) or s[start_idx] not in "{[":
        return None

    opening = s[start_idx]
    closing = "}" if opening == "{" else "]"

    depth, in_str, esc = 0, False, False
    j = start_idx
    while j < len(s):
        ch = s[j]
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
        else:
            if ch == '"':
                in_str = True
            elif ch == opening:
                depth += 1
            elif ch == closing:
                depth -= 1
                if depth == 0:
                    raw = s[start_idx : j + 1]
                    try:
                        val = json.loads(raw)
                    except (json.JSONDecodeError, ValueError):
                        val = raw
                    return {"raw": raw, "end": j + 1, "value": val}
        j += 1
    return None


def _scan_balanced_json(s: str, start_idx: int) -> Optional[Dict[str, Any]]:  # type: ignore[override]
    """Compatibility wrapper for _scan_balanced_json_like (object/array)."""
    return _scan_balanced_json_like(s, start_idx)


# ---------- Helpers ----------
def _split_server_tool(tool_name: Optional[str]) -> Tuple[str, str]:
    """Split a fully-qualified tool name into (server, tool)."""
    if not tool_name:
        return "", ""
    for sep in ("__", ".", "_"):
        if sep in tool_name:
            a, b = tool_name.split(sep, 1)
            return a, b
    return "", tool_name


def _slice_until_next_tag(text: str, start: int) -> Tuple[str, int]:
    """Return (slice, next_index) from start until next structural tag (or end)."""
    m = NEXT_TAG_RE.search(text, start)
    if not m:
        return text[start:].strip(), len(text)
    return text[start:m.start()].strip(), m.start()


def _normalize_args(val: Any) -> Dict[str, Any]:
    """Ensure arguments is a dict."""
    return val if isinstance(val, dict) else {"value": val}


# ---------- Parsers ----------
def parse_analysis(text: str) -> str:
    """
    Return the analysis as a single string.
    Analysis is the plain text at the top, before the first structural marker:
      earliest of <|start|>, 'assistantcommentary', or <|end|>.
    """
    idx_start = text.find("<|start|>")
    m_compact = re.search(r"\bassistantcommentary\b", text)
    idx_compact = m_compact.start() if m_compact else -1
    idx_end = text.find("<|end|>")

    candidates = [i for i in (idx_start, idx_compact, idx_end) if i != -1]
    cutoff = min(candidates) if candidates else -1

    prefix = text if cutoff == -1 else text[:cutoff]
    return prefix.strip()


def _parse_tool_call_harmony(text: str) -> List[Dict[str, Any]]:
    """Parse tool calls in full Harmony tag format."""
    calls: List[Dict[str, Any]] = []
    for m in COMMENTARY_HEADER_RE.finditer(text):
        header = m.group(0)  # includes <|message|>
        name_m = TO_NAME_RE.search(header)
        tool_name = name_m.group(1) if name_m else None
        server, tool = _split_server_tool(tool_name)

        i, n = m.end(), len(text)
        while i < n and text[i].isspace():
            i += 1

        if i < n and text[i] in "{[":
            parsed = _scan_balanced_json_like(text, i)
            args_val: Any = parsed["value"] if parsed else _slice_until_next_tag(text, i)[0]
        else:
            args_val = _slice_until_next_tag(text, i)[0]

        calls.append({"tool_name": tool_name, "server": server, "tool": tool, "arguments": args_val})
    return calls


def _parse_tool_call_compact(text: str) -> List[Dict[str, Any]]:
    """
    Parse blocks that start with 'assistantcommentary' and include 'to=functions.NAME'
    followed by optional 'json' token and a JSON object/array. If none found, args="".
    """
    calls: List[Dict[str, Any]] = []
    for m in COMPACT_COMMENTARY_RE.finditer(text):
        block = m.group(0)
        name_m = TO_NAME_RE.search(block)
        tool_name = name_m.group(1) if name_m else None
        server, tool = _split_server_tool(tool_name)

        start_search_pos = name_m.end() if name_m else 0
        # Find first '{' or '[' after the tool name
        brace_idx = len(block)
        for ch in "{[":
            idx = block.find(ch, start_search_pos)
            if idx != -1:
                brace_idx = min(brace_idx, idx)

        args_val: Any = ""
        if brace_idx != len(block):
            parsed = _scan_balanced_json_like(block, brace_idx)
            args_val = parsed["value"] if parsed else block[brace_idx:].strip()

        calls.append({"tool_name": tool_name, "server": server, "tool": tool, "arguments": args_val})
    return calls


def parse_tool_call(text: str) -> List[Dict[str, Any]]:
    """Parse tool calls from both Harmony and compact styles; normalize arguments to dict."""
    calls = _parse_tool_call_harmony(text)
    calls.extend(_parse_tool_call_compact(text))
    for c in calls:
        c["arguments"] = _normalize_args(c.get("arguments"))
    return calls


# ---------- Final answer (JSON-only) ----------
def _final_json_after_idx(text: str, i: int) -> Optional[Any]:
    """Return a parsed JSON value starting at i if it begins with '{' or '['; else None."""
    n = len(text)
    while i < n and text[i].isspace():
        i += 1
    if i < n and text[i] in "{[":
        parsed = _scan_balanced_json_like(text, i)
        if parsed:
            return parsed["value"]
    return None


def _parse_final_harmony_json(text: str) -> Optional[Any]:
    """Extract JSON payload from Harmony final block."""
    m = FINAL_HEADER_RE.search(text)
    return _final_json_after_idx(text, m.end()) if m else None


def _parse_final_compact_json(text: str) -> Optional[Any]:
    """Extract JSON payload from compact 'assistantfinal' block."""
    m = COMPACT_FINAL_RE.search(text)
    if not m:
        return None
    block = m.group(0)
    after = m.group("after") or ""
    block_start = text.find(block)
    after_start = block_start + (block.find(after) if after else len(block))
    return _final_json_after_idx(text, after_start)


def parse_final(text: str) -> Optional[Any]:
    """
    Returns the JSON value from the final answer:
      - dict/list when the payload is valid JSON or JSON array
      - None if no JSON final is found
    """
    return _parse_final_harmony_json(text) or _parse_final_compact_json(text)


def parse_harmony(text: str) -> Dict[str, Any]:
    """Parse Harmony-formatted transcript into analysis/tool_call/final fields."""
    return {
        "analysis": parse_analysis(text),  # always a string (may be empty)
        "tool_call": parse_tool_call(text),
        "final": parse_final(text),  # dict/list or None
        "raw": text,
    }


# ========== Qwen3 Chat Template Functions ==========

def render_qwen3_react_prompt(
    system_instruction: str,
    tools_description: str,
    user_message: str,
    history: List[Dict[str, Any]],
) -> str:
    """
    Render a ReAct prompt in Qwen3 chat template format for raw completion.

    This function formats the conversation using Qwen3's <|im_start|> and <|im_end|> tags
    for use with raw completion API (not chat completion).

    Args:
        system_instruction: System prompt with instructions
        tools_description: Description of available tools
        user_message: User's query/question
        history: List of previous interaction steps, each containing:
            - thought: The reasoning from the model
            - action (optional): Tool action dict with server/tool/arguments
            - result (optional): Result from tool execution
            - answer (optional): Final answer from model

    Returns:
        str: Formatted prompt in Qwen3 chat template format for raw completion
    """
    parts: List[str] = []

    # System message
    parts.append("<|im_start|>system\n")
    parts.append(system_instruction.strip())
    if tools_description:
        parts.append("\n\n")
        parts.append(tools_description.strip())
    parts.append("<|im_end|>\n")

    # User message
    parts.append("<|im_start|>user\n")
    parts.append(user_message.strip())
    parts.append("<|im_end|>\n")

    # History of interactions
    for step in history:
        parts.append("<|im_start|>assistant\n")

        # Build the JSON response
        response_dict = {"thought": step.get("thought", "")}

        if "answer" in step:
            response_dict["answer"] = step["answer"]
        elif "action" in step:
            response_dict["action"] = step["action"]

        parts.append(json.dumps(response_dict, ensure_ascii=False, indent=2))
        parts.append("<|im_end|>\n")

        # If there's a tool result, add it as a user message
        if "result" in step:
            parts.append("<|im_start|>user\n")
            parts.append(f"Tool execution result:\n{step['result']}")
            parts.append("<|im_end|>\n")

    # Start of assistant's next response
    parts.append("<|im_start|>assistant\n")

    return "".join(parts)


def parse_qwen3_react_response(text: str) -> Dict[str, Any]:
    """
    Parse the model's response in ReAct format.

    Expected JSON format:
    {
        "thought": "reasoning text",
        "action": {  // optional
            "reason": "why use this tool",
            "server": "server-name",
            "tool": "tool-name",
            "arguments": {...}
        }
    }
    OR
    {
        "thought": "reasoning text",
        "answer": "final answer"
    }

    Args:
        text: Raw text response from the model

    Returns:
        Dict with keys: thought, action (optional), answer (optional), raw
    """
    # Clean up the response
    text = text.strip()

    # Remove markdown code fences more robustly
    # Handle cases like: ```json\n{...}\n```<|im_end|>
    # Pattern 1: Remove opening ```json or ```
    if text.startswith("```"):
        # Find the first newline after ```
        newline_idx = text.find("\n")
        if newline_idx > 0:
            # Check if it's ```json, ```json\n, or just ```
            first_line = text[:newline_idx]
            if first_line.startswith("```"):
                text = text[newline_idx + 1:]
        else:
            # No newline, just remove ```
            text = text[3:]
        text = text.strip()

    # Pattern 2: Remove closing ``` (may be followed by <|im_end|> or other text)
    # Handle: ...}\n```<|im_end|> or ...}\n```
    if text.endswith("```"):
        text = text[:-3].strip()
    elif "```" in text:
        # Find the last occurrence of ```
        last_idx = text.rfind("```")
        if last_idx >= 0:
            # Check if it's at the end of a line or followed by <|im_end|>
            remaining = text[last_idx + 3:].strip()
            # If remaining is empty or starts with <|im_end|>, remove the ```
            if not remaining or remaining.startswith("<|im_end|>"):
                text = text[:last_idx].strip()

    # Pattern 3: Remove <|im_end|> if present (should be removed by formatter, but handle it here too)
    text = text.rstrip("<|im_end|>").strip()

    # Remove 'json' prefix if present (after cleaning code fences)
    if text.startswith("json"):
        text = text[4:].strip()

    # Try parsing the cleaned text directly first
    parsed = None
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        # If direct parsing fails, try to extract JSON object
        # This handles cases where there's extra text before/after the JSON
        # Find JSON object boundaries by matching balanced braces
        # Start from the first { and find the matching }
        brace_count = 0
        start_idx = text.find('{')
        if start_idx >= 0:
            for i in range(start_idx, len(text)):
                if text[i] == '{':
                    brace_count += 1
                elif text[i] == '}':
                    brace_count -= 1
                    if brace_count == 0:
                        # Found matching closing brace
                        json_text = text[start_idx:i+1]
                        try:
                            parsed = json.loads(json_text)
                            # Update text to the extracted JSON for consistency
                            text = json_text
                            break
                        except json.JSONDecodeError:
                            # This JSON extraction failed, continue searching
                            pass

    if parsed is not None:
        #########################################################
        # Check if this follows the expected ReAct format (has thought/action/answer)
        has_react_format = any(key in parsed for key in ["thought", "action", "answer"])

        if has_react_format:
            # Standard ReAct format
            result = {
                "thought": parsed.get("thought", ""),
                "action": parsed.get("action"),
                "answer": parsed.get("answer"),
                "raw": text,
            }
        else:
            # Model output doesn't follow ReAct format but is valid JSON
            # Treat the entire parsed object as the final answer
            # This handles cases where model directly outputs: {"place_name": "...", "address": "..."}
            result = {
                "thought": "",
                "action": None,
                "answer": parsed,  # Use the entire parsed JSON as the answer
                "raw": text,
            }
        #########################################################

        return result

    # If parsing fails after all attempts, return the raw text as thought
    return {
        "thought": text,
        "action": None,
        "answer": None,
        "raw": text,
        "error": "JSON decode error: Could not parse JSON from response"
    }
