/pentestgpt/interface/tui.py

"""Main TUI application for PentestGPT using Textual framework."""

import asyncio
import threading
from pathlib import Path
from typing import Any, ClassVar

from rich.text import Text
from textual import events
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Grid, Vertical, VerticalScroll
from textual.reactive import reactive
from textual.screen import ModalScreen
from textual.widgets import Button, Input, Label, Static

from pentestgpt.core.controller import AgentController
from pentestgpt.core.events import Event, EventBus, EventType
from pentestgpt.interface.components.activity_feed import ActivityFeed
from pentestgpt.interface.components.splash import SplashScreen


class HelpScreen(ModalScreen[None]):
    """Modal screen showing keyboard shortcuts and help."""

    def compose(self) -> ComposeResult:
        """Create the help dialog."""
        yield Grid(
            Label("🚩 PentestGPT CTF Solver - Help", id="help_title"),
            Label(
                "F1        Help\n"
                "Ctrl+P    Pause/Resume agent\n"
                "Ctrl+Q    Quit\n"
                "Ctrl+C    Quit\n"
                "↑/↓       Scroll feed\n"
                "PgUp/PgDn Fast scroll\n"
                "Enter     Send instruction (when paused)",
                id="help_content",
            ),
            id="help_dialog",
        )

    def on_key(self, _event: events.Key) -> None:
        """Close help on any key press."""
        self.app.pop_screen()


class QuitScreen(ModalScreen[None]):
    """Modal screen for quit confirmation."""

    def compose(self) -> ComposeResult:
        """Create the quit confirmation dialog."""
        yield Grid(
            Label("🚩 Quit PentestGPT CTF Solver?", id="quit_title"),
            Grid(
                Button("Yes", variant="error", id="btn_quit_confirm"),
                Button("No", variant="default", id="btn_quit_cancel"),
                id="quit_buttons",
            ),
            id="quit_dialog",
        )

    def on_mount(self) -> None:
        """Focus cancel button by default."""
        cancel_button = self.query_one("#btn_quit_cancel", Button)
        cancel_button.focus()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle button clicks."""
        if event.button.id == "btn_quit_confirm":
            self.app.exit()
        else:
            self.app.pop_screen()


class PentestGPTApp(App[None]):
    """Main TUI application for PentestGPT CTF Solver."""

    CSS_PATH = "styles.tcss"
    TITLE = "PentestGPT - CTF Challenge Solver"

    show_splash: reactive[bool] = reactive(default=True)
    agent_state: reactive[str] = reactive(default="idle")

    BINDINGS: ClassVar[list[Binding | tuple[str, str] | tuple[str, str, str]]] = [
        Binding("f1", "toggle_help", "Help", priority=True),
        Binding("ctrl+p", "toggle_pause", "Pause/Resume", priority=True),
        Binding("ctrl+q", "request_quit", "Quit", priority=True),
        Binding("ctrl+c", "request_quit", "Quit", priority=True),
    ]

    def __init__(
        self,
        target: str,
        custom_instruction: str | None = None,
        model: str | None = None,
        debug: bool = False,
        resume_session: str | None = None,
    ) -> None:
        """Initialize the TUI app."""
        super().__init__()
        self.target = target
        self.custom_instruction = custom_instruction
        self.model = model
        self._debug = debug
        self.resume_session = resume_session

        self._agent_thread: threading.Thread | None = None
        self._agent_stop_event = threading.Event()
        self._activity_feed: ActivityFeed | None = None
        self._controller: AgentController | None = None
        self._events: EventBus | None = None  # Initialize in on_mount

    def _setup_event_handlers(self) -> None:
        """Subscribe to agent events for UI updates."""
        if self._events is None:
            return
        self._events.subscribe(EventType.STATE_CHANGED, self._on_state_change)
        self._events.subscribe(EventType.MESSAGE, self._on_agent_message)
        self._events.subscribe(EventType.FLAG_FOUND, self._on_flag)
        self._events.subscribe(EventType.TOOL, self._on_tool)

    def _on_state_change(self, event: Event) -> None:
        """Handle agent state changes."""
        state = event.data.get("state", "unknown")
        details = event.data.get("details", "")

        # Update reactive state (triggers watch_agent_state)
        self.call_from_thread(setattr, self, "agent_state", state)

        # Log state change
        if self._activity_feed and details:
            self.call_from_thread(
                self._activity_feed.add_message,
                f"Agent: {details}",
                "info",
            )

    def _on_agent_message(self, event: Event) -> None:
        """Handle agent messages."""
        if not self._activity_feed:
            return
        text = event.data.get("text", "")
        msg_type = event.data.get("type", "info")
        if text:
            self.call_from_thread(
                self._activity_feed.add_message,
                text,
                msg_type,
            )

    def _on_flag(self, event: Event) -> None:
        """Handle flag found events."""
        if not self._activity_feed:
            return
        flag = event.data.get("flag", "")
        if flag:
            self.call_from_thread(
                self._activity_feed.add_message,
                f"🚩 FLAG FOUND: {flag}",
                "success",
            )

    def _on_tool(self, event: Event) -> None:
        """Handle tool events."""
        if not self._activity_feed:
            return
        status = event.data.get("status", "")
        name = event.data.get("name", "")
        args = event.data.get("args", {})
        result = event.data.get("result")

        if status == "start":
            self.call_from_thread(
                self._activity_feed.add_tool_execution,
                name,
                args,
                "running",
                None,
            )
        elif status == "complete":
            self.call_from_thread(
                self._activity_feed.add_tool_execution,
                name,
                args,
                "completed",
                result,
            )

    def watch_agent_state(self, state: str) -> None:
        """React to agent state changes - show/hide input field."""
        try:
            user_input = self.query_one("#user_input", Input)
            # Show input when paused, hide otherwise
            user_input.display = state == "paused"
            if state == "paused":
                user_input.focus()
        except Exception:
            pass

        # Update status bar
        self._update_status_bar()

    def compose(self) -> ComposeResult:
        """Create the initial UI layout."""
        if self.show_splash:
            yield SplashScreen(id="splash_screen")

    def watch_show_splash(self, show_splash: bool) -> None:
        """React to splash screen visibility changes."""
        if not show_splash and hasattr(self, "_driver") and self._driver is not None:
            # Remove splash screen
            try:
                splash = self.query_one("#splash_screen")
                splash.remove()
            except Exception:
                pass

            # Build main interface
            self._build_main_interface()

    def _build_main_interface(self) -> None:
        """Build the main application interface."""
        # Main container
        main_container = Vertical(id="main_container")
        self.mount(main_container)

        # Header
        header = self._create_header()
        main_container.mount(header)

        # Activity feed
        feed = ActivityFeed(id="activity_feed")
        self._activity_feed = feed

        content_area = VerticalScroll(feed, id="content_area")
        main_container.mount(content_area)

        # User input field (hidden by default, shown when paused)
        user_input = Input(
            placeholder="Type instruction and press Enter...",
            id="user_input",
        )
        user_input.display = False
        main_container.mount(user_input)

        # Status bar
        status_bar = self._create_status_bar()
        main_container.mount(status_bar)

    def _create_header(self) -> Static:
        """Create the header bar."""
        header_text = Text()
        header_text.append("🚩 ", style="bold #6366f1")
        header_text.append("PentestGPT", style="bold #6366f1")
        header_text.append(" CTF Solver", style="bold #6366f1")
        header_text.append(" v1.0", style="dim")
        header_text.append("  │  ", style="dim")
        header_text.append("Target: ", style="dim")
        header_text.append(self.target, style="bold")

        if self.model:
            header_text.append("  │  ", style="dim")
            header_text.append(f"Model: {self.model}", style="dim")

        header = Static(header_text, id="header")
        return header

    def _create_status_bar(self) -> Static:
        """Create the status bar."""
        status_text = self._build_status_text()
        status_bar = Static(status_text, id="status_bar")
        return status_bar

    def _build_status_text(self) -> Text:
        """Build status bar text based on current state."""
        status_text = Text()

        # Agent state
        status_text.append("Agent: ", style="dim")
        state_styles = {
            "idle": ("Idle", "dim"),
            "running": ("Running", "bold #6366f1"),
            "paused": ("PAUSED", "bold #f59e0b"),
            "completed": ("Completed", "bold #10b981"),
            "error": ("Error", "bold #ef4444"),
        }
        state_label, state_style = state_styles.get(self.agent_state, ("Unknown", "dim"))
        status_text.append(state_label, style=state_style)

        if self._debug:
            from pentestgpt.core.agent import DEBUG_LOG

            status_text.append("  │  ", style="dim")
            status_text.append("Debug: ", style="dim")
            status_text.append(str(DEBUG_LOG), style="bold #f59e0b")

        status_text.append("  │  ", style="dim")
        status_text.append("Ctrl+P", style="bold #8b5cf6")
        status_text.append(" Pause  ", style="dim")
        status_text.append("F1", style="bold #8b5cf6")
        status_text.append(" Help  ", style="dim")
        status_text.append("Ctrl+Q", style="bold #8b5cf6")
        status_text.append(" Quit", style="dim")

        return status_text

    def _update_status_bar(self) -> None:
        """Update status bar with current state."""
        try:
            status_bar = self.query_one("#status_bar", Static)
            status_bar.update(self._build_status_text())
        except Exception:
            pass

    def on_mount(self) -> None:
        """Handle app mount - show splash then start agent."""
        # Initialize event bus after app is mounted
        self._events = EventBus.get()
        self._setup_event_handlers()

        # Show splash for 2.5 seconds
        self.set_timer(2.5, self._hide_splash)

    def _hide_splash(self) -> None:
        """Hide splash screen and start main app."""
        self.show_splash = False
        # Start agent after splash
        self.call_later(self._start_agent)

    def _start_agent(self) -> None:
        """Start the agent in a background thread using AgentController."""
        from pentestgpt.core.config import load_config

        if self._activity_feed:
            self._activity_feed.add_message(
                "🚩 CTF Solver initializing...",
                message_type="info",
            )

        def run_agent() -> None:
            """Run the agent with controller."""
            try:
                # Build config
                config_kwargs: dict[str, Any] = {
                    "target": self.target,
                    "working_directory": Path("/workspace"),
                }
                if self.custom_instruction:
                    config_kwargs["custom_instruction"] = self.custom_instruction
                if self.model:
                    config_kwargs["llm_model"] = self.model

                config = load_config(**config_kwargs)

                # Create controller
                self._controller = AgentController(config)

                # Build task
                task = f"Solve this CTF challenge and capture the flag(s): {self.target}"
                if self.custom_instruction:
                    task += f"\n\nChallenge context: {self.custom_instruction}"

                # Run agent with controller
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                result = loop.run_until_complete(
                    self._controller.run(task, resume_session_id=self.resume_session)
                )
                loop.close()

                # Show final result summary
                if self._activity_feed:
                    if result.get("success"):
                        flags = result.get("flags_found", [])
                        cost = result.get("cost_usd", 0)
                        session_id = result.get("session_id", "")
                        if flags:
                            self.call_from_thread(
                                self._activity_feed.add_message,
                                f"🚩 Challenge complete! {len(flags)} flag(s) | Cost: ${cost:.4f} | Session: {session_id}",
                                "success",
                            )
                        else:
                            self.call_from_thread(
                                self._activity_feed.add_message,
                                f"⚠ Challenge ended without flags | Cost: ${cost:.4f} | Session: {session_id}",
                                "warning",
                            )
                    else:
                        error = result.get("error", "Unknown error")
                        self.call_from_thread(
                            self._activity_feed.add_message,
                            f"✗ Agent error: {error}",
                            "error",
                        )

            except Exception as e:
                if self._activity_feed:
                    self.call_from_thread(
                        self._activity_feed.add_message,
                        f"✗ Agent error: {e!s}",
                        "error",
                    )

        self._agent_thread = threading.Thread(target=run_agent, daemon=True)
        self._agent_thread.start()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        """Handle user input submission."""
        if event.input.id == "user_input":
            text = event.value.strip()
            if text and self._events:
                # Emit user input event
                self._events.emit_input(text)
                if self._activity_feed:
                    self._activity_feed.add_message(
                        f"📝 Instruction queued: {text[:50]}...",
                        message_type="info",
                    )
            event.input.value = ""

    def action_toggle_pause(self) -> None:
        """Toggle pause/resume state (Ctrl+P)."""
        if self.show_splash or not self._controller:
            return

        if self.agent_state == "paused":
            self._controller.resume()
        elif self.agent_state == "running":
            self._controller.pause()

    def action_toggle_help(self) -> None:
        """Show/hide help screen."""
        if self.show_splash:
            return

        if isinstance(self.screen, HelpScreen):
            self.pop_screen()
        else:
            self.push_screen(HelpScreen())

    def action_request_quit(self) -> None:
        """Show quit confirmation."""
        if self.show_splash:
            self.exit()
            return

        self.push_screen(QuitScreen())

    def on_unmount(self) -> None:
        """Cleanup when app closes."""
        # Signal agent to stop
        self._agent_stop_event.set()

        # Wait for agent thread
        if self._agent_thread and self._agent_thread.is_alive():
            self._agent_thread.join(timeout=2.0)


async def run_tui(
    target: str,
    custom_instruction: str | None = None,
    model: str | None = None,
    debug: bool = False,
    resume_session: str | None = None,
) -> None:
    """Run the TUI application.

    Args:
        target: Target URL/IP/domain
        custom_instruction: Optional challenge context
        model: Optional model override
        debug: Enable debug mode
        resume_session: Optional session ID to resume
    """
    app = PentestGPTApp(
        target=target,
        custom_instruction=custom_instruction,
        model=model,
        debug=debug,
        resume_session=resume_session,
    )
    await app.run_async()


# For testing
if __name__ == "__main__":
    import asyncio

    asyncio.run(run_tui("example.com"))