from __future__ import annotations

import json
from string.templatelib import Template

from markupsafe import Markup

from ..data.models import (
    ImageBlock,
    Message,
    TextBlock,
    ThinkingBlock,
    ToolResultBlock,
    ToolUseBlock,
)

TOOL_ICONS: dict[str, str] = {
    "Task": "\U0001f527",
    "Bash": "\u26a1",
    "Read": "\U0001f4c4",
    "Write": "\u270f\ufe0f",
    "Edit": "\u270f\ufe0f",
    "Glob": "\U0001f50d",
    "Grep": "\U0001f50d",
    "WebFetch": "\U0001f310",
    "WebSearch": "\U0001f310",
}

DEFAULT_TOOL_ICON = "\U0001f527"


def render_messages(messages: list[Message]) -> list[Template]:
    """Render a list of messages, pairing tool-use with tool-result blocks."""
    result: list[Template] = []
    skip_next = False

    for i, msg in enumerate(messages):
        if skip_next:
            skip_next = False
            continue

        if msg.role == "user":
            if all(isinstance(b, ToolResultBlock) for b in msg.content_blocks):
                # Tool-results-only message — already consumed by previous assistant
                continue
            result.append(_render_user_message(msg))
        else:
            # Look ahead for tool results in the next user message
            result_map: dict[str, ToolResultBlock] = {}
            if i + 1 < len(messages):
                next_msg = messages[i + 1]
                if next_msg.role == "user" and all(
                    isinstance(b, ToolResultBlock) for b in next_msg.content_blocks
                ):
                    for b in next_msg.content_blocks:
                        if isinstance(b, ToolResultBlock):
                            result_map[b.tool_use_id] = b
                    skip_next = True

            result.append(_render_assistant_message(msg, result_map))

    return result


def _render_user_message(msg: Message) -> Template:
    blocks = [_render_content_block(b) for b in msg.content_blocks]
    time = msg.time_display
    return t"""<div class="message user">
        <div class="role-label">user <span class="timestamp">{time}</span></div>
        {blocks}
    </div>"""


def _render_assistant_message(
    msg: Message, result_map: dict[str, ToolResultBlock]
) -> Template:
    parts: list[Template] = []
    time = msg.time_display

    # Collect tool_use blocks to group them
    tool_uses: list[ToolUseBlock] = []

    for block in msg.content_blocks:
        if isinstance(block, ToolUseBlock):
            tool_uses.append(block)
        else:
            # Flush any pending tool uses before rendering non-tool block
            if tool_uses:
                parts.append(_render_tool_groups(tool_uses, result_map))
                tool_uses = []
            parts.append(_render_content_block(block))

    # Flush remaining tool uses
    if tool_uses:
        parts.append(_render_tool_groups(tool_uses, result_map))

    return t"""<div class="message assistant">
        <div class="role-label">assistant <span class="timestamp">{time}</span></div>
        {parts}
    </div>"""


def _render_tool_groups(
    tool_uses: list[ToolUseBlock], result_map: dict[str, ToolResultBlock]
) -> Template:
    groups = [_render_tool_group(tu, result_map.get(tu.id)) for tu in tool_uses]

    if len(groups) == 1:
        return groups[0]

    return t"""<div class="tool-call-stack">{groups}</div>"""


def _render_tool_group(
    tool_use: ToolUseBlock, tool_result: ToolResultBlock | None
) -> Template:
    use_html = _render_tool_use(tool_use)
    result_html = _render_tool_result(tool_result) if tool_result else t""""""

    return t"""<div class="tool-group">
        {use_html}
        {result_html}
    </div>"""


def _render_content_block(block) -> Template:
    match block:
        case TextBlock():
            return _render_text(block)
        case ThinkingBlock():
            return _render_thinking(block)
        case ToolUseBlock():
            return _render_tool_use(block)
        case ToolResultBlock():
            return _render_tool_result_standalone(block)
        case ImageBlock():
            return _render_image(block)
        case _:
            return t"""<div class="unknown-block">Unknown block type</div>"""


def _render_text(block: TextBlock) -> Template:
    return t"""<div class="msg-text">{block.text}</div>"""


def _render_thinking(block: ThinkingBlock) -> Template:
    return t"""<details class="thinking-block">
        <summary class="thinking-toggle"><span class="arrow">▶</span> Thinking\u2026</summary>
        <div class="thinking-content">{block.thinking}</div>
    </details>"""


def _render_tool_use(block: ToolUseBlock) -> Template:
    icon = TOOL_ICONS.get(block.name, DEFAULT_TOOL_ICON)
    summary = _tool_summary(block)
    formatted_input = _format_json(block.input)
    return t"""<div class="tool-use-block">
        <div class="tool-header">
            <span class="tool-icon">{icon}</span>
            <span class="tool-name">{block.name}</span>
            <span class="tool-desc">{summary}</span>
        </div>
        <details class="tool-input-section">
            <summary class="tool-input-toggle"><span class="arrow">▶</span> Input</summary>
            <div class="tool-input-content"><pre>{formatted_input}</pre></div>
        </details>
    </div>"""


def _render_tool_result(block: ToolResultBlock | None) -> Template:
    if block is None:
        return t""""""
    result_text = block.result_text()
    error_class = "error" if block.is_error else ""
    icon = "\u2717" if block.is_error else "\u2713"
    return t"""<details class="tool-result-block {error_class}">
        <summary class="result-header">
            <span class="result-icon">{icon}</span>
            <span class="result-label">Result</span>
            <span class="result-arrow">▶</span>
        </summary>
        <div class="result-content">{result_text}</div>
    </details>"""


def _render_tool_result_standalone(block: ToolResultBlock) -> Template:
    """Render a tool result that wasn't paired with a tool use."""
    return _render_tool_result(block)


def _render_image(block: ImageBlock) -> Template:
    src = Markup(f"data:{block.media_type};base64,{block.data}")
    return t"""<div class="image-block"><img src="{src}" class="content-image" /></div>"""


def _tool_summary(block: ToolUseBlock) -> str:
    inp = block.input
    if block.name in ("Read", "Glob", "Grep"):
        return inp.get("file_path", "") or inp.get("pattern", "") or inp.get("path", "")
    if block.name == "Write":
        return inp.get("file_path", "")
    if block.name == "Edit":
        return inp.get("file_path", "")
    if block.name == "Bash":
        cmd = inp.get("command", "")
        return cmd[:80] if cmd else ""
    if block.name == "Task":
        return inp.get("description", "")
    # Fallback: show first string value
    for v in inp.values():
        if isinstance(v, str) and v:
            return v[:80]
    return ""


def _format_json(data: dict) -> str:
    try:
        return json.dumps(data, indent=2, ensure_ascii=False)
    except (TypeError, ValueError):
        return str(data)
