handler.py

"""
Claude Message Handler

Platform-agnostic Claude interaction logic.
Handles the core workflow of processing user messages via Claude CLI.
Uses tree-based queuing for message ordering.
"""

import time
import asyncio
import logging
from typing import List, Optional

from markdown_it import MarkdownIt

from .base import MessagingPlatform, SessionManagerInterface
from .models import IncomingMessage
from .session import SessionStore
from .tree_queue import TreeQueueManager, MessageNode, MessageState, MessageTree
from .event_parser import parse_cli_event

logger = logging.getLogger(__name__)


MDV2_SPECIAL_CHARS = set("\\_*[]()~`>#+-=|{}.!")

MDV2_LINK_ESCAPE = set("\\)")

_MD = MarkdownIt("commonmark", {"html": False, "breaks": False})
_MD.enable("strikethrough")


def escape_md_v2(text: str) -> str:
    """Escape text for Telegram MarkdownV2."""
    return "".join(f"\\{ch}" if ch in MDV2_SPECIAL_CHARS else ch for ch in text)


def escape_md_v2_code(text: str) -> str:
    """Escape text for Telegram MarkdownV2 code spans/blocks."""
    return text.replace("\\", "\\\\").replace("`", "\\`")


def escape_md_v2_link_url(text: str) -> str:
    """Escape URL for Telegram MarkdownV2 link destination."""
    return "".join(f"\\{ch}" if ch in MDV2_LINK_ESCAPE else ch for ch in text)


def mdv2_bold(text: str) -> str:
    return f"*{escape_md_v2(text)}*"


def mdv2_code_inline(text: str) -> str:
    return f"`{escape_md_v2_code(text)}`"


def format_status(emoji: str, label: str, suffix: Optional[str] = None) -> str:
    base = f"{emoji} {mdv2_bold(label)}"
    if suffix:
        return f"{base} {escape_md_v2(suffix)}"
    return base


def render_markdown_to_mdv2(text: str) -> str:
    """Render common Markdown into Telegram MarkdownV2."""
    if not text:
        return ""

    tokens = _MD.parse(text)

    def render_inline_plain(children) -> str:
        out: List[str] = []
        for tok in children:
            if tok.type == "text":
                out.append(escape_md_v2(tok.content))
            elif tok.type == "code_inline":
                out.append(escape_md_v2(tok.content))
            elif tok.type in {"softbreak", "hardbreak"}:
                out.append("\n")
        return "".join(out)

    def render_inline(children) -> str:
        out: List[str] = []
        i = 0
        while i < len(children):
            tok = children[i]
            t = tok.type
            if t == "text":
                out.append(escape_md_v2(tok.content))
            elif t in {"softbreak", "hardbreak"}:
                out.append("\n")
            elif t == "em_open":
                out.append("_")
            elif t == "em_close":
                out.append("_")
            elif t == "strong_open":
                out.append("*")
            elif t == "strong_close":
                out.append("*")
            elif t == "s_open":
                out.append("~")
            elif t == "s_close":
                out.append("~")
            elif t == "code_inline":
                out.append(f"`{escape_md_v2_code(tok.content)}`")
            elif t == "link_open":
                href = ""
                if tok.attrs:
                    for key, val in tok.attrs:
                        if key == "href":
                            href = val
                            break
                inner_tokens = []
                i += 1
                while i < len(children) and children[i].type != "link_close":
                    inner_tokens.append(children[i])
                    i += 1
                link_text = ""
                for child in inner_tokens:
                    if child.type == "text":
                        link_text += child.content
                    elif child.type == "code_inline":
                        link_text += child.content
                out.append(
                    f"[{escape_md_v2(link_text)}]({escape_md_v2_link_url(href)})"
                )
            elif t == "image":
                href = ""
                alt = tok.content or ""
                if tok.attrs:
                    for key, val in tok.attrs:
                        if key == "src":
                            href = val
                            break
                if alt:
                    out.append(f"{escape_md_v2(alt)} ({escape_md_v2_link_url(href)})")
                else:
                    out.append(escape_md_v2_link_url(href))
            else:
                out.append(escape_md_v2(tok.content or ""))
            i += 1
        return "".join(out)

    out: List[str] = []
    list_stack: List[dict] = []
    pending_prefix: Optional[str] = None
    blockquote_level = 0
    in_heading = False

    def apply_blockquote(val: str) -> str:
        if blockquote_level <= 0:
            return val
        prefix = "> " * blockquote_level
        return prefix + val.replace("\n", "\n" + prefix)

    i = 0
    while i < len(tokens):
        tok = tokens[i]
        t = tok.type
        if t == "paragraph_open":
            pass
        elif t == "paragraph_close":
            out.append("\n")
        elif t == "heading_open":
            in_heading = True
        elif t == "heading_close":
            in_heading = False
            out.append("\n")
        elif t == "bullet_list_open":
            list_stack.append({"type": "bullet", "index": 1})
        elif t == "bullet_list_close":
            if list_stack:
                list_stack.pop()
            out.append("\n")
        elif t == "ordered_list_open":
            start = 1
            if tok.attrs:
                for key, val in tok.attrs:
                    if key == "start":
                        try:
                            start = int(val)
                        except ValueError:
                            start = 1
                        break
            list_stack.append({"type": "ordered", "index": start})
        elif t == "ordered_list_close":
            if list_stack:
                list_stack.pop()
            out.append("\n")
        elif t == "list_item_open":
            if list_stack:
                top = list_stack[-1]
                if top["type"] == "bullet":
                    pending_prefix = "\\- "
                else:
                    pending_prefix = f"{top['index']}\\."
                    top["index"] += 1
                    pending_prefix += " "
        elif t == "list_item_close":
            out.append("\n")
        elif t == "blockquote_open":
            blockquote_level += 1
        elif t == "blockquote_close":
            blockquote_level = max(0, blockquote_level - 1)
            out.append("\n")
        elif t in {"code_block", "fence"}:
            code = escape_md_v2_code(tok.content.rstrip("\n"))
            out.append(f"```\n{code}\n```")
            out.append("\n")
        elif t == "inline":
            rendered = render_inline(tok.children or [])
            if in_heading:
                rendered = f"*{render_inline_plain(tok.children or [])}*"
            if pending_prefix:
                rendered = pending_prefix + rendered
                pending_prefix = None
            rendered = apply_blockquote(rendered)
            out.append(rendered)
        else:
            if tok.content:
                out.append(escape_md_v2(tok.content))
        i += 1

    return "".join(out).rstrip()


class ClaudeMessageHandler:
    """
    Platform-agnostic handler for Claude interactions.

    Uses a tree-based message queue where:
    - New messages create a tree root
    - Replies become children of the message being replied to
    - Each node has state: PENDING, IN_PROGRESS, COMPLETED, ERROR
    - Per-tree queue ensures ordered processing
    """

    def __init__(
        self,
        platform: MessagingPlatform,
        cli_manager: SessionManagerInterface,
        session_store: SessionStore,
    ):
        self.platform = platform
        self.cli_manager = cli_manager
        self.session_store = session_store
        self.tree_queue = TreeQueueManager(
            queue_update_callback=self._update_queue_positions,
            node_started_callback=self._mark_node_processing,
        )

    async def handle_message(self, incoming: IncomingMessage) -> None:
        """
        Main entry point for handling an incoming message.

        Determines if this is a new conversation or reply,
        creates/extends the message tree, and queues for processing.
        """
        # Check for commands
        if incoming.text == "/stop":
            await self._handle_stop_command(incoming)
            return

        if incoming.text == "/stats":
            await self._handle_stats_command(incoming)
            return

        # Filter out status messages (our own messages)
        if any(
            incoming.text.startswith(p)
            for p in ["⏳", "💭", "🔧", "✅", "❌", "🚀", "🤖", "📋", "📊", "🔄"]
        ):
            return

        # Check if this is a reply to an existing node in a tree
        parent_node_id = None
        tree = None

        if incoming.is_reply() and incoming.reply_to_message_id:
            # Look up if the replied-to message is in any tree (could be a node or status message)
            reply_id = incoming.reply_to_message_id
            tree = self.tree_queue.get_tree_for_node(reply_id)
            if tree:
                # Resolve to actual node ID (handles status message replies)
                parent_node_id = self.tree_queue.resolve_parent_node_id(reply_id)
                if parent_node_id:
                    logger.info(f"Found tree for reply, parent node: {parent_node_id}")
                else:
                    logger.warning(
                        f"Reply to {incoming.reply_to_message_id} found tree but no valid parent node"
                    )
                    tree = None  # Treat as new conversation

        # Generate node ID
        node_id = incoming.message_id

        # Send initial status message
        status_text = self._get_initial_status(tree, parent_node_id)
        status_msg_id = await self.platform.queue_send_message(
            incoming.chat_id,
            status_text,
            reply_to=incoming.message_id,
            fire_and_forget=False,
        )

        # Create or extend tree
        if parent_node_id and tree and status_msg_id:
            # Reply to existing node - add as child
            tree, node = await self.tree_queue.add_to_tree(
                parent_node_id=parent_node_id,
                node_id=node_id,
                incoming=incoming,
                status_message_id=status_msg_id,
            )
            # Register status message as a node too for reply chains
            self.tree_queue.register_node(status_msg_id, tree.root_id)
            self.session_store.register_node(status_msg_id, tree.root_id)
            self.session_store.register_node(node_id, tree.root_id)
        elif status_msg_id:
            # New conversation - create new tree
            tree = await self.tree_queue.create_tree(
                node_id=node_id,
                incoming=incoming,
                status_message_id=status_msg_id,
            )
            # Register status message
            self.tree_queue.register_node(status_msg_id, tree.root_id)
            self.session_store.register_node(node_id, tree.root_id)
            self.session_store.register_node(status_msg_id, tree.root_id)

        # Persist tree
        if tree:
            self.session_store.save_tree(tree.root_id, tree.to_dict())

        # Enqueue for processing
        was_queued = await self.tree_queue.enqueue(
            node_id=node_id,
            processor=self._process_node,
        )

        if was_queued and status_msg_id:
            # Update status to show queue position
            queue_size = self.tree_queue.get_queue_size(node_id)
            await self.platform.queue_edit_message(
                incoming.chat_id,
                status_msg_id,
                format_status("📋", "Queued", f"(position {queue_size}) - waiting..."),
                parse_mode="MarkdownV2",
            )

    async def _update_queue_positions(self, tree: MessageTree) -> None:
        """Refresh queued status messages after a dequeue."""
        try:
            queued_ids = await tree.get_queue_snapshot()
        except Exception as e:
            logger.warning(f"Failed to read queue snapshot: {e}")
            return

        if not queued_ids:
            return

        position = 0
        for node_id in queued_ids:
            node = tree.get_node(node_id)
            if not node or node.state != MessageState.PENDING:
                continue
            position += 1
            self.platform.fire_and_forget(
                self.platform.queue_edit_message(
                    node.incoming.chat_id,
                    node.status_message_id,
                    format_status(
                        "📋", "Queued", f"(position {position}) - waiting..."
                    ),
                    parse_mode="MarkdownV2",
                )
            )

    async def _mark_node_processing(self, tree: MessageTree, node_id: str) -> None:
        """Update the dequeued node's status to processing immediately."""
        node = tree.get_node(node_id)
        if not node or node.state == MessageState.ERROR:
            return
        self.platform.fire_and_forget(
            self.platform.queue_edit_message(
                node.incoming.chat_id,
                node.status_message_id,
                format_status("🔄", "Processing..."),
                parse_mode="MarkdownV2",
            )
        )

    async def _process_node(
        self,
        node_id: str,
        node: MessageNode,
    ) -> None:
        """Core task processor - handles a single Claude CLI interaction."""
        incoming = node.incoming
        status_msg_id = node.status_message_id
        chat_id = incoming.chat_id

        # Update node state to IN_PROGRESS
        tree = self.tree_queue.get_tree_for_node(node_id)
        if tree:
            await tree.update_state(node_id, MessageState.IN_PROGRESS)

        # Components for structured display
        components = {
            "thinking": [],
            "tools": [],
            "subagents": [],
            "content": [],
            "errors": [],
        }

        last_ui_update = 0.0
        last_displayed_text = None
        captured_session_id = None
        temp_session_id = None

        # Get parent session ID for forking (if child node)
        parent_session_id = None
        if tree and node.parent_id:
            parent_session_id = tree.get_parent_session_id(node_id)
            if parent_session_id:
                logger.info(f"Will fork from parent session: {parent_session_id}")

        async def update_ui(status: Optional[str] = None, force: bool = False) -> None:
            nonlocal last_ui_update, last_displayed_text
            now = time.time()

            # Small 1s debounce for UI sanity - we still want to avoid
            # spamming the queue with too many intermediate states
            if not force and now - last_ui_update < 1.0:
                return

            last_ui_update = now
            display = self._build_message(components, status)
            if display and display != last_displayed_text:
                last_displayed_text = display
                await self.platform.queue_edit_message(
                    chat_id, status_msg_id, display, parse_mode="MarkdownV2"
                )

        try:
            # Get or create CLI session
            try:
                (
                    cli_session,
                    session_or_temp_id,
                    is_new,
                ) = await self.cli_manager.get_or_create_session(
                    session_id=parent_session_id  # Fork from parent if available
                )
                if is_new:
                    temp_session_id = session_or_temp_id
                else:
                    captured_session_id = session_or_temp_id
            except RuntimeError as e:
                components["errors"].append(str(e))
                await update_ui(
                    format_status("⏳", "Session limit reached"), force=True
                )
                if tree:
                    await tree.update_state(
                        node_id, MessageState.ERROR, error_message=str(e)
                    )
                return

            # Process CLI events
            logger.info(f"HANDLER: Starting CLI task processing for node {node_id}")
            event_count = 0
            async for event_data in cli_session.start_task(
                incoming.text, session_id=captured_session_id
            ):
                if not isinstance(event_data, dict):
                    logger.warning(
                        f"HANDLER: Non-dict event received: {type(event_data)}"
                    )
                    continue
                event_count += 1
                if event_count % 10 == 0:
                    logger.debug(f"HANDLER: Processed {event_count} events so far")

                # Handle session_info event
                if event_data.get("type") == "session_info":
                    real_session_id = event_data.get("session_id")
                    if real_session_id and temp_session_id:
                        await self.cli_manager.register_real_session_id(
                            temp_session_id, real_session_id
                        )
                        captured_session_id = real_session_id
                        temp_session_id = None
                    continue

                parsed_list = parse_cli_event(event_data)
                logger.debug(f"HANDLER: Parsed {len(parsed_list)} events from CLI")

                for parsed in parsed_list:
                    if parsed["type"] == "thinking":
                        components["thinking"].append(parsed["text"])
                        await update_ui(format_status("🧠", "Claude is thinking..."))

                    elif parsed["type"] == "content":
                        if parsed.get("text"):
                            components["content"].append(parsed["text"])
                        await update_ui(format_status("🧠", "Claude is working..."))

                    elif parsed["type"] == "tool_start":
                        names = [t.get("name") for t in parsed.get("tools", [])]
                        components["tools"].extend(names)
                        await update_ui(format_status("⏳", "Executing tools..."))

                    elif parsed["type"] == "subagent_start":
                        tasks = parsed.get("tasks", [])
                        components["subagents"].extend(tasks)
                        await update_ui(format_status("🤖", "Subagent working..."))

                    elif parsed["type"] == "complete":
                        if not any(components.values()):
                            components["content"].append("Done.")
                        logger.info("HANDLER: Task complete, updating UI")
                        await update_ui(format_status("✅", "Complete"), force=True)

                        # Update node state and session
                        if tree and captured_session_id:
                            await tree.update_state(
                                node_id,
                                MessageState.COMPLETED,
                                session_id=captured_session_id,
                            )
                            self.session_store.save_tree(tree.root_id, tree.to_dict())

                    elif parsed["type"] == "error":
                        error_msg = parsed.get("message", "Unknown error")
                        logger.error(
                            f"HANDLER: Error event received: {error_msg[:200]}"
                        )
                        components["errors"].append(error_msg)
                        logger.info("HANDLER: Updating UI with error status")
                        await update_ui(format_status("❌", "Error"), force=True)
                        if tree:
                            await self._propagate_error_to_children(
                                node_id, error_msg, "Parent task failed"
                            )

        except asyncio.CancelledError:
            logger.warning(f"HANDLER: Task cancelled for node {node_id}")
            components["errors"].append("Task was cancelled")
            await update_ui(format_status("❌", "Cancelled"), force=True)
            if tree:
                await self._propagate_error_to_children(
                    node_id, "Cancelled by user", "Parent task was stopped"
                )
        except Exception as e:
            logger.error(
                f"HANDLER: Task failed with exception: {type(e).__name__}: {e}"
            )
            error_msg = str(e)[:200]
            components["errors"].append(error_msg)
            await update_ui(format_status("💥", "Task Failed"), force=True)
            if tree:
                await self._propagate_error_to_children(
                    node_id, error_msg, "Parent task failed"
                )
        finally:
            logger.info(
                f"HANDLER: _process_node completed for node {node_id}, errors={len(components['errors'])}"
            )

    async def _propagate_error_to_children(
        self,
        node_id: str,
        error_msg: str,
        child_status_text: str,
    ) -> None:
        """Mark node as error and propagate to pending children with UI updates."""
        affected = await self.tree_queue.mark_node_error(
            node_id, error_msg, propagate_to_children=True
        )
        # Update status messages for all affected children (skip first = current node)
        for child in affected[1:]:
            self.platform.fire_and_forget(
                self.platform.queue_edit_message(
                    child.incoming.chat_id,
                    child.status_message_id,
                    format_status("❌", "Cancelled:", child_status_text),
                    parse_mode="MarkdownV2",
                )
            )

    def _build_message(
        self,
        components: dict,
        status: Optional[str] = None,
    ) -> str:
        """
        Build unified message with specific order.
        Handles truncation while preserving markdown structure (closing code blocks).
        """
        lines = []

        # 1. Thinking
        if components["thinking"]:
            thinking_text = "".join(components["thinking"])
            # Truncate thinking if too long, it's usually less critical than final content
            if len(thinking_text) > 1000:
                thinking_text = "..." + thinking_text[-995:]

            lines.append(
                f"💭 {mdv2_bold('Thinking:')}\n```\n{escape_md_v2_code(thinking_text)}\n```"
            )

        # 2. Tools
        if components["tools"]:
            unique_tools = []
            seen = set()
            for t in components["tools"]:
                if t and t not in seen:
                    unique_tools.append(str(t))
                    seen.add(t)
            if unique_tools:
                lines.append(
                    f"🛠 {mdv2_bold('Tools:')} {mdv2_code_inline(', '.join(unique_tools))}"
                )

        # 3. Subagents
        if components["subagents"]:
            for task in components["subagents"]:
                lines.append(f"🤖 {mdv2_bold('Subagent:')} {mdv2_code_inline(task)}")

        # 4. Content
        if components["content"]:
            lines.append(render_markdown_to_mdv2("".join(components["content"])))

        # 5. Errors
        if components["errors"]:
            for err in components["errors"]:
                lines.append(f"⚠️ {mdv2_bold('Error:')} {mdv2_code_inline(err)}")

        if not any(lines) and not status:
            return format_status("⏳", "Claude is working...")

        # Telegram character limit is 4096. We leave buffer for status updates.
        LIMIT = 3900

        # Filter out empty lines first for a clean join
        lines = [l for l in lines if l]

        main_text = "\n".join(lines)
        status_text = f"\n\n{status}" if status else ""

        if len(main_text) + len(status_text) <= LIMIT:
            return (
                main_text + status_text
                if main_text + status_text
                else format_status("⏳", "Claude is working...")
            )

        # If too long, truncate the start of the content (keep the end)
        available_limit = LIMIT - len(status_text) - 20  # 20 for truncation marker
        raw_truncated = main_text[-available_limit:].lstrip()

        # Check for unbalanced code blocks
        prefix = escape_md_v2("... (truncated)\n")
        if raw_truncated.count("```") % 2 != 0:
            prefix += "```\n"

        truncated_main = prefix + raw_truncated

        return truncated_main + status_text

    def _get_initial_status(
        self,
        tree: Optional[object],
        parent_node_id: Optional[str],
    ) -> str:
        """Get initial status message text."""
        if tree and parent_node_id:
            # Reply to existing tree
            if self.tree_queue.is_node_tree_busy(parent_node_id):
                queue_size = self.tree_queue.get_queue_size(parent_node_id) + 1
                return format_status(
                    "📋", "Queued", f"(position {queue_size}) - waiting..."
                )
            return format_status("🔄", "Continuing conversation...")

        # New conversation
        stats = self.cli_manager.get_stats()
        if stats["active_sessions"] >= stats["max_sessions"]:
            return format_status(
                "⏳",
                "Waiting for slot...",
                f"({stats['active_sessions']}/{stats['max_sessions']})",
            )
        return format_status("⏳", "Launching new Claude CLI instance...")

    async def stop_all_tasks(self) -> int:
        """
        Stop all pending and in-progress tasks.

        Order of operations:
        1. Cancel tree queue tasks (uses internal locking)
        2. Stop CLI sessions
        3. Update UI for all affected nodes
        """
        # 1. Cancel tree queue tasks using the public async method
        logger.info("Cancelling tree queue tasks...")
        cancelled_nodes = await self.tree_queue.cancel_all()
        logger.info(f"Cancelled {len(cancelled_nodes)} nodes")

        # 2. Stop CLI sessions - this kills subprocesses and ensures everything is dead
        logger.info("Stopping all CLI sessions...")
        await self.cli_manager.stop_all()

        # 3. Update UI and persist state for all cancelled nodes
        for node in cancelled_nodes:
            self.platform.fire_and_forget(
                self.platform.queue_edit_message(
                    node.incoming.chat_id,
                    node.status_message_id,
                    format_status("⏹", "Stopped."),
                    parse_mode="MarkdownV2",
                )
            )

            # Persist tree state
            tree = self.tree_queue.get_tree_for_node(node.node_id)
            if tree:
                self.session_store.save_tree(tree.root_id, tree.to_dict())

        return len(cancelled_nodes)

    async def _handle_stop_command(self, incoming: IncomingMessage) -> None:
        """Handle /stop command from messaging platform."""
        count = await self.stop_all_tasks()
        await self.platform.queue_send_message(
            incoming.chat_id,
            format_status(
                "⏹", "Stopped.", f"Cancelled {count} pending or active requests."
            ),
        )

    async def _handle_stats_command(self, incoming: IncomingMessage) -> None:
        """Handle /stats command."""
        stats = self.cli_manager.get_stats()
        tree_count = self.tree_queue.get_tree_count()
        await self.platform.queue_send_message(
            incoming.chat_id,
            "📊 "
            + mdv2_bold("Stats")
            + "\n"
            + escape_md_v2(f"• Active CLI: {stats['active_sessions']}")
            + "\n"
            + escape_md_v2(f"• Max CLI: {stats['max_sessions']}")
            + "\n"
            + escape_md_v2(f"• Message Trees: {tree_count}"),
        )