"""XML / channel-style tool-call parsing and native/XML merge helpers.

Depends on agent_loop_tool_repair for JSON + path repair utilities.
"""
from __future__ import annotations

import ast
import hashlib
import json
import re
from typing import Any

from agent_loop_tool_repair import (
    _finalize_xml_parsed_arguments_string,
    _json_loads_lenient,
    _preclean_gemma_tool_call_brace_args,
    _quote_json_value,
    _repair_artifact_path_values_in_params,
    _repair_artifact_path_values_in_tool_args_dict,
    _salvage_json_tool_call,
)
from agent_tool_format import GUARD_TOOL_CALL_SUFFIX

_TOOL_CALL_BLOCK_RE = re.compile(
    r"<tool_call>\s*([\s\S]*?)\s*</tool_call>"
    r"|<\|tool_call>\s*([\s\S]*?)\s*<tool_call\|>",
    re.IGNORECASE,
)
_TOOL_CALL_TAG_RE = re.compile(r"<\s*tool_call\b|<\|tool_call>", re.IGNORECASE)
_TOOL_CALL_OPEN_TAGS = ("<tool_call>", "<|tool_call>")
_TOOL_CALL_CLOSE_TAGS = ("</tool_call>", "<tool_call|>")
_PARAM_CLOSE_RE = re.compile(r"</parameter\s*>", re.IGNORECASE)
_FUNC_CLOSE_RE = re.compile(r"</(?:function|invoke)\s*>", re.IGNORECASE)

# Bare Gemma channel-style calls in assistant content (no <tool_call> wrapper).
BARE_GEMMA_CALL_RE = re.compile(
    r"call:\s*\w+\s*(?:[\({][\s\S]*?[\}\)]|\([\s\S]*?\))",
    re.IGNORECASE,
)


def _parse_gemma_delimited_args(raw: str) -> str:
    """Parse Gemma ``call:tool{...}`` / ``(...)`` / mixed ``(...}`` args to JSON."""
    s = (raw or "").strip()
    if not s:
        return "{}"

    if len(s) >= 2 and s[0] in "{(" and s[-1] in "})":
        inner = s[1:-1].strip()
    else:
        inner = s

    inner = re.sub(r'<\|["\u2018\u2019\u201c\u201d\']\|>', '"', inner)

    if not inner:
        return "{}"

    if re.search(r"\w+\s*=", inner):
        try:
            tree = ast.parse("_f({})".format(inner), mode="eval")
            literals = {"true": True, "false": False, "null": None, "none": None}
            out: dict[str, Any] = {}
            for kw in tree.body.keywords:
                if kw.arg is None:
                    continue
                val = kw.value
                if isinstance(val, ast.Name) and val.id.lower() in literals:
                    out[kw.arg] = literals[val.id.lower()]
                else:
                    out[kw.arg] = ast.literal_eval(val)
            return json.dumps(out)
        except Exception:
            pass

    flat = re.match(r"^(\w+)\s*:\s*(.+)$", inner, re.DOTALL)
    if flat:
        key = flat.group(1)
        val_raw = flat.group(2).strip()
        if val_raw.startswith('"'):
            try:
                val: Any = json.loads(val_raw)
            except json.JSONDecodeError:
                val = val_raw
        else:
            val = val_raw
        return json.dumps({key: val})

    raw_args = inner if inner.startswith("{") else "{" + inner + "}"
    raw_args = _preclean_gemma_tool_call_brace_args(raw_args)
    if raw_args:
        try:
            return json.dumps(json.loads(raw_args))
        except (json.JSONDecodeError, ValueError):
            repaired = re.sub(r'(?<=[{,])\s*(\w+)\s*:', r'"\1":', raw_args)
            try:
                return json.dumps(json.loads(repaired))
            except (json.JSONDecodeError, ValueError):
                repaired2 = re.sub(
                    r":\s*([^{}\[\]\"',\s][^{}\[\]\",}]*)", _quote_json_value, repaired
                )
                try:
                    return json.dumps(json.loads(repaired2))
                except (json.JSONDecodeError, ValueError):
                    pass
    return "{}"


def _extract_xml_params_robust(
    block: str,
    open_pattern: str,
    key_extractor,
) -> dict[str, str]:
    """Extract parameter key/value pairs from an XML tool-call block.

    For each ``<parameter...>`` opener, looks for a matching ``</parameter>`` within
    the segment up to the next opener or ``</function>``/end-of-block.  When the
    closing tag is missing (common with local/MLX models that drop it after long
    content), the full segment is used instead.
    """
    params: dict[str, str] = {}
    open_re = re.compile(open_pattern)
    openers = list(open_re.finditer(block))
    if not openers:
        return params

    for i, po in enumerate(openers):
        key = key_extractor(po)
        start = po.end()
        if i + 1 < len(openers):
            boundary = openers[i + 1].start()
        else:
            fc = _FUNC_CLOSE_RE.search(block[start:])
            boundary = (start + fc.start()) if fc else len(block)

        segment = block[start:boundary]
        close = _PARAM_CLOSE_RE.search(segment)
        if close:
            val = segment[: close.start()].strip()
        else:
            val = segment.strip()
            val = re.sub(
                r"</(?:function|parameter|invoke|tool_call)\s*>\s*$",
                "",
                val,
                flags=re.IGNORECASE,
            ).strip()

        if val:
            params[key] = val
    return params

def _parse_xml_tool_call_block(
    block: str,
    idx: int,
    recovery_context: str | None = None,
) -> dict | None:
    """Parse one <tool_call> payload block into canonical call dict."""
    block = (block or "").strip()
    if not block:
        return None

    # --- Format 3: JSON inside <tool_call> ---
    if block.startswith("{"):
        try:
            obj = _json_loads_lenient(block)
            name = obj.get("name") or obj.get("function", "")
            args = obj.get("arguments") or obj.get("parameters") or {}
            if isinstance(args, dict):
                args = _repair_artifact_path_values_in_tool_args_dict(
                    str(name or "").strip(), args, block, recovery_context=recovery_context
                )
                args = json.dumps(args)
            elif not isinstance(args, str):
                args = json.dumps(args)
            if name:
                return {
                    "id": f"xmltc_{idx}_{hashlib.sha256(name.encode()).hexdigest()[:8]}",
                    "name": name,
                    "arguments": _finalize_xml_parsed_arguments_string(args),
                }
        except (json.JSONDecodeError, ValueError):
            salvaged = _salvage_json_tool_call(block)
            if salvaged:
                name = salvaged["name"]
                try:
                    a = _json_loads_lenient(salvaged["arguments"])
                    if isinstance(a, dict):
                        a = _repair_artifact_path_values_in_tool_args_dict(
                            str(name or "").strip(), a, block, recovery_context=recovery_context
                        )
                        salvaged = {"name": name, "arguments": json.dumps(a)}
                except Exception:
                    pass
                return {
                    "id": f"xmltc_{idx}_{hashlib.sha256(name.encode()).hexdigest()[:8]}",
                    "name": name,
                    "arguments": _finalize_xml_parsed_arguments_string(salvaged["arguments"]),
                }
        return None

    # --- Format 1: <function=name><parameter=key>val</parameter></function> ---
    fn_eq = re.search(r"<function=([^>]+)>", block)
    if fn_eq:
        fn_name = fn_eq.group(1).strip()
        params = _extract_xml_params_robust(
            block,
            r"<parameter=([^>]+)>",
            lambda m: m.group(1).strip(),
        )
        params = _repair_artifact_path_values_in_params(
            fn_name, params, block, recovery_context=recovery_context
        )
        return {
            "id": f"xmltc_{idx}_{hashlib.sha256(fn_name.encode()).hexdigest()[:8]}",
            "name": fn_name,
            "arguments": _finalize_xml_parsed_arguments_string(json.dumps(params)),
        }

    # --- Format 2: <invoke name="fn"><parameter name="key">val</parameter></invoke> ---
    inv = re.search(r'<invoke\s+name=["\']([^"\']+)["\']', block)
    if inv:
        fn_name = inv.group(1).strip()
        params = _extract_xml_params_robust(
            block,
            r'<parameter\s+name=["\']([^"\']+)["\']>',
            lambda m: m.group(1).strip(),
        )
        params = _repair_artifact_path_values_in_params(
            fn_name, params, block, recovery_context=recovery_context
        )
        return {
            "id": f"xmltc_{idx}_{hashlib.sha256(fn_name.encode()).hexdigest()[:8]}",
            "name": fn_name,
            "arguments": _finalize_xml_parsed_arguments_string(json.dumps(params)),
        }

    # --- Format 4: bare tool_name<arg_key>key</arg_key><arg_value>val</arg_value> ---
    _argkey_re = re.compile(r"<arg_key>\s*(.*?)\s*</arg_key>\s*<arg_value>\s*([\s\S]*?)\s*</arg_value>", re.IGNORECASE)
    _argkey_matches = list(_argkey_re.finditer(block))

    # Also detect an unclosed trailing <arg_key>key</arg_key><arg_value>...EOF
    # (model ran out of output tokens before emitting </arg_value>).
    _unclosed_key_re = re.compile(r"<arg_key>\s*(.*?)\s*</arg_key>\s*<arg_value>\s*", re.IGNORECASE)
    _unclosed_trailing: tuple[str, str] | None = None
    _search_start = _argkey_matches[-1].end() if _argkey_matches else 0
    _tail = block[_search_start:]
    _um = _unclosed_key_re.search(_tail)
    if _um:
        _ukey = _um.group(1).strip()
        _uval_start = _um.end()
        if _ukey and _tail.find("</arg_value>", _uval_start) == -1:
            _uval = _tail[_uval_start:].strip()
            _uval = re.sub(r"</?(tool_call|arg_key|arg_value)>", "", _uval, flags=re.IGNORECASE).strip()
            if _uval:
                _unclosed_trailing = (_ukey, _uval)

    if _argkey_matches or _unclosed_trailing:
        _anchor = _argkey_matches[0].start() if _argkey_matches else (_search_start + _um.start())  # type: ignore[union-attr]
        fn_name_text = block[: _anchor].strip()
        fn_name_text = re.sub(r"</?[a-z_]+>", "", fn_name_text, flags=re.IGNORECASE).strip()
        if fn_name_text:
            params: dict[str, str] = {}
            for m in _argkey_matches:
                key = m.group(1).strip()
                val = m.group(2).strip()
                if key:
                    # Strip stray tag fragments that leak into values.
                    val = re.sub(r"</?arg_(key|value)>", "", val, flags=re.IGNORECASE).strip()
                    params[key] = val
            if _unclosed_trailing:
                params[_unclosed_trailing[0]] = _unclosed_trailing[1]
            params = _repair_artifact_path_values_in_params(
                fn_name_text.strip(), params, block, recovery_context=recovery_context
            )
            return {
                "id": f"xmltc_{idx}_{hashlib.sha256(fn_name_text.encode()).hexdigest()[:8]}",
                "name": fn_name_text,
                "arguments": _finalize_xml_parsed_arguments_string(json.dumps(params)),
            }

    # --- Format 5: Gemma channel-style call:name{args} / call:name(args) ---
    _gemma_call_re = re.search(
        r"call:\s*(\w+)\s*([\({][\s\S]*[\}\)]|\([\s\S]*\))\s*$",
        block,
        re.IGNORECASE,
    )
    if _gemma_call_re:
        fn_name = _gemma_call_re.group(1).strip()
        args_str = _parse_gemma_delimited_args(_gemma_call_re.group(2))
        if fn_name:
            return {
                "id": f"xmltc_{idx}_{hashlib.sha256(fn_name.encode()).hexdigest()[:8]}",
                "name": fn_name,
                "arguments": _finalize_xml_parsed_arguments_string(args_str),
            }

    return None


def _extract_unclosed_tool_call_block(text: str) -> str | None:
    """Return payload for a trailing unclosed <tool_call> or <|tool_call> block, if present."""
    if not isinstance(text, str) or not text:
        return None
    lower = text.lower()
    best_idx, best_tag_len = -1, 0
    for open_tag in _TOOL_CALL_OPEN_TAGS:
        idx = lower.rfind(open_tag.lower())
        if idx != -1 and (best_idx == -1 or idx > best_idx):
            best_idx = idx
            best_tag_len = len(open_tag)
    if best_idx == -1:
        return None
    after = text[best_idx + best_tag_len:]
    for close_tag in _TOOL_CALL_CLOSE_TAGS:
        if close_tag.lower() in after.lower():
            return None
    payload = after.strip()
    return payload or None
def _parse_xml_tool_calls(text: str) -> list[dict]:
    """Extract tool calls from XML markup in model output.

    Supported formats:
      1. <tool_call><function=name><parameter=key>val</parameter>...</function></tool_call>
      2. <tool_call><invoke name="fn"><parameter name="key">val</parameter>...</invoke></tool_call>
      3. <tool_call>{"name":"fn","arguments":{...}}</tool_call>  (Hermes JSON)
      4. <tool_call>tool_name<arg_key>k</arg_key><arg_value>v</arg_value></tool_call>
      5. <|tool_call>call:tool_name{args}<tool_call|>  (Gemma channel-style)

    Returns list of {"id": str, "name": str, "arguments": str(json)}.
    Parsing is bounded and linear over explicit <tool_call> blocks only.
    """
    raw_matches = _TOOL_CALL_BLOCK_RE.findall(text)
    if not raw_matches:
        return []

    blocks = [g1 or g2 for g1, g2 in raw_matches]

    results = []
    for idx, block in enumerate(blocks):
        parsed = _parse_xml_tool_call_block(block, idx, recovery_context=text)
        if parsed:
            results.append(parsed)

    return results


def _native_xml_tool_call_key(tc: dict[str, Any]) -> str | None:
    if not isinstance(tc, dict):
        return None
    name = (tc.get("name") or "").strip()
    if not name:
        return None
    try:
        parsed = _json_loads_lenient(tc.get("arguments") or "{}")
    except Exception:
        parsed = {}
    if not isinstance(parsed, dict):
        parsed = {}
    payload = json.dumps(parsed, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
    return f"{name}|{payload}"


def narrow_merge_native_xml_tool_calls(
    native: list[dict[str, Any]],
    xml: list[dict[str, Any]],
) -> list[dict[str, Any]] | None:
    """If exactly one native and one XML call with different tool names and no duplicate
    (name,args) key vs native, return [native, xml] (max two calls). Otherwise None.

    Prefer :func:`union_merge_native_xml_tool_calls` for multi-call native output; this
    remains for tests and explicit 1+1 semantics.
    """
    if len(native) != 1 or len(xml) != 1:
        return None
    n_tc, x_tc = native[0], xml[0]
    if not isinstance(n_tc, dict) or not isinstance(x_tc, dict):
        return None
    n_name = (n_tc.get("name") or "").strip()
    x_name = (x_tc.get("name") or "").strip()
    if not n_name or not x_name or n_name == x_name:
        return None
    nk = _native_xml_tool_call_key(n_tc)
    xk = _native_xml_tool_call_key(x_tc)
    if not nk or not xk or nk == xk:
        return None
    out_xml = dict(x_tc)
    if not out_xml.get("id"):
        out_xml["id"] = f"xml_merged_{n_name}_{x_name}"
    return [dict(n_tc), out_xml]


# Max combined native + salvaged XML tool calls in one turn (after union dedupe).
_MAX_UNION_NATIVE_XML_TOOL_CALLS = 32
def union_merge_native_xml_tool_calls(
    native: list[dict[str, Any]],
    xml: list[dict[str, Any]],
    *,
    max_total: int = _MAX_UNION_NATIVE_XML_TOOL_CALLS,
) -> tuple[list[dict[str, Any]], int]:
    """Append XML-derived calls whose ``(name, canonical-args)`` key is new.

    Native order is preserved. Duplicate keys (same tool + same parsed arguments)
    are skipped. Returns ``(merged_list, xml_added_count)``.
    """
    if not native:
        return (list(xml) if xml else [], 0)
    if not xml:
        return ([dict(tc) if isinstance(tc, dict) else tc for tc in native], 0)

    out: list[dict[str, Any]] = []
    seen: set[str] = set()
    for tc in native:
        if not isinstance(tc, dict):
            continue
        k = _native_xml_tool_call_key(tc)
        if k:
            seen.add(k)
        out.append(dict(tc))

    xml_added = 0
    for tc in xml:
        if not isinstance(tc, dict):
            continue
        if len(out) >= max_total:
            break
        k = _native_xml_tool_call_key(tc)
        if not k:
            continue
        if k in seen:
            continue
        seen.add(k)
        xcopy = dict(tc)
        if not xcopy.get("id"):
            xcopy["id"] = (
                f"xml_union_{xml_added}_"
                f"{hashlib.sha256(k.encode('utf-8', errors='replace')).hexdigest()[:12]}"
            )
        out.append(xcopy)
        xml_added += 1
        if len(out) >= max_total:
            break

    return (out, xml_added)


def _xml_calls_delegate_task_only(xml: list[dict[str, Any]]) -> bool:
    if not xml:
        return False
    for tc in xml:
        if not isinstance(tc, dict):
            return False
        if (tc.get("name") or "").strip() != "delegate_task":
            return False
    return True


def _format_delegate_task_native_repair_prompt() -> str:
    return (
        "You described `delegate_task` in assistant prose but only native `tool_calls` were executed this turn. "
        "In your next message, issue `delegate_task` as a native structured tool call with valid JSON arguments. "
        "Do not duplicate the action in markup or narration.\n"
        + GUARD_TOOL_CALL_SUFFIX
    )
_PROMISED_DELEGATE_HINT_RE = re.compile(
    r"(?:\bdelegate_task\b|"
    r"\b(?:local\s+)?subagent\b|"
    r"\brunning\s+in\s+the\s+background\b|"
    r"\b(?:spin(?:ning)?|start(?:ing)?|launch(?:ing)?|spawn(?:ing)?)\s+(?:up\s+)?(?:a\s+)?(?:local\s+)?subagent\b|"
    r"\bbackground\s+(?:subagent|task|worker)\b)",
    re.IGNORECASE,
)


def _delegate_task_in_assistant_message(m: dict[str, Any]) -> bool:
    if not isinstance(m, dict) or m.get("role") != "assistant":
        return False
    tcs = m.get("tool_calls")
    if not isinstance(tcs, list):
        return False
    for tc in tcs:
        if not isinstance(tc, dict):
            continue
        fn = tc.get("function")
        if isinstance(fn, dict):
            nm = (fn.get("name") or "").strip()
        else:
            nm = (tc.get("name") or "").strip()
        if nm == "delegate_task":
            return True
    return False


def _delegate_task_called_since_last_user(messages: list[dict]) -> bool:
    if not isinstance(messages, list) or not messages:
        return False
    last_user_i = -1
    for i in range(len(messages) - 1, -1, -1):
        m = messages[i]
        if isinstance(m, dict) and (m.get("role") or "").strip() == "user":
            last_user_i = i
            break
    tail = messages[last_user_i + 1 :] if last_user_i >= 0 else messages
    for m in tail:
        if _delegate_task_in_assistant_message(m):
            return True
    return False


def _last_tool_message_content(messages: list[dict]) -> str:
    for m in reversed(messages[-32:]):
        if not isinstance(m, dict) or m.get("role") != "tool":
            continue
        c = m.get("content")
        if isinstance(c, str) and c.strip():
            return c
    return ""


def _last_tool_suggests_mesh_handshake_wait(messages: list[dict]) -> bool:
    c = _last_tool_message_content(messages).lower()
    if not c:
        return False
    return "ready for handoff" in c or "formal handoff" in c


def _promised_delegate_without_tool(
    content: str,
    messages: list[dict],
    *,
    coordinator_mode: bool = False,
) -> bool:
    if coordinator_mode:
        return False
    if not isinstance(content, str):
        return False
    text = content.strip()
    if len(text) < 40:
        return False
    if not _PROMISED_DELEGATE_HINT_RE.search(text):
        return False
    if _delegate_task_called_since_last_user(messages):
        return False
    if re.search(r"\bsubagent\s+mode\b", text, re.IGNORECASE):
        return False
    if _last_tool_suggests_mesh_handshake_wait(messages):
        return True
    return True


def _format_promised_delegate_retry_prompt() -> str:
    return (
        "Your last reply described delegating work (subagent/background task) but `delegate_task` did not run. "
        "Call `delegate_task` now with the concrete task, or answer directly without claiming delegation is in flight.\n"
        + GUARD_TOOL_CALL_SUFFIX
    )

def _strip_xml_tool_calls(text: str) -> str:
    """Remove <tool_call>...</tool_call> and <|tool_call>...<tool_call|> blocks from text.

    Also strips half-open / unclosed tool-call markup that balanced-regex removal leaves behind.
    """
    t = _TOOL_CALL_BLOCK_RE.sub("", text).strip()

    def _first_close_end(fragment: str) -> int | None:
        best: int | None = None
        low = fragment.lower()
        for ct in _TOOL_CALL_CLOSE_TAGS:
            idx = low.find(ct.lower())
            if idx != -1:
                end = idx + len(ct)
                best = end if best is None else min(best, end)
        return best

    for _ in range(32):
        m = _TOOL_CALL_TAG_RE.search(t)
        if not m:
            break
        rest = t[m.start() :]
        close_end = _first_close_end(rest)
        if close_end is None:
            t = t[: m.start()].strip()
            continue
        t = (t[: m.start()] + rest[close_end:]).strip()

    t = re.sub(r"(?:</tool_call>\s*|<tool_call\|>\s*)+$", "", t, flags=re.IGNORECASE).strip()
    return t
