SneakyCode/app/main.py

"""SneakyCode entrypoint — argument parsing, config loading, and interactive REPL."""

import argparse
import asyncio
import signal
import sys
from pathlib import Path

import structlog

from app.agent.context import SessionContext
from app.agent.loop import AgentLoop
from app.models.config import AppConfig, load_config
from app.services.llm import LLMClient, LLMConnectionError, LLMError
from app.services.permissions import PermissionsService
from app.services.session import SessionManager
from app.services.streaming import StreamHandler
from app.tools.registry import create_default_registry
from app.utils.display import (
    print_banner,
    print_error,
    print_history,
    print_info,
    print_success,
    print_user_message,
    print_warning,
)
from app.utils.logging import console, get_logger, setup_logging


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments.

    Returns:
        Parsed arguments namespace.
    """
    parser = argparse.ArgumentParser(
        prog="sneakycode",
        description="SneakyCode — A privacy-first local AI coding agent",
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=None,
        help="Path to config YAML file (default: config/config.yaml)",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        default=False,
        help="Enable verbose (DEBUG) logging",
    )
    parser.add_argument(
        "--log-file",
        type=Path,
        default=None,
        help="Path to log file for persistent logging",
    )
    return parser.parse_args()


async def _preflight(config: AppConfig) -> None:
    """Check that Ollama is reachable and the configured model is available."""
    async with LLMClient(config.llm) as client:
        await client.preflight_check()


def _offer_session_resume(
    session_mgr: SessionManager, ctx: SessionContext
) -> bool:
    """Check for a saved session and offer to resume it.

    Args:
        session_mgr: Session manager instance.
        ctx: Session context to restore into.

    Returns:
        True if a session was restored.
    """
    saved = session_mgr.load_latest()
    if saved is None:
        return False

    msg_count = len(saved.messages)
    print_info(f"Found previous session ({msg_count} messages, model: {saved.model})")

    try:
        answer = console.input("[bold cyan]Resume previous session? [y/N] [/bold cyan]").strip().lower()
    except (KeyboardInterrupt, EOFError):
        return False

    if answer in ("y", "yes"):
        session_mgr.restore(saved, ctx)
        print_success(f"Session restored ({msg_count} messages).")
        return True

    return False


def _save_session_quiet(session_mgr: SessionManager, ctx: SessionContext) -> None:
    """Save session without raising on errors."""
    try:
        if ctx.message_count > 0:
            session_mgr.save(ctx)
    except OSError as e:
        print_warning(f"Could not save session: {e}")


async def _run_repl(
    ctx: SessionContext,
    config: AppConfig,
    session_mgr: SessionManager,
    logger: structlog.stdlib.BoundLogger,
    shutdown_event: asyncio.Event,
) -> None:
    """Run the interactive REPL loop with streaming LLM responses.

    Args:
        ctx: Session context for conversation state.
        config: Application configuration.
        session_mgr: Session manager for auto-save.
        logger: Structured logger instance.
        shutdown_event: Event signalling graceful shutdown.
    """
    registry = create_default_registry(config.agent.workspace_root, config)
    permissions = PermissionsService(config.permissions)

    async with LLMClient(config.llm) as client:
        handler = StreamHandler(config.display)
        agent = AgentLoop(config, ctx, client, handler, registry, permissions)

        while not shutdown_event.is_set():
            try:
                user_input = console.input("[bold cyan]> [/bold cyan]")
            except (KeyboardInterrupt, EOFError):
                _save_session_quiet(session_mgr, ctx)
                console.print("\n[dim]Session saved. Goodbye![/dim]")
                break

            user_input = user_input.strip()
            if not user_input:
                continue

            # Handle slash commands
            if user_input.startswith("/"):
                command = user_input.lower()
                if command == "/quit":
                    _save_session_quiet(session_mgr, ctx)
                    console.print("[dim]Session saved. Goodbye![/dim]")
                    break
                elif command == "/history":
                    print_history(ctx.get_history())
                elif command == "/clear":
                    ctx.clear_history()
                    print_success("Conversation history cleared.")
                elif command == "/save":
                    try:
                        path = session_mgr.save(ctx)
                        print_success(f"Session saved to {path}")
                    except OSError as e:
                        print_error(f"Failed to save session: {e}")
                elif command == "/session":
                    print_info(f"Messages: {ctx.message_count}")
                    print_info(f"Tokens: ~{ctx.estimated_tokens:,} / {ctx.token_counter.budget:,}")
                    print_info(f"Started: {ctx.start_time.isoformat()}")
                else:
                    print_warning(f"Unknown command: {user_input}")
                continue

            print_user_message(user_input)
            await agent.run_turn(user_input)
            logger.debug("turn_complete", message_count=ctx.message_count)

            # Auto-save after each turn
            if config.session.auto_save:
                _save_session_quiet(session_mgr, ctx)

        # Shutdown triggered by signal
        if shutdown_event.is_set():
            _save_session_quiet(session_mgr, ctx)
            console.print("\n[dim]Session saved. Shutting down gracefully.[/dim]")


def main() -> None:
    """Main entrypoint: load config, setup logging, launch interactive REPL."""
    args = parse_args()

    # Setup logging first
    setup_logging(
        log_file=args.log_file,
        verbose=args.verbose,
    )
    logger = get_logger(__name__)

    # Load configuration
    try:
        config = load_config(config_path=args.config)
    except (FileNotFoundError, ValueError) as e:
        print_error(f"Configuration error: {e}")
        sys.exit(1)

    logger.info("config_loaded", model=config.llm.model, endpoint=config.llm.endpoint)

    # Print startup info
    print_banner()
    print_info(f"Model: {config.llm.model}")
    print_info(f"Endpoint: {config.llm.endpoint}")
    print_info(f"Workspace: {config.agent.workspace_root}")

    if args.verbose:
        print_info("Verbose mode enabled")

    # Preflight: check Ollama is reachable and model exists
    try:
        asyncio.run(_preflight(config))
    except LLMConnectionError as e:
        print_error(str(e))
        sys.exit(1)
    except LLMError as e:
        print_error(str(e))
        sys.exit(1)

    print_success("Ollama connected, model ready.")

    # Create session and session manager
    ctx = SessionContext(config)
    session_mgr = SessionManager(config.session, config.agent.workspace_root, config.llm.model)

    # Clean up old session files
    cleaned = session_mgr.cleanup_old()
    if cleaned > 0:
        logger.info("old_sessions_cleaned", count=cleaned)

    # Offer to resume previous session
    if config.session.offer_resume:
        _offer_session_resume(session_mgr, ctx)

    logger.info("startup_complete")

    # Setup shutdown event and SIGTERM handler
    shutdown_event = asyncio.Event()

    original_sigterm = signal.getsignal(signal.SIGTERM)

    def _sigterm_handler(signum: int, frame: object) -> None:
        shutdown_event.set()

    signal.signal(signal.SIGTERM, _sigterm_handler)

    print_info("Commands: /quit, /history, /clear, /save, /session")

    try:
        asyncio.run(_run_repl(ctx, config, session_mgr, logger, shutdown_event))
    finally:
        signal.signal(signal.SIGTERM, original_sigterm)


if __name__ == "__main__":
    main()