pentestagent/pentestagent/interface/tui.py

"""
PentestAgent TUI - Terminal User Interface
"""

import asyncio
import re
import textwrap
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast

from rich.text import Text
from textual import on, work
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import (
    Center,
    Container,
    Horizontal,
    ScrollableContainer,
    Vertical,
)
from textual.reactive import reactive
from textual.screen import ModalScreen
from textual.scrollbar import ScrollBar, ScrollBarRender
from textual.timer import Timer
from textual.widgets import Button, Input, Static, Tree
from textual.widgets.tree import TreeNode

from ..config.constants import DEFAULT_MODEL

# ANSI escape sequence pattern for stripping control codes from input
_ANSI_ESCAPE = re.compile(
    r"\x1b\[[0-9;]*[mGKHflSTABCDEFsu]|\x1b\].*?\x07|\x1b\[<[0-9;]*[Mm]"
)


# ASCII-safe scrollbar renderer to avoid Unicode glyph issues
class ASCIIScrollBarRender(ScrollBarRender):
    """Scrollbar renderer using ASCII-safe characters."""

    BLANK_GLYPH = " "
    VERTICAL_BARS = [" ", " ", " ", " ", " ", " ", " ", " "]
    HORIZONTAL_BARS = [" ", " ", " ", " ", " ", " ", " ", " "]


# Apply ASCII scrollbar globally
ScrollBar.renderer = ASCIIScrollBarRender


# Custom Tree with ASCII-safe icons for PowerShell compatibility
class CrewTree(Tree):
    """Tree widget with ASCII-compatible expand/collapse icons."""

    ICON_NODE = "> "
    ICON_NODE_EXPANDED = "v "


if TYPE_CHECKING:
    from ..agents.pa_agent import PentestAgentAgent


def wrap_text_lines(text: str, width: int = 80) -> List[str]:
    """
    Wrap text content preserving line breaks and wrapping long lines.

    Args:
        text: The text to wrap
        width: Maximum width per line (default 80 for safe terminal fit)

    Returns:
        List of wrapped lines
    """
    result = []
    for line in text.split("\n"):
        if len(line) <= width:
            result.append(line)
        else:
            # Wrap long lines
            wrapped = textwrap.wrap(
                line, width=width, break_long_words=False, break_on_hyphens=False
            )
            result.extend(wrapped if wrapped else [""])
    return result


# ----- Help Screen -----


class HelpScreen(ModalScreen):
    """Help modal"""

    BINDINGS = [
        Binding("escape", "dismiss", "Close"),
        Binding("q", "dismiss", "Close"),
    ]

    CSS = """
    HelpScreen {
        align: center middle;
        scrollbar-background: #1a1a1a;
        scrollbar-background-hover: #1a1a1a;
        scrollbar-background-active: #1a1a1a;
        scrollbar-color: #3a3a3a;
        scrollbar-color-hover: #3a3a3a;
        scrollbar-color-active: #3a3a3a;
        scrollbar-corner-color: #1a1a1a;
        scrollbar-size: 1 1;
    }

    #help-container {
        width: 60;
        height: 26;
        background: #121212;
        border: solid #3a3a3a;
        padding: 1 2;
        layout: vertical;
    }

    #help-title {
        text-align: center;
        text-style: bold;
        color: #d4d4d4;
        margin-bottom: 1;
    }

    #help-content {
        color: #9a9a9a;
    }


    #help-close {
        margin-top: 1;
        width: auto;
        min-width: 10;
        background: #1a1a1a;
        color: #9a9a9a;
        border: none;
    }

    #help-close:hover {
        background: #262626;
    }

    #help-close:focus {
        background: #262626;
        text-style: none;
    }
    """

    def compose(self) -> ComposeResult:
        yield Container(
            Static("PentestAgent Help", id="help-title"),
            Static(self._get_help_text(), id="help-content"),
            Center(Button("Close", id="help-close")),
            id="help-container",
        )

    def _get_help_text(self) -> str:
        header = (
            "[bold]Modes:[/] Assist | Agent | Crew\n"
            "[bold]Keys:[/] Enter=Send  Up/Down=History  Ctrl+Q=Quit\n\n"
            "[bold]Commands:[/]\n"
        )

        cmds = [
            ("/agent <task>", "Run in agent mode"),
            ("/crew <task>", "Run multi-agent crew mode"),
            ("/target <host>", "Set target"),
            ("/prompt", "Show system prompt"),
            ("/memory", "Show memory stats"),
            ("/token", "Show token usage & cost"),
            ("/notes", "Show saved notes"),
            ("/report", "Generate report"),
            ("/help", "Show help"),
            ("/clear", "Clear chat"),
            ("/tools", "List tools"),
            ("/quit", "Exit"),
        ]

        # Determine consistent width for command column so the dash aligns
        cmd_col_width = max(len(c) for c, _ in cmds) + 3  # padding before dash
        lines = []
        for cmd, desc in cmds:
            pad = " " * (cmd_col_width - len(cmd))
            lines.append(f"  {cmd}{pad}- {desc}")

        return header + "\n".join(lines)

    def action_dismiss(self) -> None:
        self.app.pop_screen()

    @on(Button.Pressed, "#help-close")
    def close_help(self) -> None:
        self.app.pop_screen()


class ToolsScreen(ModalScreen):
    """Interactive tools browser — split-pane layout.

    Left pane: tree of tools. Right pane: full description (scrollable).
    Selecting another tool replaces the right-pane content. Close returns
    to the main screen.
    """

    BINDINGS = [Binding("escape", "dismiss", "Close"), Binding("q", "dismiss", "Close")]

    CSS = """
    ToolsScreen { align: center middle; }
    """

    def __init__(self, tools: List[Any]) -> None:
        super().__init__()
        self.tools = tools

    def compose(self) -> ComposeResult:
        # Build a split view: left tree, right description
        with Container(id="tools-container"):
            with Horizontal(id="tools-split"):
                with Vertical(id="tools-left"):
                    yield Static("Tools", id="tools-title")
                    yield Tree("TOOLS", id="tools-tree")

                with Vertical(id="tools-right"):
                    yield Static("Description", id="tools-desc-title")
                    yield ScrollableContainer(Static("Select a tool to view details.", id="tools-desc"), id="tools-desc-scroll")

            yield Center(Button("Close", id="tools-close"))

    def on_mount(self) -> None:
        try:
            tree = self.query_one("#tools-tree", Tree)
        except Exception:
            return

        root = tree.root
        root.allow_expand = True
        root.show_root = False

        # Populate tool nodes
        for t in self.tools:
            name = getattr(t, "name", str(t))
            root.add(name, data={"tool": t})

        try:
            tree.focus()
        except Exception:
            pass

    @on(Tree.NodeSelected, "#tools-tree")
    def on_tool_selected(self, event: Tree.NodeSelected) -> None:
        node = event.node
        try:
            tool = node.data.get("tool") if node.data else None
            name = node.label or (getattr(tool, "name", str(tool)) if tool else "Unknown")

            # Prefer Tool.description (registered tools use this), then fall back
            desc = None
            if tool is not None:
                desc = getattr(tool, "description", None)
                if not desc:
                    desc = (
                        getattr(tool, "summary", None)
                        or getattr(tool, "help_text", None)
                        or getattr(tool, "__doc__", None)
                    )
            if not desc:
                desc = "No description available."

            # Update right-hand description pane
            try:
                desc_widget = self.query_one("#tools-desc", Static)
                text = Text()
                text.append(f"{name}\n", style="bold #d4d4d4")
                text.append(str(desc), style="#d4d4d4")
                desc_widget.update(text)
            except Exception:
                pass
        except Exception:
            pass

    @on(Button.Pressed, "#tools-close")
    def close_tools(self) -> None:
        self.app.pop_screen()


# ----- Main Chat Message Widgets -----


class ThinkingMessage(Static):
    """Thinking/reasoning message"""

    def __init__(self, content: str, **kwargs):
        super().__init__(**kwargs)
        self.thinking_content = content

    def render(self) -> Text:
        text = Text()
        text.append("* ", style="#9a9a9a")
        text.append("Thinking\n", style="bold #9a9a9a")

        # Wrap content
        for line in wrap_text_lines(self.thinking_content, width=90):
            text.append(f"  {line}\n", style="#6b6b6b italic")

        return text


class ToolMessage(Static):
    """Tool execution message"""

    # Standard tool icon and color (pa theme)
    TOOL_ICON = "$"
    TOOL_COLOR = "#9a9a9a"  # spirit gray

    def __init__(self, tool_name: str, args: str = "", **kwargs):
        super().__init__(**kwargs)
        self.tool_name = tool_name
        self.tool_args = args

    def render(self) -> Text:
        text = Text()
        text.append(f"{self.TOOL_ICON} ", style=self.TOOL_COLOR)
        text.append(f"{self.tool_name}", style=self.TOOL_COLOR)
        text.append("\n", style="")

        # Wrap args
        if self.tool_args:
            for line in wrap_text_lines(self.tool_args, width=110):
                text.append(f"  {line}\n", style="#6b6b6b")

        return text


class ToolResultMessage(Static):
    """Tool result/output message"""

    RESULT_ICON = "#"
    RESULT_COLOR = "#7a7a7a"

    def __init__(self, tool_name: str, result: str = "", **kwargs):
        super().__init__(**kwargs)
        self.tool_name = tool_name
        self.result = result

    def render(self) -> Text:
        text = Text()
        text.append(f"{self.RESULT_ICON} ", style=self.RESULT_COLOR)
        text.append(f"{self.tool_name} output", style=self.RESULT_COLOR)
        text.append("\n", style="")

        if self.result:
            for line in wrap_text_lines(self.result, width=110):
                text.append(f"  {line}\n", style="#5a5a5a")

        return text


class AssistantMessage(Static):
    """Assistant response message"""

    def __init__(self, content: str, **kwargs):
        super().__init__(**kwargs)
        self.message_content = content

    def render(self) -> Text:
        text = Text()
        text.append(">> ", style="#9a9a9a")
        text.append("PentestAgent\n", style="bold #d4d4d4")

        # Wrap content
        for line in wrap_text_lines(self.message_content, width=90):
            text.append(f"  {line}\n", style="#d4d4d4")

        return text


class UserMessage(Static):
    """User message"""

    def __init__(self, content: str, **kwargs):
        super().__init__(**kwargs)
        self.message_content = content

    def render(self) -> Text:
        text = Text()
        text.append("> ", style="#9a9a9a")
        text.append("You\n", style="bold #d4d4d4")
        text.append(f"  {self.message_content}\n", style="#d4d4d4")
        return text


class SystemMessage(Static):
    """System message"""

    def __init__(self, content: str, **kwargs):
        super().__init__(**kwargs)
        self.message_content = content

    def render(self) -> Text:
        text = Text()
        for line in self.message_content.split("\n"):
            text.append(f"  {line}\n", style="#6b6b6b")  # phantom - subtle system text
        return text


# ----- Status Bar -----


class StatusBar(Static):
    """Animated status bar"""

    status = reactive("idle")
    mode = reactive("assist")  # "assist" or "agent"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._frame = 0
        self._timer: Optional[Timer] = None

    def on_mount(self) -> None:
        self._timer = self.set_interval(0.2, self._tick)

    def _tick(self) -> None:
        self._frame = (self._frame + 1) % 4
        if self.status not in ["idle", "complete"]:
            self.refresh()

    def render(self) -> Text:
        dots = "." * (self._frame + 1)

        # Use fixed-width labels (pad dots to 4 chars so text doesn't jump)
        dots_padded = dots.ljust(4)

        # PA theme status colors (muted, ethereal)
        status_map = {
            "idle": ("Ready", "#6b6b6b"),
            "initializing": (f"Initializing{dots_padded}", "#9a9a9a"),
            "thinking": (f"Thinking{dots_padded}", "#9a9a9a"),
            "running": (f"Running{dots_padded}", "#9a9a9a"),
            "processing": (f"Processing{dots_padded}", "#9a9a9a"),
            "waiting": ("Waiting for input", "#9a9a9a"),
            "complete": ("Complete", "#4a9f6e"),
            "error": ("Error", "#9f4a4a"),
        }

        label, color = status_map.get(self.status, (self.status, "#6b6b6b"))

        text = Text()

        # Show mode (ASCII-safe symbols)
        if self.mode == "crew":
            text.append("  :: Crew ", style="#9a9a9a")
        elif self.mode == "agent":
            text.append("  >> Agent ", style="#9a9a9a")
        else:
            text.append("  >> Assist ", style="#9a9a9a")

        text.append(f"| {label}", style=color)

        if self.status not in ["idle", "initializing", "complete", "error"]:
            text.append("    ESC to stop", style="#525252")

        return text


class MemoryDiagnostics(Static):
    """Live memory diagnostics widget mounted into the chat area.

    This widget polls the agent's LLM memory stats periodically and
    renders a compact, updating diagnostics panel.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._timer: Optional[Timer] = None

    def on_mount(self) -> None:
        # Refresh periodically for a lively display
        self._timer = self.set_interval(0.8, self.refresh)

    def on_unmount(self) -> None:
        if self._timer:
            self._timer.stop()

    def _bar(self, ratio: float, width: int = 20) -> str:
        filled = int(max(0, min(1.0, ratio)) * width)
        return "█" * filled + "░" * (width - filled)

    def render(self) -> Text:
        text = Text()

        try:
            app = self.app
            agent = getattr(app, "agent", None)
            if not agent or not getattr(agent, "llm", None):
                text.append("Memory Diagnostics\n", style="bold #d4d4d4")
                text.append("Agent not initialized", style="#9a9a9a")
                return text

            stats = agent.llm.get_memory_stats()
            msgs = len(agent.conversation_history)
            llm_msgs = agent._format_messages_for_llm()
            current_tokens = agent.llm.memory.get_total_tokens(llm_msgs)

            budget = stats.get("token_budget") or 1
            thresh = stats.get("summarize_threshold") or budget
            recent_keep = stats.get("recent_to_keep", 5)
            has_summary = stats.get("has_summary", False)
            summarized_count = stats.get("summarized_message_count", 0)

            # Header
            text.append("Memory Diagnostics\n", style="bold #d4d4d4")

            # Use a consistent bar width for all bars and align labels
            bar_width = 28
            labels = ["Tokens:", "Messages:", "Retention:"]
            label_width = max(len(label_text) for label_text in labels)

            # Tokens line
            ratio = current_tokens / max(1, budget)
            bar = self._bar(ratio, width=bar_width)
            label = "Tokens:".ljust(label_width)
            text.append(
                f"{label} [{bar}] {current_tokens:,} / {budget:,}\n", style="#9a9a9a"
            )

            # Messages line (scale messages to an expected max window)
            expected_msgs_max = max(1, recent_keep * 6)
            mratio = min(1.0, msgs / expected_msgs_max)
            mbar = self._bar(mratio, width=bar_width)
            label = "Messages:".ljust(label_width)
            text.append(f"{label} [{mbar}] {msgs} active\n", style="#9a9a9a")

            # Retention / recent
            k_ratio = min(1.0, recent_keep / max(1, recent_keep))
            keep_bar = self._bar(k_ratio, width=bar_width)
            label = "Retention:".ljust(label_width)
            text.append(
                f"{label} [{keep_bar}] keeping last {recent_keep}\n", style="#9a9a9a"
            )

            # Summary status
            summary_state = "active" if has_summary else "inactive"
            emoji = "🟢" if has_summary else "🔴"
            text.append(f"Summary: {emoji} {summary_state}\n", style="#9a9a9a")

            # Summarized / threshold
            text.append(
                f"Summarized: {summarized_count} / {thresh:,}\n", style="#9a9a9a"
            )
            text.append(f"Threshold: {thresh:,}\n", style="#9a9a9a")

        except Exception as e:
            text.append(f"Memory diagnostics error: {e}", style="#9a9a9a")

        return text


class TokenDiagnostics(Static):
    """Live token/cost diagnostics panel mounted into the chat area.

    Reads persisted daily usage from the token_tracker, computes cost
    using environment variables, and displays a simple ASCII progress bar.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._timer: Optional[Timer] = None

    def on_mount(self) -> None:
        # Refresh periodically for a lively display
        self._timer = self.set_interval(1.0, self.refresh)

    def on_unmount(self) -> None:
        if self._timer:
            self._timer.stop()

    def _bar(self, ratio: float, width: int = 28) -> str:
        """Block-style usage bar matching MemoryDiagnostics visuals."""
        r = max(0.0, min(1.0, ratio))
        filled = int(r * width)
        return "█" * filled + "░" * (width - filled)

    def render(self) -> Text:
        text = Text()
        try:
            import os

            # Lazy import of token_tracker (best-effort)
            try:
                from ..tools import token_tracker
            except Exception:
                token_tracker = None

            text.append("Token Usage Diagnostics\n", style="bold #d4d4d4")

            if not token_tracker:
                text.append(
                    "Token tracker not available (tools/token_tracker).\n",
                    style="#9a9a9a",
                )
                return text

            stats = token_tracker.get_stats_sync()

            # If a reset is pending (date changed), perform a reset now so daily
            # usage is accurate and visible to the user.
            reset_occurred = False
            if stats.get("reset_pending"):
                try:
                    token_tracker.record_usage_sync(0, 0)
                    stats = token_tracker.get_stats_sync()
                    reset_occurred = True
                except Exception:
                    pass

            # Extract values
            last_in = int(stats.get("last_input_tokens", 0) or 0)
            last_out = int(stats.get("last_output_tokens", 0) or 0)
            last_total = int(stats.get("last_total_tokens", 0) or 0)
            daily_usage = int(stats.get("daily_usage", 0) or 0)
            last_reset = stats.get("last_reset_date")
            current_date = stats.get("current_date")

            # (env parsing moved below)

            # Environment cost config
            def _parse_env(name: str):
                v = os.getenv(name)
                if v is None or v == "":
                    return None
                try:
                    return float(v)
                except Exception:
                    return "INVALID"

            unified = _parse_env("COST_PER_MILLION")
            input_cost_per_m = _parse_env("INPUT_COST_PER_MILLION")
            output_cost_per_m = _parse_env("OUTPUT_COST_PER_MILLION")
            daily_limit = _parse_env("DAILY_TOKEN_LIMIT")

            # Determine if any env-based limits exist
            has_env_limits = any(
                v is not None
                for v in (unified, input_cost_per_m, output_cost_per_m, daily_limit)
            )

            # If nothing has been recorded yet (no tokens, no daily usage)
            # and no env limits are configured, show the concise sentinel only.
            if last_total == 0 and daily_usage == 0 and not has_env_limits:
                text.append("No token usage recorded\n", style="#9a9a9a")
                return text

            # Validate env vars
            env_errors = []
            if unified == "INVALID":
                env_errors.append("COST_PER_MILLION is not numeric")
            if input_cost_per_m == "INVALID":
                env_errors.append("INPUT_COST_PER_MILLION is not numeric")
            if output_cost_per_m == "INVALID":
                env_errors.append("OUTPUT_COST_PER_MILLION is not numeric")
            if daily_limit == "INVALID":
                env_errors.append("DAILY_TOKEN_LIMIT is not numeric")

            if env_errors:
                text.append("Environment configuration errors:\n", style="#ef4444")
                for e in env_errors:
                    text.append(f"  - {e}\n", style="#9a9a9a")
                text.append(
                    "\nSet environment variables correctly to compute costs.\n",
                    style="#9a9a9a",
                )
                return text

            # Compute costs
            if unified is not None:
                # Use unified cost for both input and output
                input_cost = (last_in / 1_000_000.0) * float(unified)
                output_cost = (last_out / 1_000_000.0) * float(unified)
            else:
                # Require per-direction costs to be present to compute
                if input_cost_per_m is None or output_cost_per_m is None:
                    text.append(
                        "Cost vars missing. Set COST_PER_MILLION or both INPUT_COST_PER_MILLION and OUTPUT_COST_PER_MILLION.\n",
                        style="#9a9a9a",
                    )
                    # Still show numeric token stats below
                    input_cost = output_cost = None
                else:
                    input_cost = (last_in / 1_000_000.0) * float(input_cost_per_m)
                    output_cost = (last_out / 1_000_000.0) * float(output_cost_per_m)

            total_cost = None
            if input_cost is not None and output_cost is not None:
                total_cost = input_cost + output_cost

            # Daily budget calculations per spec
            # Derive daily usage excluding last command (in case tracker already included it)
            daily_without_last = max(daily_usage - last_total, 0)
            new_daily_total = daily_without_last + last_total

            remaining_tokens = None
            percent_used = None
            if daily_limit is not None:
                try:
                    dl = float(daily_limit)
                    remaining_tokens = max(int(dl - new_daily_total), 0)
                    percent_used = (new_daily_total / max(1.0, dl)) * 100.0
                except Exception:
                    remaining_tokens = None

            # Render structured panel with aligned labels and block bars
            bar_width = 28
            labels = [
                "Last command:",
                "Cost:",
                "Daily usage:",
                "Remaining:",
                "Usage:",
                "Last reset:",
                "Current date:",
                "Reset occurred:",
            ]
            label_width = max(len(label_text) for label_text in labels)

            # Last command tokens
            label = "Last command:".ljust(label_width)
            text.append(
                f"{label} in={last_in:,} out={last_out:,} total={last_total:,}\n",
                style="#9a9a9a",
            )

            # Cost line
            label = "Cost:".ljust(label_width)
            if input_cost is not None and output_cost is not None:
                text.append(
                    f"{label} in=${input_cost:.6f} out=${output_cost:.6f} total=${total_cost:.6f}\n",
                    style="#9a9a9a",
                )
            else:
                text.append(
                    f"{label} not computed (missing env vars)\n",
                    style="#9a9a9a",
                )

            # Daily usage
            label = "Daily usage:".ljust(label_width)
            text.append(f"{label} {new_daily_total:,}\n", style="#9a9a9a")

            # Remaining tokens
            label = "Remaining:".ljust(label_width)
            if remaining_tokens is not None:
                text.append(f"{label} {remaining_tokens:,}\n", style="#9a9a9a")
            else:
                text.append(
                    f"{label} N/A (DAILY_TOKEN_LIMIT not set)\n",
                    style="#9a9a9a",
                )

            # Usage percent + bar
            label = "Usage:".ljust(label_width)
            if percent_used is not None:
                bar = self._bar(percent_used / 100.0, width=bar_width)
                text.append(
                    f"{label} [{bar}] {percent_used:.1f}%\n",
                    style="#9a9a9a",
                )
            else:
                text.append(f"{label} N/A\n", style="#9a9a9a")

            # Dates
            label = "Last reset:".ljust(label_width)
            text.append(f"{label} {last_reset}\n", style="#9a9a9a")
            label = "Current date:".ljust(label_width)
            text.append(f"{label} {current_date}\n", style="#9a9a9a")

            # Reset occurrence
            label = "Reset occurred:".ljust(label_width)
            text.append(
                f"{label} {'Yes' if reset_occurred else 'No'}\n",
                style="#9a9a9a",
            )

        except Exception as e:
            text.append(f"Token diagnostics error: {e}\n", style="#9a9a9a")

        return text


# ----- Main TUI App -----


class PentestAgentTUI(App):
    """Main PentestAgent TUI Application"""

    # ═══════════════════════════════════════════════════════════
    # PA THEME - Ethereal grays
    # ═══════════════════════════════════════════════════════════
    # Void:       #0a0a0a  (terminal black - the darkness)
    # Shadow:     #121212  (subtle surface)
    # Mist:       #1a1a1a  (panels, elevated)
    # Whisper:    #262626  (default borders)
    # Fog:        #3a3a3a  (hover states)
    # Apparition: #525252  (focus states)
    # Phantom:    #6b6b6b  (secondary text)
    # Spirit:     #9a9a9a  (normal text)
    # Specter:    #d4d4d4  (primary text)
    # Ectoplasm:  #f0f0f0  (highlights)
    # ═══════════════════════════════════════════════════════════

    CSS = """
    Screen {
        background: #0a0a0a;
    }

    #main-container {
        width: 100%;
        height: 100%;
        layout: horizontal;
    }

    /* Chat area - takes full width normally, fills remaining space with sidebar */
    #chat-area {
        width: 1fr;
        height: 100%;
    }

    #chat-area.with-sidebar {
        width: 1fr;
    }

    #chat-scroll {
        width: 100%;
        height: 1fr;
        background: transparent;
        padding: 1 2;
        scrollbar-background: #1a1a1a;
        scrollbar-background-hover: #1a1a1a;
        scrollbar-background-active: #1a1a1a;
        scrollbar-color: #3a3a3a;
        scrollbar-color-hover: #3a3a3a;
        scrollbar-color-active: #3a3a3a;
        scrollbar-corner-color: #1a1a1a;
        scrollbar-size: 1 1;
    }

    #input-container {
        width: 100%;
        height: 3;
        background: transparent;
        border: round #262626;
        margin: 0 2;
        padding: 0;
        layout: horizontal;
        align-vertical: middle;
    }

    #input-container:focus-within {
        border: round #525252;
    }

    #input-container:focus-within #chat-prompt {
        color: #d4d4d4;
    }

    #chat-prompt {
        width: auto;
        height: 100%;
        padding: 0 0 0 1;
        color: #6b6b6b;
        content-align-vertical: middle;
    }

    #chat-input {
        width: 1fr;
        height: 100%;
        background: transparent;
        border: none;
        padding: 0;
        margin: 0;
        color: #d4d4d4;
    }

    #chat-input:focus {
        border: none;
    }

    #chat-input > .input--placeholder {
        color: #6b6b6b;
        text-style: italic;
    }

    #status-bar {
        width: 100%;
        height: 1;
        background: transparent;
        padding: 0 3;
        margin: 0;
    }

    .message {
        margin-bottom: 1;
    }

    /* Sidebar - hidden by default */
    #sidebar {
        width: 28;
        height: 100%;
        display: none;
        padding-right: 1;
    }

    #sidebar.visible {
        display: block;
    }

    #workers-tree {
        height: 1fr;
        background: transparent;
        border: round #262626;
        padding: 0 1;
        margin-bottom: 0;
    }

    #workers-tree:focus {
        border: round #3a3a3a;
    }

    #crew-stats {
        height: auto;
        max-height: 10;
        background: transparent;
        border: round #262626;
        border-title-color: #9a9a9a;
        border-title-style: bold;
        padding: 0 1;
        margin-top: 0;
    }

    Tree {
        background: transparent;
        color: #d4d4d4;
        scrollbar-background: #1a1a1a;
        scrollbar-background-hover: #1a1a1a;
        scrollbar-background-active: #1a1a1a;
        scrollbar-color: #3a3a3a;
        scrollbar-color-hover: #3a3a3a;
        scrollbar-color-active: #3a3a3a;
        scrollbar-size: 1 1;
    }

    Tree > .tree--cursor {
        background: transparent;
    }

    Tree > .tree--highlight {
        background: transparent;
    }

    Tree > .tree--highlight-line {
        background: transparent;
    }

    .tree--node-label {
        padding: 0 1;
    }

    .tree--node:hover .tree--node-label {
        background: transparent;
    }

    .tree--node.-selected .tree--node-label {
        background: transparent;
        color: #d4d4d4;
    }
    """

    BINDINGS = [
        Binding("ctrl+q", "quit_app", "Quit", priority=True),
        Binding("ctrl+c", "stop_agent", "Stop", priority=True, show=False),
        Binding("escape", "stop_agent", "Stop", priority=True),
        Binding("f1", "show_help", "Help"),
        Binding("tab", "focus_next", "Next", show=False),
        Binding("up", "history_up", "Prev", show=False),
        Binding("down", "history_down", "Next", show=False),
    ]

    TITLE = "PentestAgent"
    SUB_TITLE = "AI Penetration Testing"

    def __init__(
        self,
        target: Optional[str] = None,
        model: Optional[str] = None,
        use_docker: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.target = target
        self.model = model or DEFAULT_MODEL
        self.use_docker = use_docker

        # Agent components
        self.agent: Optional["PentestAgentAgent"] = None
        self.runtime = None
        self.mcp_manager = None
        self.all_tools = []
        self.rag_engine = None  # RAG engine

        # State
        self._mode = "assist"  # "assist", "agent", or "crew"
        self._is_running = False
        self._is_initializing = True  # Block input during init
        self._should_stop = False
        # Worker handle returned by `@work` or an `asyncio.Task` (keep generic)
        self._current_worker: Optional[Any] = None  # Track running worker for cancellation
        self._current_crew = None  # Track crew orchestrator for cancellation

        # Crew mode state
        self._crew_workers: Dict[str, Dict[str, Any]] = {}
        self._crew_worker_nodes: Dict[str, TreeNode] = {}
        self._crew_orchestrator_node: Optional[TreeNode] = None
        self._crew_findings_count = 0
        self._viewing_worker_id: Optional[str] = None
        self._worker_events: Dict[str, List[Dict]] = {}
        self._crew_start_time: Optional[float] = None
        self._crew_tokens_used: int = 0
        self._crew_stats_timer: Optional[Timer] = None
        self._spinner_timer: Optional[Timer] = None
        self._spinner_frame: int = 0
        self._spinner_frames = [
            "⠋",
            "⠙",
            "⠹",
            "⠸",
            "⠼",
            "⠴",
            "⠦",
            "⠧",
            "⠇",
            "⠏",
        ]  # Braille dots spinner

        # Command history
        self._cmd_history: List[str] = []
        self._history_index: int = 0

    def compose(self) -> ComposeResult:
        with Horizontal(id="main-container"):
            # Chat area (left side)
            with Vertical(id="chat-area"):
                yield ScrollableContainer(id="chat-scroll")
                yield StatusBar(id="status-bar")
                with Horizontal(id="input-container"):
                    yield Static("> ", id="chat-prompt")
                    yield Input(placeholder="Enter task or type /help", id="chat-input")

            # Sidebar (right side, hidden by default)
            with Vertical(id="sidebar"):
                yield CrewTree("CREW", id="workers-tree")
                yield Static("", id="crew-stats")

    async def on_mount(self) -> None:
        """Initialize on mount"""
        # Call the textual worker - decorator returns a Worker, not a coroutine
        _ = cast(Any, self._initialize_agent())

    @work(thread=False)
    async def _initialize_agent(self) -> None:
        """Initialize agent"""
        self._set_status("initializing")

        try:
            import os

            from ..agents.pa_agent import PentestAgentAgent
            from ..knowledge import RAGEngine
            from ..llm import LLM, ModelConfig
            from ..mcp import MCPManager
            from ..runtime.docker_runtime import DockerRuntime
            from ..runtime.runtime import LocalRuntime
            from ..tools import get_all_tools, register_tool_instance

            # RAG Engine - auto-load knowledge sources
            rag_doc_count = 0
            knowledge_path = None

            # Check local knowledge dir first (must have files, not just exist)
            local_knowledge = Path("knowledge")
            bundled_path = Path(__file__).parent.parent / "knowledge" / "sources"

            if local_knowledge.exists() and any(local_knowledge.rglob("*.*")):
                knowledge_path = local_knowledge
            elif bundled_path.exists():
                knowledge_path = bundled_path

            if knowledge_path:
                try:
                    # Determine embedding method: env var > auto-detect
                    embeddings_setting = os.getenv(
                        "PENTESTAGENT_EMBEDDINGS", ""
                    ).lower()
                    if embeddings_setting == "local":
                        use_local = True
                    elif embeddings_setting == "openai":
                        use_local = False
                    else:
                        # Auto: use OpenAI if key available, else local
                        use_local = not os.getenv("OPENAI_API_KEY")

                    self.rag_engine = RAGEngine(
                        knowledge_path=knowledge_path, use_local_embeddings=use_local
                    )
                    await asyncio.to_thread(self.rag_engine.index)
                    rag_doc_count = self.rag_engine.get_document_count()
                except Exception as e:
                    self._add_system(f"[!] RAG: {e}")
                    self.rag_engine = None

            # MCP - auto-load if config exists
            mcp_server_count = 0
            try:
                self.mcp_manager = MCPManager()
                if self.mcp_manager.config_path.exists():
                    mcp_tools = await self.mcp_manager.connect_all()
                    for tool in mcp_tools:
                        register_tool_instance(tool)
                    mcp_server_count = len(self.mcp_manager.servers)
            except Exception as e:
                self._add_system(f"[!] MCP: {e}")

            # Runtime - Docker or Local
            if self.use_docker:
                self._add_system("+ Starting Docker container...")
                self.runtime = DockerRuntime(mcp_manager=self.mcp_manager)
            else:
                self.runtime = LocalRuntime(mcp_manager=self.mcp_manager)
            await self.runtime.start()

            # LLM
            # Ensure types for static analysis: runtime and model are set
            assert self.model is not None
            assert self.runtime is not None
            llm = LLM(
                model=self.model,
                config=ModelConfig(temperature=0.7),
                rag_engine=self.rag_engine,
            )

            # Tools
            self.all_tools = get_all_tools()  # Ensure tools are loaded

            # Agent
            self.agent = PentestAgentAgent(
                llm=llm,
                tools=self.all_tools,
                runtime=self.runtime,
                target=self.target,
                rag_engine=self.rag_engine,
            )

            self._set_status("idle", "assist")
            self._is_initializing = False  # Allow input now

            # Show ready message
            tools_str = ", ".join(t.name for t in self.all_tools[:5])
            if len(self.all_tools) > 5:
                tools_str += f", +{len(self.all_tools) - 5} more"

            runtime_str = "Docker" if self.use_docker else "Local"
            self._add_system(
                f"+ PentestAgent ready\n"
                f"  Model: {self.model} | Tools: {len(self.all_tools)} | MCP: {mcp_server_count} | RAG: {rag_doc_count}\n"
                f"  Runtime: {runtime_str} | Mode: Assist (use /agent or /crew for autonomous modes)"
            )

            # Show target if provided (but don't auto-start)
            if self.target:
                self._add_system(f"  Target: {self.target}")

        except Exception as e:
            import traceback

            self._add_system(f"[!] Init failed: {e}\n{traceback.format_exc()}")
            self._set_status("error")
            self._is_initializing = False  # Allow input even on error

    def _set_status(self, status: str, mode: Optional[str] = None) -> None:
        """Update status bar"""
        try:
            bar = self.query_one("#status-bar", StatusBar)
            bar.status = status
            if mode:
                bar.mode = mode
                self._mode = mode
        except Exception:
            pass

    def _add_message(self, widget: Static) -> None:
        """Add a message widget to chat"""
        try:
            scroll = self.query_one("#chat-scroll", ScrollableContainer)
            widget.add_class("message")
            scroll.mount(widget)
            scroll.scroll_end(animate=False)
        except Exception:
            pass

    def _add_system(self, content: str) -> None:
        self._add_message(SystemMessage(content))

    def _add_user(self, content: str) -> None:
        self._add_message(UserMessage(content))

    def _add_assistant(self, content: str) -> None:
        self._add_message(AssistantMessage(content))

    def _add_thinking(self, content: str) -> None:
        self._add_message(ThinkingMessage(content))

    def _add_tool(self, name: str, action: str = "") -> None:
        self._add_message(ToolMessage(name, action))

    def _add_tool_result(self, name: str, result: str) -> None:
        """Display tool execution result"""
        # Hide tool output - LLM will synthesize it in its response
        # This prevents duplication and keeps the chat clean
        pass

    def _show_system_prompt(self) -> None:
        """Display the current system prompt"""
        if self.agent:
            prompt = self.agent.get_system_prompt()
            self._add_system(f"=== System Prompt ===\n{prompt}")
        else:
            self._add_system("Agent not initialized")

    def _show_memory_stats(self) -> None:
        """Mount a live memory diagnostics widget into the chat area."""
        try:
            scroll = self.query_one("#chat-scroll", ScrollableContainer)
        except Exception:
            self._add_system("Agent not initialized")
            return
        # Mount a new diagnostics panel with a unique ID and scroll into view
        try:
            import uuid

            panel_id = f"memory-diagnostics-{uuid.uuid4().hex}"
        except Exception:
            panel_id = None

        widget = MemoryDiagnostics(id=panel_id)
        scroll.mount(widget)
        try:
            scroll.scroll_end(animate=False)
        except Exception:
            pass

    def _show_token_stats(self) -> None:
        """Mount a live token diagnostics widget into the chat area."""
        try:
            scroll = self.query_one("#chat-scroll", ScrollableContainer)
        except Exception:
            self._add_system("Agent not initialized")
            return
        # Mount a new diagnostics panel with a unique ID and scroll into view
        try:
            import uuid

            panel_id = f"token-diagnostics-{uuid.uuid4().hex}"
        except Exception:
            panel_id = None

        widget = TokenDiagnostics(id=panel_id)
        scroll.mount(widget)
        try:
            scroll.scroll_end(animate=False)
        except Exception:
            pass

    async def _show_notes(self) -> None:
        """Display saved notes"""
        from ..tools.notes import get_all_notes

        notes = await get_all_notes()
        if not notes:
            self._add_system(
                "=== Notes ===\nNo notes saved.\n\nThe AI can save key findings using the notes tool."
            )
            return

        lines = [f"=== Notes ({len(notes)} entries) ==="]
        for key, value in notes.items():
            # Show full value, indent multi-line content
            if "\n" in value:
                indented = str(value).replace("\n", "\n    ")
                lines.append(f"\n[{key}]\n    {indented}")
            else:
                lines.append(f"[{key}] {value}")
        lines.append("\nFile: loot/notes.json")
        lines.append("Reports: loot/reports/")

        self._add_system("\n".join(lines))

    def _build_prior_context(self) -> str:
        """Build a summary of prior findings for crew mode.

        Extracts:
        - Tool results (nmap scans, etc.) - the actual findings
        - Assistant analyses - interpretations and summaries
        - Last user task - what they were working on

        Excludes:
        - Raw user messages (noise)
        - Tool call declarations (just names/args, not results)
        - Very short responses
        """
        if not self.agent or not self.agent.conversation_history:
            return ""

        findings = []
        last_user_task = ""

        for msg in self.agent.conversation_history:
            # Track user tasks/questions
            if msg.role == "user" and msg.content:
                last_user_task = msg.content[:200]

            # Extract tool results (the actual findings)
            elif msg.tool_results:
                for result in msg.tool_results:
                    if result.success and result.result:
                        content = (
                            result.result[:1500]
                            if len(result.result) > 1500
                            else result.result
                        )
                        findings.append(f"[{result.tool_name}]\n{content}")

            # Include assistant analyses (but not tool call messages)
            elif msg.role == "assistant" and msg.content and not msg.tool_calls:
                if len(msg.content) > 50:
                    findings.append(f"[Analysis]\n{msg.content[:1000]}")

        if not findings and not last_user_task:
            return ""

        # Build context with last user task + recent findings
        parts = []
        if last_user_task:
            parts.append(f"Last task: {last_user_task}")
        if findings:
            parts.append("Findings:\n" + "\n\n".join(findings[-5:]))

        context = "\n\n".join(parts)
        if len(context) > 4000:
            context = context[:4000] + "\n... (truncated)"

        return context

    def _set_target(self, cmd: str) -> None:
        """Set the target for the engagement"""
        # Remove /target prefix
        target = cmd[7:].strip()

        if not target:
            if self.target:
                self._add_system(
                    f"Current target: {self.target}\nUsage: /target <host>"
                )
            else:
                self._add_system(
                    "No target set.\nUsage: /target <host>\nExample: /target 192.168.1.1"
                )
            return

        self.target = target

        # Update agent's target if agent exists
        if self.agent:
            self.agent.target = target
        # Update the initial ready SystemMessage (if present) so Target appears under Runtime
        try:
            scroll = self.query_one("#chat-scroll", ScrollableContainer)
            updated = False
            for child in scroll.children:
                if isinstance(child, SystemMessage) and "PentestAgent ready" in getattr(
                    child, "message_content", ""
                ):
                    # Append Target line if not already present
                    if "Target:" not in child.message_content:
                        child.message_content = (
                            child.message_content + f"\n  Target: {target}"
                        )
                        try:
                            child.refresh()
                        except Exception:
                            pass
                    updated = True
                    break
            if not updated:
                # Fallback: add system line near top by inserting at beginning
                try:
                    first = scroll.children[0] if scroll.children else None
                    msg = SystemMessage(f"  Target: {target}")
                    if first:
                        scroll.mount_before(msg, first)
                    else:
                        scroll.mount(msg)
                except Exception:
                    self._add_system(f"  Target: {target}")
        except Exception:
            # Last resort: append a subtle system line
            self._add_system(f"  Target: {target}")

    @work(exclusive=True)
    async def _run_report_generation(self) -> None:
        """Generate a pentest report from notes and conversation"""
        from pathlib import Path

        from ..tools.notes import get_all_notes

        if not self.agent or not self.agent.llm:
            self._add_system("[!] Agent not initialized")
            return

        notes = await get_all_notes()
        if not notes:
            self._add_system(
                "No notes found. PentestAgent saves findings using the notes tool during testing."
            )
            return

        self._add_system("Generating report...")

        # Format notes
        notes_text = "\n".join(f"### {k}\n{v}\n" for k, v in notes.items())

        # Build conversation summary from full history
        conversation_summary = ""
        if self.agent.conversation_history:
            # Summarize key actions from conversation
            actions = []
            for msg in self.agent.conversation_history:
                if msg.role == "assistant" and msg.tool_calls:
                    for tc in msg.tool_calls:
                        actions.append(f"- Tool: {tc.name}")
                elif msg.role == "tool_result" and msg.tool_results:
                    for tr in msg.tool_results:
                        # Include truncated result
                        result = tr.result or ""
                        output = result[:200] + "..." if len(result) > 200 else result
                        actions.append(f"  Result: {output}")
            if actions:
                conversation_summary = "\n".join(actions[-30:])  # Last 30 actions

        report_prompt = f"""Generate a penetration test report in Markdown from the notes below.

# Notes
{notes_text}

# Activity Log
{conversation_summary if conversation_summary else "N/A"}

# Target
{self.target or "Not specified"}

Output a report with:
1. Executive Summary (2-3 sentences)
2. Findings (use notes, include severity: Critical/High/Medium/Low/Info)
3. Recommendations

Be concise. Use the actual data from notes."""

        try:
            report_content = await self.agent.llm.simple_completion(
                prompt=report_prompt,
                system="You are a penetration tester writing a security report. Be concise and factual.",
            )

            if not report_content or not report_content.strip():
                self._add_system(
                    "[!] Report generation returned empty. Check LLM connection."
                )
                return

            # Save to loot/reports/
            reports_dir = Path("loot/reports")
            reports_dir.mkdir(parents=True, exist_ok=True)

            # Append Shadow Graph if available
            try:
                from ..knowledge.graph import ShadowGraph
                from ..tools.notes import get_all_notes_sync

                # Rehydrate graph from notes
                graph = ShadowGraph()
                notes = get_all_notes_sync()
                if notes:
                    graph.update_from_notes(notes)
                    mermaid_code = graph.to_mermaid()

                    if mermaid_code:
                        report_content += (
                            "\n\n## Attack Graph (Visual)\n\n```mermaid\n"
                            + mermaid_code
                            + "\n```\n"
                        )
            except Exception as e:
                self._add_system(f"[!] Graph generation error: {e}")

            timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
            report_path = reports_dir / f"report_{timestamp}.md"
            report_path.write_text(report_content, encoding="utf-8")

            self._add_system(f"+ Report saved: {report_path}")

        except Exception as e:
            self._add_system(f"[!] Report error: {e}")

    @on(Input.Submitted, "#chat-input")
    async def handle_submit(self, event: Input.Submitted) -> None:
        """Handle input submission"""
        # Block input while initializing or AI is processing
        if self._is_initializing or self._is_running:
            return

        # Strip ANSI escape sequences and control codes
        message = _ANSI_ESCAPE.sub("", event.value).strip()
        if not message:
            return

        # Save to history (de-duplicate consecutive duplicates)
        if not (self._cmd_history and self._cmd_history[-1] == message):
            self._cmd_history.append(message)
        # Reset index to one past the end (blank)
        self._history_index = len(self._cmd_history)

        event.input.value = ""

        # Commands
        if message.startswith("/"):
            await self._handle_command(message)
            return

        self._add_user(message)

        # Hide crew sidebar when entering assist mode
        self._hide_sidebar()

        # Use assist mode by default
        if self.agent and not self._is_running:
            # Schedule assist run and keep task handle
            self._current_worker = asyncio.create_task(self._run_assist(message))

    async def _handle_command(self, cmd: str) -> None:
        """Handle slash commands"""
        cmd_lower = cmd.lower().strip()
        cmd_original = cmd.strip()

        if cmd_lower in ["/help", "/h", "/?"]:
            await self.push_screen(HelpScreen())
        elif cmd_lower == "/clear":
            scroll = self.query_one("#chat-scroll", ScrollableContainer)
            await scroll.remove_children()
            self._hide_sidebar()
            # Clear agent conversation history for fresh start
            if self.agent:
                self.agent.conversation_history.clear()
            self._add_system("Chat cleared")
        elif cmd_lower == "/tools":
            # Open the interactive tools browser (split-pane).
            try:
                await self.push_screen(ToolsScreen(tools=self.all_tools))
            except Exception:
                # Fallback: list tools in the system area if UI push fails
                from ..runtime.runtime import detect_environment

                names = [t.name for t in self.all_tools]
                msg = f"Tools ({len(names)}): " + ", ".join(names)

                # Add detected CLI tools
                env = detect_environment()
                if env.available_tools:
                    # Group by category
                    by_category = {}
                    for tool_info in env.available_tools:
                        if tool_info.category not in by_category:
                            by_category[tool_info.category] = []
                        by_category[tool_info.category].append(tool_info.name)

                    cli_sections = []
                    for category in sorted(by_category.keys()):
                        tools_list = ", ".join(sorted(by_category[category]))
                        cli_sections.append(f"{category}: {tools_list}")

                    msg += f"\n\nCLI Tools ({len(env.available_tools)}):\n" + "\n".join(
                        cli_sections
                    )

                self._add_system(msg)
        elif cmd_lower in ["/quit", "/exit", "/q"]:
            self.exit()
        elif cmd_lower == "/prompt":
            self._show_system_prompt()
        elif cmd_lower == "/memory":
            self._show_memory_stats()
        elif cmd_lower == "/token":
            self._show_token_stats()
        elif cmd_lower == "/notes":
            await self._show_notes()
        elif cmd_lower == "/report":
            # Call the textual worker - decorator returns a Worker
            _ = cast(Any, self._run_report_generation())
        elif cmd_original.startswith("/target"):
            self._set_target(cmd_original)
        elif cmd_original.startswith("/agent"):
            await self._parse_agent_command(cmd_original)
        elif cmd_original.startswith("/crew"):
            await self._parse_crew_command(cmd_original)
        else:
            self._add_system(f"Unknown command: {cmd}\nType /help for commands.")

    async def _parse_agent_command(self, cmd: str) -> None:
        """Parse and execute /agent command"""

        # Remove /agent prefix
        rest = cmd[6:].strip()

        if not rest:
            self._add_system(
                "Usage: /agent <task>\n"
                "Example: /agent scan 192.168.1.1\n"
                "         /agent enumerate SSH on target"
            )
            return

        task = rest

        if not task:
            self._add_system("Error: No task provided. Usage: /agent <task>")
            return

        self._add_user(f"/agent {task}")
        self._add_system(">> Agent Mode")

        # Hide crew sidebar when entering agent mode
        self._hide_sidebar()

        if self.agent and not self._is_running:
            # Schedule agent mode and keep task handle
            self._current_worker = asyncio.create_task(self._run_agent_mode(task))

    async def _parse_crew_command(self, cmd: str) -> None:
        """Parse and execute /crew command"""
        # Remove /crew prefix
        rest = cmd[5:].strip()

        if not rest:
            self._add_system(
                "Usage: /crew <task>\n"
                "Example: /crew https://example.com\n"
                "         /crew 192.168.1.100\n\n"
                "Crew mode spawns specialized workers in parallel:\n"
                "  - recon: Reconnaissance and mapping\n"
                "  - sqli: SQL injection testing\n"
                "  - xss: Cross-site scripting testing\n"
                "  - ssrf: Server-side request forgery\n"
                "  - auth: Authentication testing\n"
                "  - idor: Insecure direct object references\n"
                "  - info: Information disclosure"
            )
            return

        target = rest

        if not self._is_running:
            self._add_user(f"/crew {target}")
            self._show_sidebar()
            # Schedule crew mode and keep handle
            self._current_worker = asyncio.create_task(self._run_crew_mode(target))

    def _show_sidebar(self) -> None:
        """Show the sidebar for crew mode."""
        try:
            import time

            sidebar = self.query_one("#sidebar")
            sidebar.add_class("visible")

            chat_area = self.query_one("#chat-area")
            chat_area.add_class("with-sidebar")

            # Setup tree
            tree = self.query_one("#workers-tree", CrewTree)
            tree.root.expand()
            tree.show_root = False

            # Clear old nodes
            tree.root.remove_children()
            self._crew_worker_nodes.clear()
            self._crew_workers.clear()
            self._worker_events.clear()
            self._crew_findings_count = 0

            # Start tracking time and tokens
            self._crew_start_time = time.time()
            self._crew_tokens_used = 0

            # Start stats timer (update every second)
            if self._crew_stats_timer:
                self._crew_stats_timer.stop()
            self._crew_stats_timer = self.set_interval(1.0, self._update_crew_stats)

            # Start spinner timer for running workers (faster interval for smooth animation)
            if self._spinner_timer:
                self._spinner_timer.stop()
            self._spinner_timer = self.set_interval(0.15, self._update_spinner)

            # Add crew root node (no orchestrator - just "CREW" header)
            self._crew_orchestrator_node = tree.root.add(
                "CREW", data={"type": "crew", "id": "crew"}
            )
            if self._crew_orchestrator_node:
                try:
                    self._crew_orchestrator_node.expand()
                    tree.select_node(self._crew_orchestrator_node)
                except Exception:
                    pass
            self._viewing_worker_id = None

            # Update stats
            self._update_crew_stats()
        except Exception as e:
            self._add_system(f"[!] Sidebar error: {e}")

    def _hide_sidebar(self) -> None:
        """Hide the sidebar."""
        try:
            # Stop stats timer
            if self._crew_stats_timer:
                self._crew_stats_timer.stop()
                self._crew_stats_timer = None

            sidebar = self.query_one("#sidebar")
            sidebar.remove_class("visible")

            chat_area = self.query_one("#chat-area")
            chat_area.remove_class("with-sidebar")
        except Exception:
            pass

    def _update_crew_stats(self) -> None:
        """Update crew stats panel."""
        try:
            import time

            text = Text()

            # Elapsed time
            text.append("Time:   ", style="bold #d4d4d4")
            if self._crew_start_time:
                elapsed = time.time() - self._crew_start_time
                if elapsed < 60:
                    time_str = f"{int(elapsed)}s"
                elif elapsed < 3600:
                    mins = int(elapsed // 60)
                    secs = int(elapsed % 60)
                    time_str = f"{mins}m {secs}s"
                else:
                    hrs = int(elapsed // 3600)
                    mins = int((elapsed % 3600) // 60)
                    time_str = f"{hrs}h {mins}m"
                text.append(time_str, style="#9a9a9a")
            else:
                text.append("--", style="#525252")

            text.append("\n")

            # Tokens used
            text.append("Tokens: ", style="bold #d4d4d4")
            if self._crew_tokens_used > 0:
                if self._crew_tokens_used >= 1000:
                    token_str = f"{self._crew_tokens_used / 1000:.1f}k"
                else:
                    token_str = str(self._crew_tokens_used)
                text.append(token_str, style="#9a9a9a")
            else:
                text.append("--", style="#525252")

            stats = self.query_one("#crew-stats", Static)
            stats.update(text)
            stats.border_title = "# Stats"
        except Exception:
            pass

    def _update_spinner(self) -> None:
        """Update spinner animation for running workers."""
        try:
            # Advance spinner frame
            self._spinner_frame += 1

            # Only update labels for running workers (efficient)
            has_running = False
            for worker_id, worker in self._crew_workers.items():
                if worker.get("status") == "running":
                    has_running = True
                    # Update the tree node label
                    if worker_id in self._crew_worker_nodes:
                        node = self._crew_worker_nodes[worker_id]
                        node.set_label(self._format_worker_label(worker_id))

            # Stop spinner if no workers are running (save resources)
            if not has_running and self._spinner_timer:
                self._spinner_timer.stop()
                self._spinner_timer = None
        except Exception:
            pass

    def _add_crew_worker(self, worker_id: str, worker_type: str, task: str) -> None:
        """Add a worker to the sidebar tree."""
        self._crew_workers[worker_id] = {
            "worker_type": worker_type,
            "task": task,
            "status": "pending",
            "findings": 0,
        }

        try:
            label = self._format_worker_label(worker_id)
            if self._crew_orchestrator_node:
                node = self._crew_orchestrator_node.add(
                    label, data={"type": "worker", "id": worker_id}
                )
                self._crew_worker_nodes[worker_id] = node
                try:
                    self._crew_orchestrator_node.expand()
                except Exception:
                    pass
            self._update_crew_stats()
        except Exception:
            pass

    def _update_crew_worker(self, worker_id: str, **updates) -> None:
        """Update a worker's state."""
        if worker_id not in self._crew_workers:
            return

        self._crew_workers[worker_id].update(updates)

        # Restart spinner if a worker started running
        if updates.get("status") == "running" and not self._spinner_timer:
            self._spinner_timer = self.set_interval(0.15, self._update_spinner)

        try:
            if worker_id in self._crew_worker_nodes:
                label = self._format_worker_label(worker_id)
                self._crew_worker_nodes[worker_id].set_label(label)
            self._update_crew_stats()
        except Exception:
            pass

    def _format_worker_label(self, worker_id: str) -> Text:
        """Format worker label for tree."""
        worker = self._crew_workers.get(worker_id, {})
        status = worker.get("status", "pending")
        wtype = worker.get("worker_type", "worker")
        findings = worker.get("findings", 0)

        # 4-state icons: working (braille), done (checkmark), warning (!), error (X)
        if status in ("running", "pending"):
            # Animated braille spinner for all in-progress states
            icon = self._spinner_frames[self._spinner_frame % len(self._spinner_frames)]
            color = "#d4d4d4"  # white
        elif status == "complete":
            icon = "✓"
            color = "#22c55e"  # green
        elif status == "warning":
            icon = "!"
            color = "#f59e0b"  # amber/orange
        else:  # error, cancelled, unknown
            icon = "✗"
            color = "#ef4444"  # red

        text = Text()
        text.append(f"{icon} ", style=color)
        text.append(wtype.upper(), style="bold")

        if status == "complete" and findings > 0:
            text.append(f" [{findings}]", style="#22c55e")  # green
        elif status in ("error", "cancelled"):
            # Don't append " !" here since we already have the X icon
            pass

        return text

    def _handle_worker_event(
        self, worker_id: str, event_type: str, data: Dict[str, Any]
    ) -> None:
        """Handle worker events from CrewAgent - updates tree sidebar only."""
        try:
            if event_type == "spawn":
                worker_type = data.get("worker_type", "unknown")
                task = data.get("task", "")
                self._add_crew_worker(worker_id, worker_type, task)
            elif event_type == "status":
                status = data.get("status", "running")
                self._update_crew_worker(worker_id, status=status)
            elif event_type == "tool":
                # Add tool as child node under the agent
                tool_name = data.get("tool", "unknown")
                self._add_tool_to_worker(worker_id, tool_name)
            elif event_type == "tokens":
                # Track token usage
                tokens = data.get("tokens", 0)
                self._crew_tokens_used += tokens
            elif event_type == "complete":
                findings_count = data.get("findings_count", 0)
                self._update_crew_worker(
                    worker_id, status="complete", findings=findings_count
                )
                self._crew_findings_count += findings_count
                self._update_crew_stats()
            elif event_type == "warning":
                # Worker hit max iterations but has results
                self._update_crew_worker(worker_id, status="warning")
                reason = data.get("reason", "Partial completion")
                worker = self._crew_workers.get(worker_id, {})
                wtype = worker.get("worker_type", "worker")
                self._add_system(f"[!] {wtype.upper()} stopped: {reason}")
                self._update_crew_stats()
            elif event_type == "failed":
                # Worker determined task infeasible
                self._update_crew_worker(worker_id, status="failed")
                reason = data.get("reason", "Task infeasible")
                worker = self._crew_workers.get(worker_id, {})
                wtype = worker.get("worker_type", "worker")
                self._add_system(f"[!] {wtype.upper()} failed: {reason}")
                self._update_crew_stats()
            elif event_type == "error":
                self._update_crew_worker(worker_id, status="error")
                worker = self._crew_workers.get(worker_id, {})
                wtype = worker.get("worker_type", "worker")
                error_msg = data.get("error", "Unknown error")
                # Only show errors in chat - they're important
                self._add_system(f"[!] {wtype.upper()} failed: {error_msg}")
        except Exception as e:
            self._add_system(f"[!] Worker event error: {e}")

    def _add_tool_to_worker(self, worker_id: str, tool_name: str) -> None:
        """Add a tool usage as child node under worker in tree."""
        try:
            node = self._crew_worker_nodes.get(worker_id)
            if node:
                node.add_leaf(f"  {tool_name}")
                node.expand()
        except Exception:
            pass

    @on(Tree.NodeSelected, "#workers-tree")
    def on_worker_tree_selected(self, event: Tree.NodeSelected) -> None:
        """Handle tree node selection."""
        node = event.node
        if node.data:
            node_type = node.data.get("type")
            if node_type == "crew":
                self._viewing_worker_id = None
            elif node_type == "worker":
                self._viewing_worker_id = node.data.get("id")

    @work(thread=False)
    async def _run_crew_mode(self, target: str) -> None:
        """Run crew mode with sidebar."""
        self._is_running = True
        self._should_stop = False
        self._set_status("thinking", "crew")

        try:
            from ..agents.base_agent import AgentMessage
            from ..agents.crew import CrewOrchestrator
            from ..llm import LLM, ModelConfig

            # Build prior context from assist/agent conversation history
            prior_context = self._build_prior_context()

            # Ensure model/runtime are available for static analysis
            assert self.model is not None
            assert self.runtime is not None

            llm = LLM(model=self.model, config=ModelConfig(temperature=0.7))

            crew = CrewOrchestrator(
                llm=llm,
                tools=self.all_tools,
                runtime=self.runtime,
                on_worker_event=self._handle_worker_event,
                rag_engine=self.rag_engine,
                target=target,
                prior_context=prior_context,
            )
            self._current_crew = crew  # Track for cancellation

            self._add_system(f"@ Task: {target}")

            # Track crew results for memory
            crew_report = None

            async for update in crew.run(target):
                if self._should_stop:
                    await crew.cancel()
                    self._add_system("[!] Stopped by user")
                    break

                phase = update.get("phase", "")

                if phase == "starting":
                    self._set_status("thinking", "crew")

                elif phase == "tokens":
                    # Track orchestrator token usage
                    tokens = update.get("tokens", 0)
                    self._crew_tokens_used += tokens
                    self._update_crew_stats()

                elif phase == "thinking":
                    # Show the orchestrator's reasoning
                    content = update.get("content", "")
                    if content:
                        self._add_thinking(content)

                elif phase == "tool_call":
                    # Show orchestration tool calls
                    tool = update.get("tool", "")
                    args = update.get("args", {})
                    self._add_tool(tool, str(args))

                elif phase == "tool_result":
                    # Tool results are tracked via worker events
                    pass

                elif phase == "complete":
                    crew_report = update.get("report", "")
                    if crew_report:
                        self._add_assistant(crew_report)

                elif phase == "error":
                    error = update.get("error", "Unknown error")
                    self._add_system(f"[!] Crew error: {error}")

            # Add crew results to main agent's conversation history
            # so assist mode can reference what happened
            if self.agent and crew_report:
                # Add the crew task as a user message
                self.agent.conversation_history.append(
                    AgentMessage(
                        role="user",
                        content=f"[CREW MODE] Run parallel analysis on target: {target}",
                    )
                )
                # Add the crew report as assistant response
                self.agent.conversation_history.append(
                    AgentMessage(role="assistant", content=crew_report)
                )

            self._set_status("complete", "crew")
            self._add_system("+ Crew task complete.")

            # Stop timers
            if self._crew_stats_timer:
                self._crew_stats_timer.stop()
                self._crew_stats_timer = None
            if self._spinner_timer:
                self._spinner_timer.stop()
                self._spinner_timer = None

            # Clear crew reference
            self._current_crew = None

        except asyncio.CancelledError:
            # Cancel crew workers first
            if self._current_crew:
                await self._current_crew.cancel()
                self._current_crew = None
            self._add_system("[!] Cancelled")
            self._set_status("idle", "crew")
            # Stop timers on cancel
            if self._crew_stats_timer:
                self._crew_stats_timer.stop()
                self._crew_stats_timer = None
            if self._spinner_timer:
                self._spinner_timer.stop()
                self._spinner_timer = None

        except Exception as e:
            import traceback

            # Cancel crew workers on error too
            if self._current_crew:
                try:
                    await self._current_crew.cancel()
                except Exception:
                    pass
                self._current_crew = None
            self._add_system(f"[!] Crew error: {e}\n{traceback.format_exc()}")
            self._set_status("error")
            # Stop timers on error too
            if self._crew_stats_timer:
                self._crew_stats_timer.stop()
                self._crew_stats_timer = None
            if self._spinner_timer:
                self._spinner_timer.stop()
                self._spinner_timer = None
        finally:
            self._is_running = False

    @work(thread=False)
    async def _run_assist(self, message: str) -> None:
        """Run in assist mode - single response"""
        if not self.agent:
            self._add_system("[!] Agent not ready")
            return

        self._is_running = True
        self._should_stop = False
        self._set_status("thinking", "assist")

        try:
            async for response in self.agent.assist(message):
                if self._should_stop:
                    self._add_system("[!] Stopped by user")
                    break

                self._set_status("processing")

                # Show thinking/plan FIRST if there's content with tool calls
                if response.content:
                    content = response.content.strip()
                    if response.tool_calls:
                        self._add_thinking(content)
                    else:
                        self._add_assistant(content)

                # Show tool calls (skip 'finish' - internal control)
                if response.tool_calls:
                    for call in response.tool_calls:
                        if call.name == "finish":
                            continue  # Skip - summary shown as final message
                        args_str = str(call.arguments)
                        self._add_tool(call.name, args_str)

                # Show tool results (displayed after execution completes)
                # Skip 'finish' tool - its result is shown as the final summary
                if response.tool_results:
                    for result in response.tool_results:
                        if result.tool_name == "finish":
                            continue  # Skip - summary shown separately
                        if result.success:
                            self._add_tool_result(
                                result.tool_name, result.result or "Done"
                            )
                        else:
                            self._add_tool_result(
                                result.tool_name, f"Error: {result.error}"
                            )

            self._set_status("idle", "assist")

        except asyncio.CancelledError:
            self._add_system("[!] Cancelled")
            self._set_status("idle", "assist")
        except Exception as e:
            self._add_system(f"[!] Error: {e}")
            self._set_status("error")
        finally:
            self._is_running = False

    @work(thread=False)
    async def _run_agent_mode(self, task: str) -> None:
        """Run in agent mode - autonomous until task complete or user stops"""
        if not self.agent:
            self._add_system("[!] Agent not ready")
            return

        self._is_running = True
        self._should_stop = False

        self._set_status("thinking", "agent")

        try:
            async for response in self.agent.agent_loop(task):
                if self._should_stop:
                    self._add_system("[!] Stopped by user")
                    break

                self._set_status("processing")

                # Show thinking/plan FIRST if there's content with tool calls
                if response.content:
                    content = response.content.strip()
                    # If it has tool calls, it's thinking.
                    # If it's marked as intermediate, it's thinking.
                    if response.tool_calls or response.metadata.get("intermediate"):
                        self._add_thinking(content)
                    else:
                        # Check if this is a task completion message
                        if response.metadata.get("task_complete"):
                            self._add_assistant(content)
                        else:
                            self._add_assistant(content)

                # Show tool calls AFTER thinking
                if response.tool_calls:
                    for call in response.tool_calls:
                        # Show all tools including finish
                        args_str = str(call.arguments)
                        self._add_tool(call.name, args_str)

                # Show tool results
                if response.tool_results:
                    for result in response.tool_results:
                        if result.tool_name == "finish":
                            # Skip showing result for finish tool as it's redundant with the tool call display
                            continue

                        if result.success:
                            self._add_tool_result(
                                result.tool_name, result.result or "Done"
                            )
                        else:
                            self._add_tool_result(
                                result.tool_name, f"Error: {result.error}"
                            )

                # Check state
                if self.agent.state.value == "waiting_input":
                    self._set_status("waiting")
                    self._add_system("? Awaiting input...")
                    break
                elif self.agent.state.value == "complete":
                    break

                self._set_status("thinking")

            self._set_status("complete", "agent")
            self._add_system("+ Agent task complete. Back to assist mode.")

            # Return to assist mode
            await asyncio.sleep(1)
            self._set_status("idle", "assist")

        except asyncio.CancelledError:
            self._add_system("[!] Cancelled")
            self._set_status("idle", "assist")
        except Exception as e:
            self._add_system(f"[!] Error: {e}")
            self._set_status("error")
        finally:
            self._is_running = False

    def action_quit_app(self) -> None:
        # Stop any running tasks first
        if self._is_running:
            self._should_stop = True
            if self._current_worker and not getattr(self._current_worker, "done", lambda: False)():
                cancel = getattr(self._current_worker, "cancel", None)
                if cancel:
                    cancel()
            if self._current_crew:
                # Schedule cancel but don't wait - we're exiting
                asyncio.create_task(self._cancel_crew())
        self.exit()

    def action_stop_agent(self) -> None:
        if self._is_running:
            self._should_stop = True
            self._add_system("[!] Stopping...")

            # Cancel the running worker to interrupt blocking awaits
            if self._current_worker and not getattr(self._current_worker, "done", lambda: False)():
                cancel = getattr(self._current_worker, "cancel", None)
                if cancel:
                    cancel()

            # Cancel crew orchestrator if running
            if self._current_crew:
                asyncio.create_task(self._cancel_crew())

            # Clean up agent state to prevent stale tool responses
            if self.agent:
                self.agent.cleanup_after_cancel()

            # Reconnect MCP servers (they may be in a bad state after cancellation)
            if self.mcp_manager:
                asyncio.create_task(self._reconnect_mcp_after_cancel())

    async def _cancel_crew(self) -> None:
        """Cancel crew orchestrator and all workers."""
        try:
            if self._current_crew:
                await self._current_crew.cancel()
                self._current_crew = None
                # Mark all running workers as cancelled in the UI
                for worker_id, worker in self._crew_workers.items():
                    if worker.get("status") in ("running", "pending"):
                        self._update_crew_worker(worker_id, status="cancelled")
        except Exception:
            pass  # Best effort

    async def _reconnect_mcp_after_cancel(self) -> None:
        """Reconnect MCP servers after cancellation to restore clean state."""
        await asyncio.sleep(0.5)  # Brief delay for cancellation to propagate
        try:
            if self.mcp_manager:
                await self.mcp_manager.reconnect_all()
        except Exception:
            pass  # Best effort - don't crash if reconnect fails

    def action_show_help(self) -> None:
        self.push_screen(HelpScreen())

    # ----- History navigation -----

    def action_history_up(self) -> None:
        """Recall previous input into the chat field."""
        try:
            inp = self.query_one("#chat-input", Input)
        except Exception:
            return

        if not self._cmd_history:
            return

        # Move back but not below zero
        if self._history_index > 0:
            self._history_index -= 1

        inp.value = self._cmd_history[self._history_index]

    def action_history_down(self) -> None:
        """Recall next input (or clear when at end)."""
        try:
            inp = self.query_one("#chat-input", Input)
        except Exception:
            return

        if not self._cmd_history:
            return

        if self._history_index < len(self._cmd_history) - 1:
            self._history_index += 1
            inp.value = self._cmd_history[self._history_index]
        else:
            # Past the end: clear input and set index to end
            self._history_index = len(self._cmd_history)
            inp.value = ""

    async def on_unmount(self) -> None:
        """Cleanup"""
        if self.mcp_manager:
            try:
                await self.mcp_manager.disconnect_all()
                await asyncio.sleep(0.1)
            except Exception:
                pass

        if self.runtime:
            try:
                await self.runtime.stop()
            except Exception:
                pass


# ----- Entry Point -----


def run_tui(
    target: Optional[str] = None,
    model: Optional[str] = None,
    use_docker: bool = False,
):
    """Run the PentestAgent TUI"""
    app = PentestAgentTUI(
        target=target,
        model=model,
        use_docker=use_docker,
    )
    app.run()


if __name__ == "__main__":
    run_tui()