"""Patch preciz SDK's _parse_tool_calls to handle malformed XML gracefully."""

import site
import re

P = f"{site.getsitepackages()[0]}/summarizer/llm_client.py"

with open(P) as f:
    src = f.read()

marker_old = (
    '                raw_text = re.sub(r"\\s+handle$", "", raw_text).strip()\n'
    '                return [(fallback.group(1), raw_text)]\n'
    '            return []'
)

marker_new = (
    '                raw_text = re.sub(r"\\s+handle$", "", raw_text).strip()\n'
    '                return [(fallback.group(1), raw_text)]\n'
    '            # Ultra-fallback: strip all tags, extract key-value pairs\n'
    '            # from even badly malformed XML\n'
    '            stripped = re.sub(r"<[^>]+>", "\\n", content)\n'
    '            lines = [l.strip() for l in stripped.split("\\n") if l.strip()]\n'
    '            known_keys = {"name", "question", "dataset"}\n'
    '            tool = "rag_qa"\n'
    '            question = ""\n'
    '            dataset = "developer-portfolio"\n'
    '            for i, line in enumerate(lines):\n'
    '                if line not in known_keys:\n'
    '                    continue\n'
    '                if i + 1 < len(lines) and lines[i + 1] not in known_keys:\n'
    '                    val = lines[i + 1]\n'
    '                    if line == "name":\n'
    '                        tool = LONGCAT_NAME_MAP.get(val, val)\n'
    '                    elif line == "question":\n'
    '                        question = val\n'
    '                    elif line == "dataset":\n'
    '                        dataset = val\n'
    '            if not question:\n'
    '                vals = [l for l in lines if l not in known_keys]\n'
    '                if vals:\n'
    '                    question = " ".join(vals)\n'
    '            if question:\n'
    '                return [_normalize_tool_args(tool, {"question": question, "dataset": dataset})]\n'
    '            return []'
)

assert marker_old in src, "Marker not found"
src = src.replace(marker_old, marker_new, 1)

with open(P, "w") as f:
    f.write(src)

print("Patched _parse_tool_calls with robust XML fallback")