webber/webber-cli/webber_cli/main.py

#!/usr/bin/env python3
"""
Webber CLI - Client for the Webber API.

Usage:
    webber-cli --help
    webber-cli chat [OPTIONS]
    webber-cli status
"""
import asyncio
import sys
from pathlib import Path

import typer
from prompt_toolkit import PromptSession
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.completion import Completer, Completion, PathCompleter
from prompt_toolkit.history import FileHistory
from prompt_toolkit.styles import Style
from rich.markdown import Markdown
from rich.panel import Panel
from rich.prompt import Confirm

from webber_cli.client import WebberClient, PermissionMode
from webber_cli.config import get_config, init_config, get_config_path, CONFIG_FILE
from webber_cli.theme import get_console, get_theme


# === Prompt Toolkit Setup ===

def _get_history_file() -> Path:
    """Get history file path from config."""
    config = get_config()
    return Path(config.history.file).expanduser()

# Built-in commands for completion
BUILTIN_COMMANDS = [
    "exit",
    "quit",
    "clear",
    "mode plan",
    "mode default",
    "mode auto_accept",
    "cd ",
]


class WebberCompleter(Completer):
    """Custom completer for Webber CLI commands."""

    def __init__(self, working_dir: str):
        self.working_dir = working_dir
        self.path_completer = PathCompleter(expanduser=True)

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

        # Complete built-in commands
        if not text or not text.startswith("cd "):
            for cmd in BUILTIN_COMMANDS:
                if cmd.startswith(text):
                    yield Completion(
                        cmd,
                        start_position=-len(text),
                        display_meta="command",
                    )

        # Complete file paths after "cd "
        if text.startswith("cd "):
            path_text = text[3:]
            # Create a sub-document for path completion
            from prompt_toolkit.document import Document
            path_doc = Document(path_text, len(path_text))
            for completion in self.path_completer.get_completions(path_doc, complete_event):
                yield Completion(
                    "cd " + (path_text + completion.text),
                    start_position=-len(text),
                    display_meta="directory",
                )


# Prompt style matching Rich theme
PROMPT_STYLE = Style.from_dict({
    "prompt": "#5f87d7 bold",  # info color
    "": "",  # default text
})


def create_prompt_session(working_dir: str) -> PromptSession:
    """Create a configured prompt session with history and completion."""
    return PromptSession(
        history=FileHistory(str(_get_history_file())),
        auto_suggest=AutoSuggestFromHistory(),
        completer=WebberCompleter(working_dir),
        style=PROMPT_STYLE,
        complete_while_typing=False,  # Only complete on Tab
    )

app = typer.Typer(
    name="webber-cli",
    help="CLI client for the Webber API",
    no_args_is_help=True,
    add_completion=False,
)

console = get_console()


def _get_api_url() -> str:
    """Get API URL from config (with env override already applied)."""
    return get_config().api.url


def _get_api_key() -> str:
    """Get API key from config (with env override already applied)."""
    return get_config().api.key or "webber-cli-dev-key"


def version_callback(value: bool) -> None:
    """Display version and exit."""
    if value:
        from webber_cli import __version__
        console.print(f"[title]webber-cli[/] version [success]{__version__}[/]")
        raise typer.Exit()


def _confirm_auto_accept() -> bool:
    """
    Prompt user to confirm auto_accept mode.

    Returns True if user confirms, False otherwise.
    """
    console.print()
    console.print("[warning]WARNING:[/] auto_accept mode bypasses all safety prompts.")
    console.print("The agent will execute write operations without confirmation.")
    console.print()
    return Confirm.ask(
        "[warning]Are you sure you want to enable auto_accept mode?[/]",
        default=False,
        console=console,
    )


@app.callback()
def main(
    version: bool = typer.Option(
        False,
        "--version",
        "-v",
        callback=version_callback,
        is_eager=True,
        help="Show version and exit",
    ),
) -> None:
    """Webber CLI - Talk to the Webber API."""
    pass


@app.command()
def sessions(
    api_url: str = typer.Option(
        None,
        "--api",
        "-a",
        help="Webber API URL (default from config)",
    ),
    limit: int = typer.Option(
        20,
        "--limit",
        "-n",
        help="Maximum number of sessions to show",
    ),
) -> None:
    """
    List previous conversation sessions.

    Shows recent sessions that can be resumed with 'chat --resume <id>'.
    """
    url = api_url or _get_api_url()
    asyncio.run(_list_sessions(url, limit))


async def _list_sessions(api_url: str, limit: int) -> None:
    """List conversation sessions."""
    async with WebberClient(api_url, api_key=_get_api_key()) as client:
        # Check API health
        if not await client.health_check():
            console.print(f"[error]Error:[/] Cannot connect to Webber API at {api_url}")
            return

        try:
            conversations, total = await client.list_conversations(limit=limit)
        except Exception as e:
            console.print(f"[error]Error listing sessions:[/] {e}")
            return

        if not conversations:
            console.print("[dim]No sessions found. Start one with 'webber-cli chat'[/]")
            return

        console.print(f"[title]Sessions[/] [dim]({len(conversations)} of {total})[/]\n")

        for conv in conversations:
            # Format the date
            date_str = conv.created_at.strftime("%Y-%m-%d %H:%M")

            # Title or first message preview
            title = conv.title or "[dim]untitled[/]"

            # Truncate ID for display
            short_id = conv.id[:8]

            console.print(
                f"  [info]{short_id}[/]  {date_str}  "
                f"[path]{conv.working_dir}[/]  {title}  "
                f"[dim]({conv.total_tokens} tokens)[/]"
            )

        console.print()
        console.print("[dim]Resume with: webber-cli chat --resume <id>[/]")


@app.command()
def chat(
    directory: str = typer.Option(
        ".",
        "--directory",
        "-d",
        help="Working directory for the agent",
    ),
    api_url: str = typer.Option(
        None,
        "--api",
        "-a",
        help="Webber API URL (default from config)",
    ),
    mode: str = typer.Option(
        None,
        "--mode",
        "-m",
        help="Permission mode: default, plan, auto_accept (default from config)",
    ),
    resume: str = typer.Option(
        None,
        "--resume",
        "-r",
        help="Resume a previous session by ID (use 'sessions' to list)",
    ),
    stream: bool = typer.Option(
        None,
        "--stream/--no-stream",
        "-s",
        help="Stream responses in real-time (default from config)",
    ),
) -> None:
    """
    Start interactive chat session with the Task agent.

    The Task agent is the main orchestrator that can:
    - Explore and analyze codebases
    - Plan implementation strategies
    - Execute code modifications (in default/auto_accept modes)
    - Spawn sub-agents for focused tasks

    Permission modes:
    - default: Full capabilities with approval prompts for writes
    - plan: Read-only mode for safe exploration and planning
    - auto_accept: Full capabilities without approval prompts (use with caution)

    Use --resume to continue a previous session.
    """
    # Resolve config defaults
    config = get_config()
    url = api_url or config.api.url
    mode_str = mode or config.cli.mode
    use_stream = stream if stream is not None else config.cli.stream

    working_dir = str(Path(directory).resolve())

    if not Path(working_dir).exists():
        console.print(f"[error]Error:[/] Directory not found: {working_dir}")
        raise typer.Exit(1)

    # Parse and validate mode
    try:
        permission_mode = PermissionMode(mode_str)
    except ValueError:
        console.print(f"[error]Error:[/] Invalid mode: {mode_str}")
        console.print("[dim]Valid modes: default, plan, auto_accept[/]")
        raise typer.Exit(1)

    # Confirm auto_accept mode (security risk)
    if permission_mode == PermissionMode.auto_accept:
        if not _confirm_auto_accept():
            console.print("[dim]Cancelled. Using default mode instead.[/]")
            permission_mode = PermissionMode.default

    try:
        asyncio.run(_chat_loop(url, working_dir, permission_mode, use_stream, resume))
    except KeyboardInterrupt:
        console.print("\n[dim]Goodbye![/]")


async def _chat_loop(
    api_url: str,
    working_dir: str,
    mode: PermissionMode,
    stream: bool = True,
    resume_id: str | None = None,
) -> None:
    """Interactive chat loop with the Task agent."""
    theme = get_theme()
    agent_type = "task"
    conversation_id: str | None = None
    conversation_title: str | None = None

    async with WebberClient(api_url, api_key=_get_api_key()) as client:
        # Check API health
        if not await client.health_check():
            console.print(f"[error]Error:[/] Cannot connect to Webber API at {api_url}")
            console.print("[dim]Make sure the server is running: ./wakeup.sh[/]")
            return

        # Get agent info
        agent_info = await client.get_agent(agent_type)
        if not agent_info:
            console.print(f"[error]Error:[/] Task agent not found")
            return

        # Handle resume or create new conversation
        if resume_id:
            # Try to find conversation by ID prefix
            try:
                conversations, _ = await client.list_conversations(limit=100)
                matching = [c for c in conversations if c.id.startswith(resume_id)]
                if not matching:
                    console.print(f"[error]Error:[/] Session not found: {resume_id}")
                    console.print("[dim]Use 'webber-cli sessions' to list available sessions[/]")
                    return
                if len(matching) > 1:
                    console.print(f"[error]Error:[/] Ambiguous ID, multiple matches: {resume_id}")
                    for m in matching:
                        console.print(f"  - {m.id[:8]} ({m.title or 'untitled'})")
                    return

                # Load the conversation with messages
                conv = await client.get_conversation(matching[0].id)
                if not conv:
                    console.print(f"[error]Error:[/] Could not load session")
                    return

                conversation_id = conv.id
                conversation_title = conv.title
                working_dir = conv.working_dir  # Use the session's working directory

                # Display conversation history
                console.print(f"\n[title]Resuming session[/] [dim]{conv.id[:8]}[/]")
                if conv.messages:
                    console.print(f"[dim]({len(conv.messages)} messages, {conv.total_tokens} tokens)[/]\n")
                    for msg in conv.messages[-6:]:  # Show last 6 messages
                        if msg.role == "user":
                            console.print(f"[prompt]>[/] {msg.content[:100]}{'...' if len(msg.content) > 100 else ''}")
                        else:
                            preview = msg.content[:200].replace('\n', ' ')
                            console.print(f"[dim]{preview}{'...' if len(msg.content) > 200 else ''}[/]\n")

            except Exception as e:
                console.print(f"[error]Error resuming session:[/] {e}")
                return
        else:
            # Create a new conversation
            try:
                conv = await client.create_conversation(
                    agent_type=agent_type,
                    working_dir=working_dir,
                    title=None,  # Will be set later based on first message
                )
                conversation_id = conv.id
                console.print(f"[dim]Session: {conv.id[:8]}[/]")
            except Exception as e:
                # If conversation API fails, continue without persistence
                console.print(f"[dim]Note: Session persistence unavailable ({e})[/]")

        # Mode display
        mode_display = {
            PermissionMode.default: "[info]default[/] (full with approvals)",
            PermissionMode.plan: "[success]plan[/] (read-only)",
            PermissionMode.auto_accept: "[warning]auto_accept[/] (no prompts)",
        }

        # Welcome message
        console.print()
        console.print(f"[title]Webber CLI[/] [dim]→ {api_url}[/]")
        console.print(f"[dim]Working in:[/] [path]{working_dir}[/]")
        console.print(f"[dim]Mode:[/] {mode_display[mode]}")
        console.print(f"[dim]Streaming:[/] {'enabled' if stream else 'disabled'}")
        console.print()
        console.print("[dim]Commands: 'exit' to quit, 'clear' to clear, 'mode <plan|default|auto_accept>' to switch[/]")
        console.print("[dim]Tab for completion, Up/Down for history[/]")
        console.print()

        current_mode = mode

        # Create prompt session with history and completion
        session = create_prompt_session(working_dir)

        # Chat loop
        while True:
            try:
                # Use prompt_toolkit for input (with history and completion)
                try:
                    user_input = await session.prompt_async(
                        [("class:prompt", "> ")],
                    )
                    user_input = user_input.strip()
                except EOFError:
                    # Ctrl+D pressed
                    console.print("[dim]Goodbye![/]")
                    break

                if not user_input:
                    continue

                if user_input.lower() in ("exit", "quit", "/exit", "/quit"):
                    console.print("[dim]Goodbye![/]")
                    break

                if user_input.lower() in ("clear", "/clear"):
                    console.clear()
                    continue

                if user_input.lower().startswith("cd "):
                    new_dir = user_input[3:].strip()
                    # Handle ~ expansion
                    new_path = Path(new_dir).expanduser().resolve()
                    if new_path.exists() and new_path.is_dir():
                        working_dir = str(new_path)
                        # Update completer's working directory
                        session.completer.working_dir = working_dir
                        console.print(f"[info]Changed to:[/] [path]{working_dir}[/]")
                    else:
                        console.print(f"[error]Directory not found:[/] {new_dir}")
                    continue

                # Mode switching
                if user_input.lower().startswith("mode "):
                    new_mode_str = user_input[5:].strip()
                    try:
                        new_mode = PermissionMode(new_mode_str)
                        if new_mode == PermissionMode.auto_accept:
                            if not _confirm_auto_accept():
                                console.print("[dim]Mode unchanged.[/]")
                                continue
                        current_mode = new_mode
                        console.print(f"[info]Mode changed to:[/] {mode_display[current_mode]}")
                    except ValueError:
                        console.print(f"[error]Invalid mode:[/] {new_mode_str}")
                        console.print("[dim]Valid modes: default, plan, auto_accept[/]")
                    continue

                console.print()

                if stream:
                    # Stream response in real-time, buffering for persistence
                    response_chunks: list[str] = []
                    try:
                        async for chunk in client.run_agent_stream(
                            agent_type, user_input, working_dir, current_mode
                        ):
                            sys.stdout.write(chunk)
                            sys.stdout.flush()
                            response_chunks.append(chunk)
                        console.print()  # Newline after streaming

                        # Save messages to conversation if we have a session
                        if conversation_id and response_chunks:
                            try:
                                full_response = "".join(response_chunks)
                                await client.save_messages(
                                    conversation_id,
                                    user_input,
                                    full_response,
                                )
                            except Exception as save_error:
                                # Log but don't fail the interaction
                                console.print(f"[dim]Note: Could not save to session ({save_error})[/]")

                    except Exception as e:
                        console.print(f"\n[error]Stream error:[/] {e}")
                else:
                    # Batch mode with spinner
                    with console.status("[info]Thinking...[/]", spinner=theme.spinner):
                        result = await client.run_agent(
                            agent_type, user_input, working_dir, current_mode
                        )

                    if result.success:
                        console.print(Markdown(result.response))

                        # Save messages to conversation if we have a session
                        if conversation_id:
                            try:
                                await client.save_messages(
                                    conversation_id,
                                    user_input,
                                    result.response,
                                )
                            except Exception as save_error:
                                console.print(f"[dim]Note: Could not save to session ({save_error})[/]")
                    else:
                        console.print(f"[error]Error:[/] {result.error}")

                console.print()

            except KeyboardInterrupt:
                console.print("\n[dim]Use 'exit' to quit.[/]")

            except Exception as e:
                console.print(f"[error]Error:[/] {e}")


@app.command()
def status(
    api_url: str = typer.Option(
        None,
        "--api",
        "-a",
        help="Webber API URL (default from config)",
    ),
) -> None:
    """Check API status and list available agents."""
    url = api_url or _get_api_url()
    asyncio.run(_status(url))


async def _status(api_url: str) -> None:
    """Check API status."""
    async with WebberClient(api_url) as client:
        console.print(f"[dim]API URL:[/] {api_url}")

        if await client.health_check():
            console.print("[success]Status:[/] Connected")

            agents = await client.list_agents()
            console.print(f"\n[dim]Available agents ({len(agents)}):[/]")
            for agent in agents:
                console.print(f"  [info]{agent.name}[/]: {agent.description}")
        else:
            console.print("[error]Status:[/] Cannot connect")


@app.command()
def config(
    init: bool = typer.Option(
        False,
        "--init",
        "-i",
        help="Initialize config file with defaults",
    ),
) -> None:
    """
    Show or initialize configuration.

    Without --init, displays current config and source.
    With --init, creates ~/.webber/config.toml with defaults.
    """
    if init:
        path = init_config()
        console.print(f"[success]Config initialized:[/] {path}")
        console.print("[dim]Edit this file to customize settings.[/]")
        return

    # Show current config
    cfg = get_config()
    config_path = get_config_path()

    console.print("[title]Webber Configuration[/]\n")

    if config_path:
        console.print(f"[dim]Config file:[/] {config_path}")
    else:
        console.print(f"[dim]Config file:[/] [warning]Not found[/] (using defaults)")
        console.print(f"[dim]Run 'webber-cli config --init' to create {CONFIG_FILE}[/]")

    console.print()
    console.print("[info]API Settings[/]")
    console.print(f"  url: {cfg.api.url}")
    console.print(f"  key: {'***' if cfg.api.key else '[dim]not set[/]'}")

    console.print()
    console.print("[info]CLI Settings[/]")
    console.print(f"  mode: {cfg.cli.mode}")
    console.print(f"  stream: {cfg.cli.stream}")

    console.print()
    console.print("[info]History[/]")
    console.print(f"  file: {cfg.history.file}")


# Keep 'explore' as an alias for 'chat --mode plan' for backwards compatibility
@app.command(hidden=True)
def explore(
    query: str = typer.Argument(..., help="What to search for"),
    directory: str = typer.Option(
        ".",
        "--directory",
        "-d",
        help="Working directory",
    ),
    api_url: str = typer.Option(
        None,
        "--api",
        "-a",
        help="Webber API URL (default from config)",
    ),
    stream: bool = typer.Option(
        None,
        "--stream/--no-stream",
        "-s",
        help="Stream responses in real-time (default from config)",
    ),
) -> None:
    """
    [DEPRECATED] One-shot exploration (use 'chat --mode plan' instead).

    Runs the Task agent in plan (read-only) mode for a single query.
    """
    # Resolve config defaults
    config = get_config()
    url = api_url or config.api.url
    use_stream = stream if stream is not None else config.cli.stream

    console.print("[dim]Note: 'explore' is deprecated. Use 'chat --mode plan' for interactive mode.[/]")
    console.print()

    working_dir = str(Path(directory).resolve())

    if not Path(working_dir).exists():
        console.print(f"[error]Error:[/] Directory not found: {working_dir}")
        raise typer.Exit(1)

    asyncio.run(_explore(url, query, working_dir, use_stream))


async def _explore(
    api_url: str, query: str, working_dir: str, stream: bool = True
) -> None:
    """Execute exploration query in plan mode."""
    theme = get_theme()
    mode = PermissionMode.plan

    async with WebberClient(api_url) as client:
        # Check API health
        if not await client.health_check():
            console.print(f"[error]Error:[/] Cannot connect to Webber API at {api_url}")
            console.print("[dim]Make sure the server is running: ./wakeup.sh[/]")
            return

        console.print(f"[dim]Exploring:[/] [path]{working_dir}[/]")
        console.print(f"[dim]Query:[/] {query}")
        console.print(f"[dim]Mode:[/] [success]plan[/] (read-only)")
        console.print()

        if stream:
            # Stream response in real-time
            try:
                async for chunk in client.run_agent_stream(
                    "task", query, working_dir, mode
                ):
                    sys.stdout.write(chunk)
                    sys.stdout.flush()
                console.print()  # Newline after streaming
            except Exception as e:
                console.print(f"\n[error]Stream error:[/] {e}")
        else:
            # Batch mode with spinner
            with console.status("[info]Searching...[/]", spinner=theme.spinner):
                result = await client.run_agent("task", query, working_dir, mode)

            if result.success:
                console.print(Panel(
                    Markdown(result.response),
                    title="[success]Findings[/]",
                    border_style=theme.colors.border_success,
                ))
            else:
                console.print(Panel(
                    f"[error]{result.error}[/]",
                    title="[error]Error[/]",
                    border_style=theme.colors.border_error,
                ))


if __name__ == "__main__":
    app()