"""Display helpers — make raw model output safe and readable in the Chatbot.

The Gradio Chatbot renders content as markdown and runs it through an HTML
sanitizer. Content that starts with a non-standard tag like `<tool_call>` is
treated as a raw HTML block and silently dropped, leaving an empty bubble. Some
models emit exactly this when they call a tool, e.g. to end a call:

    <tool_call>
    {"name": "end_call", "arguments": {"final_message": "..."}}
    </tool_call>

`format_for_display` converts such content into readable markdown (the
final_message as plain text + the tool call in a fenced JSON block) and escapes
any other stray leading tag so the bubble always renders. This is for DISPLAY
only — the raw content is what gets stored and fed back into each thread.
"""

import json
import re

_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)


def _safe_segment(text: str) -> str:
    """Wrap a text segment in a code block if it would trip the HTML sanitizer.

    Content that starts with a tag like `<tool_call` or `<think` (e.g. an
    unterminated tool call seen mid-stream) is treated as a raw HTML block by
    Gradio and dropped — wrapping it guarantees it renders.
    """
    if re.match(r"^<[a-zA-Z]", text.strip()):
        return "```\n" + text + "\n```"
    return text


def _render_tool_call(payload: dict) -> str:
    name = payload.get("name", "tool")
    args = payload.get("arguments", {})
    if isinstance(args, str):
        # arguments sometimes arrive as a JSON-encoded string.
        try:
            args = json.loads(args)
        except (json.JSONDecodeError, TypeError):
            args = {"raw": args}

    parts = []
    final_message = args.get("final_message") if isinstance(args, dict) else None
    if final_message:
        parts.append(str(final_message))
    pretty = json.dumps({"name": name, "arguments": args}, ensure_ascii=False, indent=2)
    parts.append(f"🛠️ **tool call: `{name}`**\n```json\n{pretty}\n```")
    return "\n\n".join(parts)


def format_for_display(raw) -> str:
    """Return markdown-safe, always-renderable text for a chatbot bubble."""
    if raw is None:
        return "_(empty response)_"
    text = str(raw).strip()
    if not text:
        return "_(empty response)_"

    # Replace any <tool_call>{...}</tool_call> blocks with readable markdown.
    if "<tool_call>" in text:
        rendered = []
        last = 0
        for m in _TOOL_CALL_RE.finditer(text):
            before = text[last:m.start()].strip()
            if before:
                rendered.append(_safe_segment(before))
            try:
                payload = json.loads(m.group(1))
                rendered.append(_render_tool_call(payload))
            except (json.JSONDecodeError, TypeError):
                # Couldn't parse — show the inner text in a code block, never raw tags.
                rendered.append(f"```\n{m.group(0)}\n```")
            last = m.end()
        tail = text[last:].strip()
        if tail:
            # The tail may be an unterminated "<tool_call>{..." (mid-stream) — wrap
            # any stray-tag content in a code block so it always renders.
            rendered.append(_safe_segment(tail))
        return "\n\n".join(rendered) if rendered else "_(empty response)_"

    # Guard against any other content that starts with a tag the sanitizer would
    # eat (e.g. "<think>"): if it leads with "<word", show it in a code block.
    return _safe_segment(text)