#!/usr/bin/env python3
"""
MundiX CLI - AI-Powered Cybersecurity Copilot for Kali Linux
An intelligent terminal assistant that connects to DeepHat/WhiteRabbitNeo models
via HuggingFace API, with collective memory, project workspaces, and pentest playbooks.

Usage:
    mundix                     # Start interactive mode
    mundix "scan 10.0.0.1"     # One-shot query
    mundix --model wrn-v3-7b   # Use specific model
    mundix --autopilot 3       # Set interactivity level (1-5)
    mundix --project mytest    # Open/create a project workspace
    mundix --update            # Check for updates
"""

import sys
import os
import argparse
import re
import time
import shutil
import subprocess as sp

from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.text import Text
from rich.theme import Theme
from rich.live import Live
from rich.table import Table
from rich.syntax import Syntax
from rich.columns import Columns
from rich.rule import Rule
from prompt_toolkit import PromptSession
from prompt_toolkit.history import FileHistory
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.formatted_text import ANSI
from prompt_toolkit.completion import Completer, Completion

from ai_provider import AIProvider, MODELS, DEFAULT_MODEL, FALLBACK_ORDER, auto_register_device, request_device_activation, check_activation, check_credits, load_config, save_config, get_device_info, get_device_status
from memory import Memory
from system_integrator import SystemIntegrator
from autopilot import AutoPilot, LEVELS, diagnose_error
from project import Project
from playbooks import PlaybookRunner, PLAYBOOKS

# ─── Version ──────────────────────────────────────────────────────────
VERSION = "0.7.1"
REPO_URL = "https://github.com/k0k4/mundix-cli"
RAW_URL = "https://raw.githubusercontent.com/k0k4/mundix-cli/main"

# ─── Theme ────────────────────────────────────────────────────────────
MUNDIX_THEME = Theme({
    "info": "cyan",
    "warning": "yellow",
    "danger": "bold red",
    "success": "bold green",
    "mundix": "bold red",
    "command": "bold yellow",
    "output": "dim white",
    "hacker": "green",
    "dim": "dim white",
    "fix": "bold cyan",
    "project": "bold magenta",
})

console = Console(theme=MUNDIX_THEME)

# ─── Slash Command Completer ─────────────────────────────────────────

SLASH_COMMANDS = {
    "/help": "Show help message",
    "/models": "List available AI models",
    "/model": "Switch model: /model <key>",
    "/autopilot": "Set AutoPilot level: /autopilot <1-5>",
    "/project": "Manage projects",
    "/project new": "Create new project: /project new <name>",
    "/project open": "Open existing project: /project open <name>",
    "/project info": "Show current project details",
    "/project target": "Set target: /project target <ip>",
    "/project type": "Set type: /project type <web|internal|app>",
    "/project close": "Close current project",
    "/projects": "List all projects",
    "/report": "Generate pentest report",
    "/note": "Add note: /note <text>",
    "/playbook": "List or start playbook: /playbook [web|internal|app]",
    "/playbook web": "Start Web Application Pentest playbook",
    "/playbook internal": "Start Internal Network Pentest playbook",
    "/playbook app": "Start Application/API Pentest playbook",
    "/next": "Execute next playbook command",
    "/skip": "Skip current playbook command",
    "/phase": "Show current playbook phase",
    "/progress": "Show playbook progress",
    "/memory": "Show memory statistics",
    "/sync": "Sync collective memory",
    "/clear": "Clear conversation history",
    "/system": "Show system information",
    "/history": "Show command execution history",
    "/exec": "Execute command: /exec <cmd>",
    "/analyze": "Analyze last command output with AI",
    "/activate": "Activate MundiX CLI with your account",
    "/credits": "Check remaining credits",
    "/update": "Check for updates",
    "/exit": "Exit MundiX CLI",
    "/quit": "Exit MundiX CLI",
}


class MundixCompleter(Completer):
    """Custom completer for MundiX slash commands.
    Activates when the user types '/' and provides intelligent
    completion with descriptions for each command."""

    def __init__(self, cli_instance=None):
        self.cli = cli_instance

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor

        # Only complete if text starts with '/'
        if not text.startswith("/"):
            return

        text_lower = text.lower()

        # Dynamic completions for specific sub-commands
        parts = text.split()

        # /model <key> - complete model keys
        if len(parts) == 2 and parts[0].lower() == "/model":
            prefix = parts[1].lower()
            for key in MODELS:
                if key.startswith(prefix):
                    yield Completion(
                        key,
                        start_position=-len(parts[1]),
                        display_meta=MODELS[key]["name"],
                    )
            return

        # /autopilot <1-5> - complete levels
        if len(parts) == 2 and parts[0].lower() == "/autopilot":
            prefix = parts[1]
            for lvl, info in LEVELS.items():
                lvl_str = str(lvl)
                if lvl_str.startswith(prefix):
                    yield Completion(
                        lvl_str,
                        start_position=-len(parts[1]),
                        display_meta=info["name"],
                    )
            return

        # /project type <type> - complete types
        if len(parts) == 3 and parts[0].lower() == "/project" and parts[1].lower() == "type":
            prefix = parts[2].lower()
            for t in ["web", "internal", "app"]:
                if t.startswith(prefix):
                    yield Completion(t, start_position=-len(parts[2]))
            return

        # /project open <name> - complete existing project names
        if len(parts) >= 2 and parts[0].lower() == "/project" and parts[1].lower() == "open":
            prefix = parts[2].lower() if len(parts) > 2 else ""
            try:
                projects = Project.list_projects()
                for p in projects:
                    if p["name"].lower().startswith(prefix):
                        yield Completion(
                            p["name"],
                            start_position=-len(prefix),
                            display_meta=p.get("target", ""),
                        )
            except Exception:
                pass
            return

        # /playbook <type> - complete playbook types
        if len(parts) == 2 and parts[0].lower() == "/playbook":
            prefix = parts[1].lower()
            for pb_id, pb in PLAYBOOKS.items():
                if pb_id.startswith(prefix):
                    yield Completion(
                        pb_id,
                        start_position=-len(parts[1]),
                        display_meta=pb["name"],
                    )
            return

        # General slash command completion
        for cmd, desc in SLASH_COMMANDS.items():
            if cmd.startswith(text_lower):
                yield Completion(
                    cmd,
                    start_position=-len(text),
                    display_meta=desc,
                )


# ─── Banner ──────────────────────────────────────────────────────────
BANNER_LINES = [
    "  ███╗   ███╗██╗   ██╗███╗   ██╗██████╗ ██╗██╗  ██╗",
    "  ████╗ ████║██║   ██║████╗  ██║██╔══██╗██║╚██╗██╔╝",
    "  ██╔████╔██║██║   ██║██╔██╗ ██║██║  ██║██║ ╚███╔╝ ",
    "  ██║╚██╔╝██║██║   ██║██║╚██╗██║██║  ██║██║ ██╔██╗ ",
    "  ██║ ╚═╝ ██║╚██████╔╝██║ ╚████║██████╔╝██║██╔╝ ██╗",
    "  ╚═╝     ╚═╝ ╚═════╝ ╚═╝  ╚═══╝╚═════╝ ╚═╝╚═╝  ╚═╝",
]

TAGLINE = "  [ AI-Powered Cybersecurity Copilot ]"

HELP_TEXT = """
[bold red]MundiX CLI Commands:[/bold red]

  [bold green]/help[/bold green]                Show this help message
  [bold green]/models[/bold green]              List available AI models
  [bold green]/model <key>[/bold green]         Switch to a different model
  [bold green]/autopilot[/bold green]           Show current AutoPilot level
  [bold green]/autopilot <1-5>[/bold green]     Set AutoPilot level

[bold red]Project / Client:[/bold red]
  [bold green]/project new <name>[/bold green]  Create a new project workspace
  [bold green]/project open <name>[/bold green] Open an existing project
  [bold green]/project info[/bold green]        Show current project details
  [bold green]/project target <ip>[/bold green] Set the target for current project
  [bold green]/project type <type>[/bold green] Set project type (web/internal/app)
  [bold green]/project close[/bold green]       Close the current project
  [bold green]/projects[/bold green]            List all projects
  [bold green]/report[/bold green]              Generate pentest report for current project
  [bold green]/note <text>[/bold green]         Add a note to the project

[bold red]Playbooks (Automated Pentest):[/bold red]
  [bold green]/playbook[/bold green]            List available playbooks
  [bold green]/playbook <type>[/bold green]     Start a playbook (web/internal/app)
  [bold green]/next[/bold green]                Execute next playbook command
  [bold green]/skip[/bold green]                Skip current playbook command
  [bold green]/phase[/bold green]               Show current playbook phase
  [bold green]/progress[/bold green]            Show playbook progress

[bold red]Memory & System:[/bold red]
  [bold green]/memory[/bold green]              Show memory statistics
  [bold green]/sync[/bold green]                Sync collective memory
  [bold green]/clear[/bold green]               Clear conversation history
  [bold green]/system[/bold green]              Show system information
  [bold green]/history[/bold green]             Show command execution history
  [bold green]/exec <cmd>[/bold green]          Execute a command directly
  [bold green]/analyze[/bold green]             Send last command output to AI for analysis
  [bold green]/update[/bold green]              Check for updates
  [bold green]/exit[/bold green]                Exit MundiX CLI

[bold red]AutoPilot Levels:[/bold red]
  [cyan]1 MANUAL[/cyan]      Ask before every action
  [green]2 GUIDED[/green]      Auto-execute safe commands (default)
  [yellow]3 BALANCED[/yellow]    Auto-execute all, ask only dangerous
  [red]4 AGGRESSIVE[/red]  Auto-execute everything, minimal prompts
  [bold red]5 AUTOPILOT[/bold red]   Full autonomy. YOLO mode.

[dim]Type ? for help | / for commands | Just type naturally to hack.[/dim]
"""


def clear_screen():
    """Clear the terminal screen."""
    os.system("clear" if os.name != "nt" else "cls")


def get_terminal_width():
    """Get terminal width."""
    return shutil.get_terminal_size((80, 24)).columns


def typing_effect(text, delay=0.008, style="bold red"):
    """Print text with a hacker typing effect."""
    for char in text:
        console.print(char, end="", style=style, highlight=False)
        time.sleep(delay)
    print()


def animated_banner():
    """Display animated startup banner with hacker aesthetic."""
    clear_screen()
    import random
    glitch_chars = "█▓▒░╔╗╚╝║═╬╣╠╩╦"

    for line in BANNER_LINES:
        glitched = ""
        for ch in line:
            if ch == " ":
                glitched += " "
            else:
                glitched += random.choice(glitch_chars)
        console.print(glitched, style="dim red", highlight=False)
        time.sleep(0.03)

    sys.stdout.write(f"\033[{len(BANNER_LINES)}A")
    sys.stdout.flush()
    time.sleep(0.15)

    for line in BANNER_LINES:
        console.print(line, style="bold red", highlight=False)
        time.sleep(0.04)

    print()
    typing_effect(TAGLINE, delay=0.015, style="bold green")
    print()
    time.sleep(0.1)


def build_prompt_prefix(cwd: str, model_name: str = "", mem_entries: int = 0,
                        ap_level: int = 2, ap_name: str = "GUIDED",
                        project_name: str = None, project_target: str = None):
    """Build the prompt with Copilot-style dividers and hacker identity."""
    red = "\033[1;31m"
    green = "\033[1;32m"
    yellow = "\033[1;33m"
    cyan = "\033[1;36m"
    magenta = "\033[1;35m"
    dim = "\033[2m"
    dim_red = "\033[2;31m"
    reset = "\033[0m"
    bold = "\033[1m"

    # AutoPilot color based on level
    ap_colors = {1: cyan, 2: green, 3: yellow, 4: red, 5: f"\033[1;31m"}
    ap_color = ap_colors.get(ap_level, green)

    width = get_terminal_width()
    short_cwd = cwd.replace(os.path.expanduser("~"), "~")

    # Top info line: CWD left, model+AP+mem right
    left = f" {short_cwd}"

    # Add project info if active
    if project_name:
        target_str = f" -> {project_target}" if project_target else ""
        right = f"[{project_name}{target_str}] | {model_name} | AP:{ap_level} {ap_name} | mem:{mem_entries} "
    else:
        right = f"{model_name} | AP:{ap_level} {ap_name} | mem:{mem_entries} "

    padding = width - len(left) - len(right)
    if padding < 1:
        padding = 1

    # Color the project name in magenta if present
    if project_name:
        target_str = f" -> {project_target}" if project_target else ""
        right_colored = f"{magenta}[{project_name}{target_str}]{reset}{dim} | {model_name} | AP:{ap_level} {ap_name} | mem:{mem_entries} {reset}"
    else:
        right_colored = f"{dim}{right}{reset}"

    info_line = f"{dim}{left}{reset}{' ' * padding}{right_colored}"

    # Divider
    divider = f"{dim_red}{'─' * width}{reset}"

    prompt = f"{info_line}\n{divider}\n{red}❯{reset} "
    return prompt


def print_bottom_bar(ap_level: int = 2, has_project: bool = False, has_playbook: bool = False):
    """Print the bottom divider + hints bar after prompt submission."""
    width = get_terminal_width()
    dim_red = "\033[2;31m"
    dim = "\033[2m"
    reset = "\033[0m"
    divider = f"{dim_red}{'─' * width}{reset}"

    hints = ["ctrl+c exit"]
    if has_playbook:
        hints.append("/next step")
    if has_project:
        hints.append("/report gen")
    hints.extend(["/autopilot <1-5>", "/models switch", "? help"])

    hint = f" {dim}{'  │  '.join(hints)}{reset}"
    print(divider)
    print(hint)


def check_for_updates():
    """Check GitHub for newer version."""
    import requests
    console.print("[info]Checking for updates...[/info]")
    try:
        resp = requests.get(
            f"{RAW_URL}/mundix_cli.py",
            timeout=10,
            headers={"Cache-Control": "no-cache"},
        )
        if resp.status_code == 200:
            for line in resp.text.split("\n"):
                if line.strip().startswith("VERSION"):
                    remote_version = line.split('"')[1]
                    from packaging.version import Version
                    try:
                        is_newer = Version(remote_version) > Version(VERSION)
                    except Exception:
                        is_newer = remote_version != VERSION
                    if is_newer:
                        console.print(f"[success]New version available: v{remote_version} (current: v{VERSION})[/success]")
                        console.print(f"[info]Update with:[/info]")
                        console.print(f"[command]  curl -sSL {RAW_URL}/install.sh | bash[/command]")
                        return True
                    else:
                        console.print(f"[success]You are running the latest version (v{VERSION})[/success]")
                        return False
        console.print("[warning]Could not check for updates.[/warning]")
        return False
    except Exception as e:
        console.print(f"[warning]Update check failed: {e}[/warning]")
        return False


class MundixCLI:
    """Main CLI application class."""

    def __init__(self, model_key: str = None,
                 autopilot_level: int = None, project_name: str = None):
        self.ai = AIProvider(model_key=model_key)
        self.mode = self.ai.mode  # cloud
        self.memory = Memory()
        self.system = SystemIntegrator()
        self.autopilot = AutoPilot()
        self.project = Project(name=project_name)  # Load active or create new
        self.playbook = PlaybookRunner()
        self.running = True

        # Override autopilot level if passed via CLI
        if autopilot_level is not None:
            self.autopilot.set_level(autopilot_level)

        # Prompt session with history and autocomplete
        history_path = os.path.expanduser("~/.mundix/prompt_history")
        os.makedirs(os.path.dirname(history_path), exist_ok=True)
        self.completer = MundixCompleter(cli_instance=self)
        self.session = PromptSession(
            history=FileHistory(history_path),
            auto_suggest=AutoSuggestFromHistory(),
            completer=self.completer,
            complete_while_typing=False,  # Only complete on Tab
        )

    def show_startup(self):
        """Display the animated startup sequence."""
        animated_banner()

        # Show device status and free queries
        config = load_config()
        status = config.get("status", "anonymous")
        total_queries = config.get("total_queries", 0)
        max_free = config.get("max_free_queries", 3)
        credits = config.get("credits", 0)

        if status == "anonymous":
            remaining = max(0, max_free - total_queries)
            if remaining > 0:
                console.print(f"  [dim]Free queries remaining: [bold green]{remaining}[/bold green] of {max_free}  |  Register for +5 more: [cyan]mundix --activate[/cyan][/dim]")
            else:
                console.print(f"  [warning]Free queries used. Register to continue: mundix --activate[/warning]")
        elif status == "registered":
            if credits > 0:
                est_queries = credits // 2000
                console.print(f"  [dim]Credits: {credits:,}  |  ~{est_queries} queries left  |  Top up: [cyan]chat.mundix.dev/billing[/cyan][/dim]")
            else:
                console.print(f"  [warning]Credits exhausted. Top up at: chat.mundix.dev/billing[/warning]")
        elif status == "paid":
            est_queries = credits // 2000
            console.print(f"  [dim]Credits: {credits:,}  |  ~{est_queries} queries left[/dim]")

        # Show active project if any
        if self.project.name:
            info = self.project.get_info()
            console.print(f"  [project]Active Project:[/project] {info['name']}", highlight=False)
            if info.get("target"):
                console.print(f"  [dim]Target: {info['target']}  |  Hosts: {info['hosts']}  |  Vulns: {info['vulns']}  |  Context: {info['context_entries']} entries[/dim]")
        else:
            console.print(f"  [dim]No active project. Use /project new <name> to start one.[/dim]")

        console.print()

    # ─── Slash Commands ─────────────────────────────────────────────

    def handle_slash_command(self, command: str) -> bool:
        """Handle slash commands. Returns True if command was handled."""
        parts = command.strip().split(maxsplit=1)
        cmd = parts[0].lower()
        args = parts[1] if len(parts) > 1 else ""

        if cmd in ("/help", "/?", "?"):
            console.print(HELP_TEXT)
            return True

        elif cmd == "/autopilot":
            if not args:
                table = Table(border_style="red", title="[bold red]AutoPilot Levels[/bold red]")
                table.add_column("Level", style="bold", justify="center")
                table.add_column("Name", style="bold")
                table.add_column("Description", style="dim")
                table.add_column("Auto-Exec", justify="center")
                table.add_column("Auto-Fix", justify="center")
                table.add_column("Auto-Install", justify="center")
                table.add_column("Status", justify="center")
                for lvl, info in LEVELS.items():
                    is_active = lvl == self.autopilot.level
                    status = f"[bold green]ACTIVE[/bold green]" if is_active else ""
                    ae = "[green]YES[/green]" if info["auto_execute"] else "[red]NO[/red]"
                    af = "[green]YES[/green]" if info["auto_fix"] else "[red]NO[/red]"
                    ai = "[green]YES[/green]" if info["auto_install"] else "[red]NO[/red]"
                    table.add_row(
                        str(lvl), f"[{info['color']}]{info['name']}[/{info['color']}]",
                        info["description"], ae, af, ai, status
                    )
                console.print(table)
                return True
            else:
                try:
                    level = int(args.strip())
                    if self.autopilot.set_level(level):
                        s = self.autopilot.settings
                        console.print(f"[success]AutoPilot set to level {level}: [{s['color']}]{s['name']}[/{s['color']}][/success]")
                        console.print(f"[dim]  {s['description']}[/dim]")
                    else:
                        console.print("[danger]Invalid level. Use 1-5.[/danger]")
                except ValueError:
                    console.print("[danger]Usage: /autopilot <1-5>[/danger]")
                return True

        elif cmd == "/models":
            table = Table(border_style="red", title="[bold red]Available Models[/bold red]")
            table.add_column("Key", style="bold green")
            table.add_column("Name", style="white")
            table.add_column("Description", style="dim")
            table.add_column("Status", style="green", justify="center")
            for key, info in MODELS.items():
                if key == self.ai.model_key:
                    status = "[bold green]ACTIVE[/bold green]"
                elif key in self.ai._failed_models:
                    status = "[dim red]OFFLINE[/dim red]"
                else:
                    status = "[dim]ready[/dim]"
                table.add_row(key, info["name"], info["description"], status)
            console.print(table)
            return True

        elif cmd == "/model":
            if not args:
                console.print("[warning]Usage: /model <model_key>[/warning]")
                return True
            if self.ai.set_model(args.strip()):
                console.print(f"[success]Switched to: {self.ai.model_info['name']}[/success]")
            else:
                console.print(f"[danger]Unknown model: {args}. Use /models to see options.[/danger]")
            return True

        # ── Project Commands ──
        elif cmd == "/project":
            return self._handle_project_command(args)

        elif cmd == "/projects":
            projects = Project.list_projects()
            if not projects:
                console.print("[dim]No projects found. Use /project new <name> to create one.[/dim]")
                return True
            table = Table(border_style="magenta", title="[bold magenta]Projects[/bold magenta]")
            table.add_column("Name", style="bold green")
            table.add_column("Target", style="white")
            table.add_column("Type", style="cyan")
            table.add_column("Status", style="yellow")
            table.add_column("Created", style="dim")
            for p in projects:
                active = " [bold green]<<[/bold green]" if p["name"] == self.project.name else ""
                table.add_row(
                    p["name"] + active,
                    p.get("target", ""),
                    p.get("type", ""),
                    p.get("status", ""),
                    p.get("created", "")[:10],
                )
            console.print(table)
            return True

        elif cmd == "/report":
            if not self.project.name:
                console.print("[warning]No active project. Use /project new <name> first.[/warning]")
                return True
            console.print("[info]Generating pentest report...[/info]")
            report_path = self.project.generate_report()
            console.print(f"[success]Report saved: {report_path}[/success]")
            console.print(f"[dim]View with: cat {report_path}[/dim]")
            return True

        elif cmd == "/note":
            if not args:
                console.print("[warning]Usage: /note <your note text>[/warning]")
                return True
            if not self.project.name:
                console.print("[warning]No active project. Use /project new <name> first.[/warning]")
                return True
            self.project.add_note(args)
            console.print(f"[success]Note added to project.[/success]")
            return True

        # ── Playbook Commands ──
        elif cmd == "/playbook":
            return self._handle_playbook_command(args)

        elif cmd == "/next":
            return self._playbook_next()

        elif cmd == "/skip":
            return self._playbook_skip()

        elif cmd == "/phase":
            if not self.playbook.playbook:
                console.print("[warning]No active playbook. Use /playbook <type> to start one.[/warning]")
                return True
            console.print(self.playbook.get_phase_summary())
            return True

        elif cmd == "/progress":
            if not self.playbook.playbook:
                console.print("[warning]No active playbook. Use /playbook <type> to start one.[/warning]")
                return True
            prog = self.playbook.get_progress()
            console.print(Panel(
                f"[bold]{prog['playbook']}[/bold]\n"
                f"Phase: {prog['phase']}/{prog['total_phases']} - {prog['phase_name']}\n"
                f"Commands: {prog['completed_commands']} done, {prog['skipped_commands']} skipped / {prog['total_commands']} total\n"
                f"Progress: [bold green]{prog['progress_pct']}%[/bold green]",
                title="[bold red]Playbook Progress[/bold red]",
                border_style="red",
            ))
            return True

        # ── Memory & System Commands ──
        elif cmd == "/memory":
            stats = self.memory.get_stats()
            table = Table(border_style="red", title="[bold red]Memory Statistics[/bold red]")
            table.add_column("Metric", style="bold green")
            table.add_column("Value", style="cyan")
            table.add_row("Local Entries", str(stats["local_entries"]))
            table.add_row("Collective Entries", str(stats["collective_entries"]))
            table.add_row("Total Queries", str(stats["total_queries"]))
            table.add_row("Total Commands", str(stats["total_commands"]))
            table.add_row("Last Sync", stats["last_sync"] or "Never")
            if self.project.name:
                pinfo = self.project.get_info()
                table.add_row("Project Context Entries", str(pinfo["context_entries"]))
                table.add_row("Project Findings (Hosts)", str(pinfo["hosts"]))
                table.add_row("Project Findings (Vulns)", str(pinfo["vulns"]))
            console.print(table)
            return True

        elif cmd == "/sync":
            console.print("[info]Syncing collective memory...[/info]")
            new_count = self.memory.sync_collective()
            console.print(f"[success]Sync complete. {new_count} new entries.[/success]")
            return True

        elif cmd == "/clear":
            self.ai.clear_history()
            console.print("[success]Conversation history cleared.[/success]")
            return True

        elif cmd == "/system":
            info = self.system.get_system_info()
            table = Table(border_style="red", title="[bold red]System Info[/bold red]")
            table.add_column("Property", style="bold green")
            table.add_column("Value", style="white")
            for key, value in info.items():
                table.add_row(key.replace("_", " ").title(), value)
            console.print(table)
            return True

        elif cmd == "/history":
            history = self.system.get_command_history()
            if not history:
                console.print("[dim]No commands executed yet.[/dim]")
                return True
            for i, entry in enumerate(history[-10:], 1):
                status = "[green]OK[/green]" if entry["returncode"] == 0 else f"[red]Exit {entry['returncode']}[/red]"
                console.print(f"  {status} [command]{entry['command']}[/command]")
            return True

        elif cmd == "/exec":
            if not args:
                console.print("[warning]Usage: /exec <command>[/warning]")
                return True
            self._execute_command(args)
            return True

        elif cmd == "/analyze":
            last_output = self.system.get_last_output()
            if not last_output:
                console.print("[warning]No previous command output to analyze.[/warning]")
                return True
            console.print("[info]Analyzing last output...[/info]")
            prompt = f"Analyze the following command output and provide insights, findings, and next steps:\n\n```\n{last_output[:3000]}\n```"
            self._stream_ai_response(prompt)
            return True

        elif cmd == "/activate":
            _handle_activate()
            return True

        elif cmd == "/credits":
            _handle_credits()
            return True

        elif cmd == "/update":
            check_for_updates()
            return True

        elif cmd in ("/exit", "/quit", "/q"):
            self.running = False
            console.print("\n[dim]Session terminated. Stay sharp, hacker.[/dim]\n")
            return True

        return False

    # ─── Project Sub-Commands ───────────────────────────────────────

    def _handle_project_command(self, args: str) -> bool:
        """Handle /project sub-commands."""
        parts = args.strip().split(maxsplit=1)
        subcmd = parts[0].lower() if parts else ""
        subargs = parts[1] if len(parts) > 1 else ""

        if subcmd == "new":
            if not subargs:
                console.print("[warning]Usage: /project new <name>[/warning]")
                return True
            self.project = Project(name=subargs)
            console.print(f"[success]Project created: [project]{self.project.name}[/project][/success]")
            console.print(f"[dim]Workspace: {self.project.project_dir}[/dim]")
            console.print(f"[dim]Set target with: /project target <ip_or_domain>[/dim]")
            return True

        elif subcmd == "open":
            if not subargs:
                console.print("[warning]Usage: /project open <name>[/warning]")
                return True
            self.project = Project(name=subargs)
            if self.project.project_dir and self.project.project_dir.exists():
                info = self.project.get_info()
                console.print(f"[success]Opened project: [project]{info['name']}[/project][/success]")
                if info.get("target"):
                    console.print(f"[dim]Target: {info['target']}  |  Hosts: {info['hosts']}  |  Vulns: {info['vulns']}  |  Context: {info['context_entries']} entries[/dim]")
            else:
                console.print(f"[info]New project created: {self.project.name}[/info]")
            return True

        elif subcmd == "info":
            if not self.project.name:
                console.print("[warning]No active project. Use /project new <name>[/warning]")
                return True
            info = self.project.get_info()
            table = Table(border_style="magenta", title=f"[bold magenta]Project: {info['name']}[/bold magenta]")
            table.add_column("Property", style="bold green")
            table.add_column("Value", style="white")
            table.add_row("Name", info["name"])
            table.add_row("Target", info.get("target", "Not set"))
            table.add_row("Type", info.get("type", "Not set"))
            table.add_row("Status", info.get("status", "active"))
            table.add_row("Hosts Found", str(info.get("hosts", 0)))
            table.add_row("Open Ports", str(info.get("ports", 0)))
            table.add_row("Vulnerabilities", str(info.get("vulns", 0)))
            table.add_row("Credentials", str(info.get("creds", 0)))
            table.add_row("Context Entries", str(info.get("context_entries", 0)))
            table.add_row("Timeline Events", str(info.get("timeline_events", 0)))
            table.add_row("Workspace", info.get("path", ""))
            table.add_row("Created", info.get("created", "")[:19])
            console.print(table)
            return True

        elif subcmd == "target":
            if not subargs:
                console.print("[warning]Usage: /project target <ip_or_domain>[/warning]")
                return True
            if not self.project.name:
                console.print("[warning]No active project. Use /project new <name> first.[/warning]")
                return True
            self.project.set_target(subargs.strip())
            console.print(f"[success]Target set: {subargs.strip()}[/success]")

            # Also set as playbook variable if playbook is active
            if self.playbook.playbook:
                self.playbook.set_variable("target", subargs.strip())

            return True

        elif subcmd == "type":
            if not subargs:
                console.print("[warning]Usage: /project type <web|internal|app>[/warning]")
                return True
            if not self.project.name:
                console.print("[warning]No active project. Use /project new <name> first.[/warning]")
                return True
            self.project.set_type(subargs.strip())
            console.print(f"[success]Project type set: {subargs.strip()}[/success]")
            return True

        elif subcmd == "close":
            if not self.project.name:
                console.print("[warning]No active project.[/warning]")
                return True
            self.project.close()
            console.print(f"[success]Project '{self.project.name}' marked as completed.[/success]")
            return True

        else:
            # No subcommand - show info or help
            if self.project.name:
                return self._handle_project_command("info")
            console.print("[warning]Usage: /project new|open|info|target|type|close[/warning]")
            return True

    # ─── Playbook Sub-Commands ──────────────────────────────────────

    def _handle_playbook_command(self, args: str) -> bool:
        """Handle /playbook sub-commands."""
        if not args:
            # List available playbooks
            playbooks = PlaybookRunner.list_playbooks()
            table = Table(border_style="red", title="[bold red]Available Playbooks[/bold red]")
            table.add_column("ID", style="bold green")
            table.add_column("Name", style="white")
            table.add_column("Description", style="dim")
            table.add_column("Phases", justify="center")
            table.add_column("Commands", justify="center")
            for pb in playbooks:
                table.add_row(pb["id"], f"{pb['icon']} {pb['name']}", pb["description"],
                              str(pb["phases"]), str(pb["commands"]))
            console.print(table)
            console.print("[dim]Start with: /playbook <id>  (e.g., /playbook web)[/dim]")
            return True

        pb_id = args.strip().lower()
        if pb_id not in PLAYBOOKS:
            console.print(f"[danger]Unknown playbook: {pb_id}. Available: web, internal, app[/danger]")
            return True

        # Need a project first
        if not self.project.name:
            console.print("[warning]Create a project first: /project new <name>[/warning]")
            return True

        self.playbook.set_playbook(pb_id)

        # Set target variable from project
        if self.project.metadata.get("target"):
            self.playbook.set_variable("target", self.project.metadata["target"])

        # Set project type
        self.project.set_type(pb_id)

        pb = PLAYBOOKS[pb_id]
        console.print(Panel(
            f"[bold]{pb['icon']} {pb['name']}[/bold]\n\n"
            f"{pb['description']}\n\n"
            f"Phases: {len(pb['phases'])}  |  "
            f"Target: {self.project.metadata.get('target', '[yellow]Not set - use /project target <ip>[/yellow]')}\n\n"
            f"[dim]Use /next to execute commands step by step\n"
            f"Use /skip to skip a command\n"
            f"Use /phase to see current phase details[/dim]",
            title=f"[bold red]Playbook Started[/bold red]",
            border_style="red",
        ))

        # Show first phase
        console.print()
        console.print(self.playbook.get_phase_summary())
        return True

    def _playbook_next(self) -> bool:
        """Execute the next command in the active playbook."""
        if not self.playbook.playbook:
            console.print("[warning]No active playbook. Use /playbook <type> to start one.[/warning]")
            return True

        if self.playbook.is_complete():
            console.print("[success]Playbook complete! Use /report to generate the final report.[/success]")
            return True

        # Check if target is set
        if "{target}" in str(self.playbook.get_current_command()):
            if not self.playbook.variables.get("target"):
                console.print("[warning]Target not set. Use /project target <ip_or_domain>[/warning]")
                return True

        cmd_info = self.playbook.get_current_command()
        if not cmd_info:
            console.print("[success]Playbook complete![/success]")
            return True

        resolved_cmd = self.playbook.resolve_command(cmd_info["cmd"])

        # Check for unresolved placeholders
        unresolved = re.findall(r'\{(\w+)\}', resolved_cmd)
        if unresolved:
            console.print(f"[warning]Missing variables: {', '.join(unresolved)}[/warning]")
            console.print(f"[dim]Set them with the playbook or skip this command with /skip[/dim]")
            return True

        phase = self.playbook.get_current_phase()
        console.print(Panel(
            f"[bold]Phase {phase['id']}:[/bold] {phase['name']}\n"
            f"[dim]{cmd_info['description']}[/dim]\n\n"
            f"[command]{resolved_cmd}[/command]",
            title="[bold red]Next Command[/bold red]",
            border_style="red",
            padding=(0, 1),
        ))

        # Execute based on AutoPilot level
        if self.autopilot.should_auto_execute() and not self.system.is_dangerous(resolved_cmd):
            self._execute_command(resolved_cmd)
            self.playbook.mark_completed(resolved_cmd)
            self.playbook.advance_command()
        else:
            try:
                choice = self.session.prompt(
                    ANSI("\033[1;32mExecute? (yes/skip/stop): \033[0m")
                )
                if choice.lower() in ("yes", "y", "sim", "s", ""):
                    self._execute_command(resolved_cmd)
                    self.playbook.mark_completed(resolved_cmd)
                    self.playbook.advance_command()
                elif choice.lower() in ("skip", "sk", "pular"):
                    self.playbook.skip_command()
                    console.print("[dim]Command skipped.[/dim]")
                else:
                    console.print("[dim]Playbook paused. Use /next to continue.[/dim]")
            except (KeyboardInterrupt, EOFError):
                console.print("\n[dim]Playbook paused.[/dim]")

        # Show progress
        if not self.playbook.is_complete():
            prog = self.playbook.get_progress()
            console.print(f"[dim]  Progress: {prog['progress_pct']}% | Phase {prog['phase']}/{prog['total_phases']}[/dim]")
        else:
            console.print("[success]Playbook complete! Use /report to generate the final report.[/success]")

        return True

    def _playbook_skip(self) -> bool:
        """Skip the current playbook command."""
        if not self.playbook.playbook:
            console.print("[warning]No active playbook.[/warning]")
            return True
        cmd = self.playbook.get_current_command()
        if cmd:
            console.print(f"[dim]Skipped: {cmd['description']}[/dim]")
        self.playbook.skip_command()
        return True

    # ─── Command Execution (with context saving) ────────────────────

    def _execute_command(self, command: str, retry_count: int = 0):
        """Execute a command with safety checks, AutoPilot awareness, auto-fix, and context saving."""
        is_dangerous = self.system.is_dangerous(command)
        is_install = command.strip().startswith("apt") or "install" in command

        # ── Safety check based on AutoPilot level ──
        if is_dangerous and not self.autopilot.should_skip_dangerous_confirm():
            console.print(Panel(
                f"[danger]DANGEROUS COMMAND:[/danger]\n[command]{command}[/command]\n\n"
                "This could cause irreversible damage.",
                title="[bold red]WARNING[/bold red]",
                border_style="red",
            ))
            try:
                confirm = self.session.prompt(
                    ANSI("\033[1;31mExecute anyway? (yes/NO): \033[0m")
                )
                if confirm.lower() not in ("yes", "y", "sim", "s"):
                    console.print("[dim]Cancelled.[/dim]")
                    return
            except (KeyboardInterrupt, EOFError):
                console.print("\n[dim]Cancelled.[/dim]")
                return

        tool = self.system.is_pentest_tool(command)
        if tool:
            console.print(f"[hacker]  Running: {tool}[/hacker]")

        console.print(f"[dim]$ {command}[/dim]")
        stdout, stderr, returncode = self.system.execute_command(command, interactive=True)

        if not stdout and stderr:
            console.print(f"[output]{stderr}[/output]")

        # ── SAVE TO PROJECT CONTEXT (THE KEY FEATURE) ──
        output_text = stdout or stderr or ""
        if self.project.name:
            # Save command output to context so AI can "remember" it
            self.project.add_context(
                entry_type="command_output",
                content=output_text,
                command=command,
                exit_code=returncode,
            )

            # Save raw scan output to project scans/ directory
            if tool or len(output_text) > 200:
                self.project.save_scan_output(command, output_text)

            # Auto-parse nmap output
            if "nmap" in command.lower() and returncode == 0 and output_text:
                try:
                    target_ip = self.project.metadata.get("target")
                    self.project.parse_nmap_output(output_text, target_ip=target_ip)
                    hosts = self.project.findings.get("hosts", {})
                    total_ports = sum(len(h.get("ports", [])) for h in hosts.values())
                    if hosts:
                        console.print(f"[success]  Auto-parsed: {len(hosts)} host(s), {total_ports} open port(s) saved to project[/success]")
                except Exception as e:
                    console.print(f"[dim]  Auto-parse warning: {e}[/dim]")

        if returncode != 0 and returncode != -1:
            console.print(f"[warning]Exit code: {returncode}[/warning]")

            # ── Auto-Fix Logic ──
            if self.autopilot.should_auto_fix() and retry_count < self.autopilot.auto_fix_retries:
                self._handle_error(command, stderr, returncode, retry_count)

    def _handle_error(self, original_command: str, stderr: str, returncode: int, retry_count: int):
        """Diagnose and attempt to fix a command error."""
        diagnosis = diagnose_error(original_command, stderr, returncode)

        console.print()
        console.print(Panel(
            f"[fix]AUTO-FIX[/fix] [dim](attempt {retry_count + 1}/{self.autopilot.auto_fix_retries})[/dim]\n\n"
            f"[dim]Error type:[/dim] [warning]{diagnosis['error_type']}[/warning]\n"
            f"[dim]Diagnosis:[/dim] {diagnosis['description']}",
            title="[bold cyan]Self-Healing[/bold cyan]",
            border_style="cyan",
            padding=(0, 1),
        ))

        if diagnosis["needs_ai"]:
            console.print("[info]Consulting AI for a fix...[/info]")
            fix_prompt = (
                f"A command failed and I need you to fix it. "
                f"Provide ONLY the corrected command(s) in a bash code block, no explanation needed.\n\n"
                f"Original command: `{original_command}`\n"
                f"Error output: ```\n{stderr[:1500]}\n```\n"
                f"Exit code: {returncode}\n\n"
                f"Fix this error. If a tool is missing, provide the install command first, then the original command."
            )
            self._stream_ai_response(fix_prompt, is_auto_fix=True)

        elif diagnosis["fix_command"]:
            fix_cmd = diagnosis["fix_command"]
            console.print(f"[fix]Applying fix:[/fix] [command]{fix_cmd}[/command]")

            if self.autopilot.should_auto_execute() or self.autopilot.should_auto_install():
                console.print(f"[dim]$ {fix_cmd}[/dim]")
                fix_stdout, fix_stderr, fix_rc = self.system.execute_command(fix_cmd, interactive=True)

                if not fix_stdout and fix_stderr:
                    console.print(f"[output]{fix_stderr}[/output]")

                if fix_rc == 0:
                    console.print(f"[success]Fix applied successfully.[/success]")
                    if diagnosis["fix_strategy"] != "try_sudo" and diagnosis["fix_strategy"] != "remove_sudo":
                        console.print(f"[info]Retrying original command...[/info]")
                        self._execute_command(original_command, retry_count=retry_count + 1)
                else:
                    console.print(f"[warning]Fix failed (exit {fix_rc}). Consulting AI...[/warning]")
                    fix_prompt = (
                        f"I tried to fix a command error but the fix also failed.\n\n"
                        f"Original command: `{original_command}`\n"
                        f"Original error: ```\n{stderr[:500]}\n```\n"
                        f"Fix attempted: `{fix_cmd}`\n"
                        f"Fix error: ```\n{fix_stderr[:500]}\n```\n\n"
                        f"Provide the correct fix as a bash code block."
                    )
                    self._stream_ai_response(fix_prompt, is_auto_fix=True)
            else:
                try:
                    confirm = self.session.prompt(
                        ANSI(f"\033[1;36mApply fix? (yes/no): \033[0m")
                    )
                    if confirm.lower() in ("yes", "y", "sim", "s"):
                        self._execute_command(fix_cmd, retry_count=retry_count + 1)
                        if diagnosis["fix_strategy"] == "install_package":
                            console.print(f"[info]Retrying original command...[/info]")
                            self._execute_command(original_command, retry_count=retry_count + 1)
                except (KeyboardInterrupt, EOFError):
                    console.print("\n[dim]Fix skipped.[/dim]")

    # ─── AI Response (with project context) ─────────────────────────

    def _stream_ai_response(self, user_input: str, is_auto_fix: bool = False):
        """Stream AI response with formatting, project context, and memory."""
        # Collect memory context
        memory_results = self.memory.search_memory(user_input)
        memory_context = self.memory.format_memory_context(memory_results)

        # Collect project context (THE KEY: feeds command outputs to AI)
        project_context = ""
        if self.project.name:
            project_context = self.project.get_context_for_ai()

        if memory_results and not is_auto_fix:
            console.print(f"[dim]  {len(memory_results)} related entries found in memory[/dim]")

        if project_context and not is_auto_fix:
            ctx_count = len(self.project.context.get("entries", []))
            console.print(f"[dim]  Project context: {ctx_count} entries loaded[/dim]")

        # Save user query to project context
        if self.project.name and not is_auto_fix:
            self.project.add_context(
                entry_type="user_query",
                content=user_input,
            )

        full_response = ""
        console.print()

        try:
            for token in self.ai.chat(
                user_input,
                memory_context=memory_context,
                project_context=project_context,
                stream=True,
            ):
                full_response += token
                # Check for funnel signals (don't print them)
                if "__FUNNEL_REGISTER__" in full_response:
                    print()  # newline
                    _show_registration_prompt()
                    return
                if "__FUNNEL_PAY__" in full_response:
                    print()  # newline
                    _show_payment_prompt()
                    return
                print(token, end="", flush=True)
            print()
        except KeyboardInterrupt:
            print("\n")
            console.print("[dim]Response interrupted.[/dim]")
            return

        # Save AI response to project context
        if self.project.name and not is_auto_fix:
            self.project.add_context(
                entry_type="ai_response",
                content=full_response,
            )

        # Extract commands from response
        commands = self.system.extract_commands(full_response)
        if commands:
            console.print()
            cmd_list = "\n".join([f"[bold green]{i+1}.[/bold green] [command]{cmd}[/command]" for i, cmd in enumerate(commands)])
            console.print(Panel(
                cmd_list,
                title="[bold red]Extracted Commands[/bold red]",
                border_style="red",
                padding=(0, 1),
            ))

            # ── AutoPilot-aware execution ──
            if self.autopilot.should_auto_execute() and not any(self.system.is_dangerous(c) for c in commands):
                console.print(f"[dim]  AutoPilot {self.autopilot.name}: auto-executing...[/dim]")
                for cmd in commands:
                    self._execute_command(cmd)
            elif self.autopilot.level == 5:
                console.print(f"[dim]  AutoPilot AUTOPILOT: executing all...[/dim]")
                for cmd in commands:
                    self._execute_command(cmd)
            else:
                try:
                    choice = self.session.prompt(
                        ANSI("\033[1;32mExecute? (number/all/no): \033[0m")
                    )

                    if choice.lower() in ("all", "a", "todos", "t"):
                        for cmd in commands:
                            self._execute_command(cmd)
                    elif choice.lower() in ("no", "n", "nao", "não", ""):
                        pass
                    elif choice.isdigit():
                        idx = int(choice) - 1
                        if 0 <= idx < len(commands):
                            self._execute_command(commands[idx])
                        else:
                            console.print("[warning]Invalid number.[/warning]")
                except (KeyboardInterrupt, EOFError):
                    console.print("\n[dim]Skipped.[/dim]")

        # Save interaction to collective memory
        if not is_auto_fix:
            executed_cmds = [e["command"] for e in self.system.get_command_history()[-len(commands):]] if commands else []
            self.memory.save_interaction(
                query=user_input,
                response=full_response,
                commands_executed=executed_cmds,
            )

    # ─── Prompt Builder ─────────────────────────────────────────────

    def _get_prompt_text(self):
        """Build prompt text with current context."""
        cwd = os.getcwd().replace(os.path.expanduser("~"), "~")
        stats = self.memory.get_stats()
        mem_total = stats["local_entries"] + stats["collective_entries"]

        project_name = self.project.name if self.project.name else None
        project_target = self.project.metadata.get("target") if self.project.name else None

        return build_prompt_prefix(
            cwd,
            model_name=self.ai.model_info["name"],
            mem_entries=mem_total,
            ap_level=self.autopilot.level,
            ap_name=self.autopilot.name,
            project_name=project_name,
            project_target=project_target,
        )

    # ─── Main Loop ──────────────────────────────────────────────────

    def run_interactive(self):
        """Run the interactive REPL loop."""
        self.show_startup()

        while self.running:
            try:
                prompt_text = self._get_prompt_text()

                user_input = self.session.prompt(
                    ANSI(prompt_text),
                )

                print_bottom_bar(
                    self.autopilot.level,
                    has_project=bool(self.project.name),
                    has_playbook=bool(self.playbook.playbook),
                )

                if not user_input.strip():
                    continue

                if user_input.strip() == "?":
                    self.handle_slash_command("/help")
                    continue

                if user_input.startswith("/"):
                    if self.handle_slash_command(user_input):
                        continue

                self._stream_ai_response(user_input)
                console.print()

            except KeyboardInterrupt:
                console.print("\n[dim]Press Ctrl+C again to exit, or keep typing.[/dim]")
                try:
                    prompt_text = self._get_prompt_text()
                    self.session.prompt(ANSI(prompt_text))
                    print_bottom_bar(
                        self.autopilot.level,
                        has_project=bool(self.project.name),
                        has_playbook=bool(self.playbook.playbook),
                    )
                except (KeyboardInterrupt, EOFError):
                    self.running = False
                    console.print("\n[dim]Session terminated. Stay sharp, hacker.[/dim]\n")
            except EOFError:
                self.running = False
                console.print("\n[dim]Session terminated. Stay sharp, hacker.[/dim]\n")

    def run_oneshot(self, query: str):
        """Run a single query and exit."""
        self._stream_ai_response(query)


def _handle_activate():
    """Handle the --activate flow (link device to user account).
    Used after 3 free queries when registration is required, or anytime."""
    console.print()
    console.print("[bold red]  MundiX CLI - Account Registration[/bold red]")
    console.print()
    console.print("[dim]  Generating device code...[/dim]")

    result = request_device_activation()

    if "error" in result:
        console.print(f"[danger]  Error: {result['error']}[/danger]")
        return

    device_code = result.get("device_code", "")

    console.print()
    console.print(Panel(
        f"[bold]Step 1:[/bold] Open this URL in your browser:\n\n"
        f"  [bold cyan]https://chat.mundix.dev/activate[/bold cyan]\n\n"
        f"[bold]Step 2:[/bold] Enter this code:\n\n"
        f"  [bold green]{device_code}[/bold green]\n\n"
        f"[dim]Waiting for activation...[/dim]",
        title="[bold red]Device Activation[/bold red]",
        border_style="red",
        padding=(1, 2),
    ))

    # Poll for activation
    import time as _time
    max_attempts = 60  # 5 minutes
    for i in range(max_attempts):
        try:
            _time.sleep(5)
            status = check_activation(device_code)

            if status.get("status") == "activated":
                api_key = status.get("api_key")
                if api_key:
                    config = load_config()
                    config["api_key"] = api_key
                    config["mode"] = "cloud"
                    config.pop("device_code", None)
                    save_config(config)

                    console.print()
                    console.print("[success]  Account activated successfully![/success]")
                    console.print(f"[dim]  API key saved to ~/.mundix/config.json[/dim]")
                    console.print()
                    console.print("  Now run: [bold]mundix[/bold]")
                    console.print()
                    return

            if i % 6 == 0 and i > 0:  # Every 30 seconds
                console.print(f"[dim]  Still waiting... ({i * 5}s)[/dim]")

        except KeyboardInterrupt:
            console.print("\n[dim]  Activation cancelled. Run 'mundix --activate' to try again.[/dim]")
            return

    console.print("[warning]  Activation timed out. Run 'mundix --activate' to try again.[/warning]")


def _handle_credits():
    """Handle the --credits command."""
    config = load_config()
    api_key = config.get("api_key")

    if not api_key:
        console.print("[warning]No API key found. MundiX will auto-register on first use.[/warning]")
        console.print("[dim]Run 'mundix' first to auto-register your device.[/dim]")
        return

    result = check_credits(api_key)

    if "error" in result:
        console.print(f"[danger]Error: {result['error']}[/danger]")
        return

    credits = result.get("credits", 0)
    used_tokens = result.get("total_tokens_used", 0)
    total_queries = result.get("total_queries", 0)
    queries_left = credits // 2000  # ~2000 credits per query
    status = result.get("status", "anonymous")
    max_free = result.get("max_free_queries", 3)
    is_blocked = result.get("is_blocked", False)

    # Determine display plan
    if status == "anonymous":
        plan_display = "FREE (Anonymous)"
    elif status == "registered":
        plan_display = "FREE (Registered)"
    elif status == "paid":
        plan_display = result.get("plan", "paid").upper()
    else:
        plan_display = status.upper()

    table = Table(border_style="green", title="[bold green]MundiX Credits[/bold green]")
    table.add_column("Property", style="bold")
    table.add_column("Value", style="white")
    table.add_row("Status", plan_display)
    table.add_row("Credits Remaining", f"{credits:,}")
    table.add_row("Tokens Used", f"{used_tokens:,}")
    table.add_row("Total Queries", f"{total_queries}")
    table.add_row("Free Queries Limit", f"{max_free}")
    table.add_row("Estimated Queries Left", f"~{queries_left}")
    if is_blocked:
        table.add_row("Blocked", "[bold red]YES[/bold red] - Purchase credits to unblock")
    console.print(table)

    if status == "anonymous" and total_queries >= max_free:
        console.print()
        console.print("[warning]  Register to get 5 more free queries: mundix --activate[/warning]")
    elif credits < 5000 and status != "anonymous":
        console.print()
        console.print("[warning]  Low credits! Top up at: https://chat.mundix.dev/billing[/warning]")


def _show_registration_prompt():
    """Display a friendly prompt to register (after 3 free queries)."""
    console.print()
    console.print(Panel(
        "[bold yellow]FREE QUERIES USED[/bold yellow]\n\n"
        "You've used your 3 free queries. Nice!\n"
        "Register with your email to unlock 5 more free queries.\n\n"
        "[bold]Option 1:[/bold] Run this command:\n"
        "  [bold green]mundix --activate[/bold green]\n\n"
        "[bold]Option 2:[/bold] Register at:\n"
        "  [bold cyan]https://chat.mundix.dev/activate[/bold cyan]\n\n"
        "[dim]Registration is free and takes 30 seconds.\n"
        "After registration, you get 5 more free queries![/dim]",
        title="[bold red]MundiX CLI[/bold red]",
        border_style="yellow",
        padding=(1, 2),
    ))
    console.print()


def _show_payment_prompt():
    """Display a prompt to purchase credits (after all free queries used)."""
    console.print()
    console.print(Panel(
        "[bold red]CREDITS EXHAUSTED[/bold red]\n\n"
        "You've used all your free queries.\n"
        "Purchase credits to continue using MundiX AI.\n\n"
        "[bold]Packages:[/bold]\n"
        "  [green]Starter[/green]    R$ 49.90   ~250 queries\n"
        "  [cyan]Pro[/cyan]        R$ 99.90   ~600 queries\n"
        "  [yellow]Enterprise[/yellow] R$ 199.90  ~1500 queries\n\n"
        "[bold]Purchase at:[/bold]\n"
        "  [bold cyan]https://chat.mundix.dev/billing[/bold cyan]\n\n"
        "[dim]Pay-as-you-go. No subscriptions. No BS.[/dim]",
        title="[bold red]MundiX CLI[/bold red]",
        border_style="red",
        padding=(1, 2),
    ))
    console.print()


def main():
    parser = argparse.ArgumentParser(
        description="MundiX CLI - AI-Powered Cybersecurity Copilot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  mundix                              Start interactive mode
  mundix "how to scan for open ports"  One-shot query
  mundix --model mx-redhawk-7b        Use MX-RedHawk model
  mundix --autopilot 3                Set AutoPilot to BALANCED
  mundix --project client1            Open/create project workspace
  mundix --exec "nmap -sV 10.0.0.1"   Execute and analyze with AI
  mundix --update                      Check for updates

AutoPilot Levels:
  1 MANUAL      Ask before every action
  2 GUIDED      Auto-execute safe commands (default)
  3 BALANCED    Auto-execute all, ask only dangerous
  4 AGGRESSIVE  Auto-execute everything, minimal prompts
  5 AUTOPILOT   Full autonomy, YOLO mode
        """,
    )
    parser.add_argument("query", nargs="?", help="One-shot query (optional)")
    parser.add_argument("--model", "-m", default=None, choices=list(MODELS.keys()),
                        help="AI model to use")
    # --token removed: all traffic goes through MundiX backend now
    parser.add_argument("--exec", "-e", dest="exec_cmd", default=None,
                        help="Execute a command and send output to AI for analysis")
    parser.add_argument("--autopilot", "-a", type=int, default=None, choices=[1, 2, 3, 4, 5],
                        help="Set AutoPilot level (1-5)")
    parser.add_argument("--project", "-p", default=None,
                        help="Open or create a project workspace")
    parser.add_argument("--update", "-u", action="store_true",
                        help="Check for updates")
    parser.add_argument("--activate", action="store_true",
                        help="Register your email to unlock more free queries")
    parser.add_argument("--credits", action="store_true",
                        help="Check remaining credits")
    parser.add_argument("--version", "-v", action="version", version=f"MundiX CLI v{VERSION}")

    args = parser.parse_args()

    # ── Commands that DON'T require activation ──
    if args.update:
        check_for_updates()
        return

    if args.activate:
        _handle_activate()
        return

    if args.credits:
        _handle_credits()
        return

    # ── ZERO FRICTION: Auto-register device and proceed ──
    # No activation gate! Users can start immediately.
    # The backend enforces the funnel (3 free → register → 5 more → pay)
    try:
        cli = MundixCLI(
            model_key=args.model,
            autopilot_level=args.autopilot,
            project_name=args.project,
        )
    except RuntimeError as e:
        error_msg = str(e)
        if "REGISTRATION_FAILED" in error_msg:
            console.print()
            console.print(Panel(
                f"[bold red]CONNECTION ERROR[/bold red]\n\n"
                f"Could not connect to MundiX servers.\n"
                f"[dim]{error_msg.split(': ', 1)[-1] if ': ' in error_msg else error_msg}[/dim]\n\n"
                f"Check your internet connection and try again.\n"
                f"If the problem persists, visit: [cyan]https://chat.mundix.dev[/cyan]",
                title="[bold red]MundiX CLI[/bold red]",
                border_style="red",
                padding=(1, 2),
            ))
            return
        raise

    if args.exec_cmd:
        console.print(f"[dim]$ {args.exec_cmd}[/dim]")
        stdout, stderr, rc = cli.system.execute_command(args.exec_cmd)
        if stdout:
            console.print(f"[output]{stdout}[/output]")
        prompt = f"I just ran this command: `{args.exec_cmd}`\n\nOutput:\n```\n{stdout[:3000]}\n```\n\nAnalyze the results and suggest next steps."
        cli._stream_ai_response(prompt)
    elif args.query:
        cli.run_oneshot(args.query)
    else:
        cli.run_interactive()


if __name__ == "__main__":
    main()
