Mayacontigo/.mise/tasks/container/start.py

#!/usr/bin/env -S uv run --script
#MISE silent=true
#MISE description="Run the container for a project"
# /// script
# dependencies = ["rich"]
# ///

import argparse
import json
import subprocess
import sys
from typing import Optional

from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.prompt import Prompt, Confirm
from rich.table import Table

console = Console()

# Configuration
VAULT_ADDRESS = "https://vault.ia-innovacion.work"
DOCKER_REGISTRY = "mayacontigo"
DEFAULT_PORT = 9000
CONTAINER_PORT = 80


def run_command(command: list[str], description: str, capture_output: bool = True) -> tuple[bool, str]:
    """Run a command with rich progress indication."""
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
        transient=True,
    ) as progress:
        task = progress.add_task(description, total=None)

        try:
            result = subprocess.run(
                command,
                capture_output=capture_output,
                text=True,
                check=True
            )
            progress.update(task, description=f"✅ {description}")
            return True, result.stdout.strip() if capture_output else ""
        except subprocess.CalledProcessError as e:
            progress.update(task, description=f"❌ {description}")
            error_msg = e.stderr if capture_output else str(e)
            console.print(f"[red]Error running command: {' '.join(command)}[/red]")
            if error_msg:
                console.print(f"[red]Error output: {error_msg}[/red]")
            return False, ""
        except FileNotFoundError:
            progress.update(task, description=f"❌ {description}")
            console.print(f"[red]Command not found: {command[0]}. Please ensure it's installed and in your PATH.[/red]")
            return False, ""


def get_vault_token() -> Optional[str]:
    """Get vault token from Vault CLI."""
    success, output = run_command(
        ["vault", "token", "lookup", "-format=json", f"-address={VAULT_ADDRESS}"],
        "Fetching Vault token"
    )

    if not success:
        return None

    try:
        data = json.loads(output)
        return data.get("data", {}).get("id")
    except json.JSONDecodeError:
        console.print("[red]Failed to parse Vault token response[/red]")
        return None


def check_docker_image_exists(app_name: str) -> bool:
    """Check if the Docker image exists locally or can be pulled."""
    image_name = f"{DOCKER_REGISTRY}/{app_name}"

    # First check if image exists locally
    success, _ = run_command(
        ["docker", "images", "-q", image_name],
        f"Checking if {image_name} exists locally"
    )

    if success:
        return True

    # If not local, try to pull it
    console.print(f"[yellow]Image {image_name} not found locally. Attempting to pull...[/yellow]")
    success, _ = run_command(
        ["docker", "pull", image_name],
        f"Pulling {image_name}"
    )

    return success


def get_available_apps() -> list[str]:
    """Get list of available Docker images from the registry."""
    try:
        result = subprocess.run(
            ["docker", "images", "--format", "{{.Repository}}:{{.Tag}}"],
            capture_output=True,
            text=True,
            check=True
        )
        images = [line.strip() for line in result.stdout.split('\n') if line.strip()]
        # Filter for banortegpt registry images
        apps = []
        for img in images:
            if img.startswith(f"{DOCKER_REGISTRY}/"):
                app_name = img.replace(f"{DOCKER_REGISTRY}/", "").split(":")[0]
                if app_name not in apps:
                    apps.append(app_name)
        return apps
    except subprocess.CalledProcessError:
        return []


def display_available_apps(apps: list[str]) -> None:
    """Display available apps in a nice table."""
    if not apps:
        console.print("[yellow]No apps found locally.[/yellow]")
        console.print("[dim]Try pulling an image first or check your Docker registry connection.[/dim]")
        return

    table = Table(title="Available Apps")
    table.add_column("App Name", style="cyan")
    table.add_column("Image", style="magenta")
    table.add_column("Status", style="green")

    for app in apps:
        image = f"{DOCKER_REGISTRY}/{app}"
        status = "🚀 Ready to run"
        table.add_row(app, image, status)

    console.print(table)


def interactive_mode() -> Optional[str]:
    """Run in interactive mode to select an app."""
    console.print(Panel.fit(
        "[bold blue]🐳 Docker Container Starter[/bold blue]\n"
        "[dim]Interactive mode - Let's select an app to run[/dim]",
        border_style="blue"
    ))

    # Get available apps
    console.print("\n[bold]Fetching available apps...[/bold]")
    apps = get_available_apps()

    if not apps:
        console.print("[red]No apps found. Please ensure you have Docker images available.[/red]")
        return None

    display_available_apps(apps)

    # Prompt for app selection
    console.print("\n[bold]Select an app to run:[/bold]")
    app_name = Prompt.ask(
        "Enter the app name (or press Enter to see suggestions)",
        default="",
        show_default=False
    )

    if not app_name:
        # Show suggestions
        console.print("\n[dim]Here are your available apps:[/dim]")
        for i, app in enumerate(apps, 1):
            console.print(f"  {i}. [cyan]{app}[/cyan]")

        choice = Prompt.ask(
            "Enter app name or number from suggestions",
            default="1"
        )

        try:
            choice_num = int(choice)
            if 1 <= choice_num <= len(apps):
                app_name = apps[choice_num - 1]
            else:
                console.print("[red]Invalid choice.[/red]")
                return None
        except ValueError:
            app_name = choice

    return app_name


def start_container(app_name: str, port: int = DEFAULT_PORT, force: bool = False) -> bool:
    """Start a Docker container for the specified app."""
    # Validate app name
    if not app_name:
        console.print("[red]Error: App name cannot be empty[/red]")
        return False

    # Check if image exists
    if not check_docker_image_exists(app_name):
        console.print(f"[red]Error: Image for app '{app_name}' not found[/red]")
        return False

    # Get vault token
    console.print("\n[bold]Getting Vault token...[/bold]")
    vault_token = get_vault_token()
    if not vault_token:
        console.print("[red]Failed to get Vault token[/red]")
        return False

    # Prepare container configuration
    image_name = f"{DOCKER_REGISTRY}/{app_name}"

    # Show what we're about to do
    info_panel = Panel.fit(
        f"[bold]App:[/bold] {app_name}\n"
        f"[bold]Image:[/bold] {image_name}\n"
        f"[bold]Port:[/bold] {port}:{CONTAINER_PORT}\n"
        f"[bold]Vault:[/bold] {VAULT_ADDRESS}\n"
        f"[bold]Token:[/bold] {'*' * 8}...{vault_token[-4:] if len(vault_token) > 8 else '****'}",
        title="🚀 Container Configuration",
        border_style="green"
    )
    console.print(info_panel)

    # Confirm action unless force is specified
    if not force:
        if not Confirm.ask(f"\nStart container for [cyan]{app_name}[/cyan] on port {port}?"):
            console.print("[yellow]Operation cancelled.[/yellow]")
            return False

    console.print("\n[bold]Starting container...[/bold]")

    # Build docker run command
    docker_cmd = [
        "docker", "run",
        "-p", f"{port}:{CONTAINER_PORT}",
        "--rm",
        "-it",
        "-e", f"VAULT_TOKEN={vault_token}",
        image_name
    ]

    # Run the container (this will be interactive)
    console.print(f"[green]Running: {' '.join(docker_cmd[:7])}... {image_name}[/green]")
    console.print(f"[dim]Container will be available at: http://localhost:{port}[/dim]")
    console.print("[dim]Press Ctrl+C to stop the container[/dim]\n")

    try:
        subprocess.run(docker_cmd, check=True)
        return True
    except subprocess.CalledProcessError as e:
        console.print(f"[red]Container failed to start: {e}[/red]")
        return False
    except KeyboardInterrupt:
        console.print("\n[yellow]Container stopped by user.[/yellow]")
        return True


def main():
    parser = argparse.ArgumentParser(
        description="Start Docker containers for Banorte GPT apps",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python start.py my-app                     # Start specific app
  python start.py --interactive              # Interactive mode
  python start.py my-app --port 8080         # Custom port
  python start.py my-app --force             # Skip confirmation
        """
    )

    parser.add_argument(
        "app_name",
        nargs="?",
        help="Name of the app to start"
    )

    parser.add_argument(
        "-i", "--interactive",
        action="store_true",
        help="Run in interactive mode"
    )

    parser.add_argument(
        "-p", "--port",
        type=int,
        default=DEFAULT_PORT,
        help=f"Port to bind to (default: {DEFAULT_PORT})"
    )

    parser.add_argument(
        "-f", "--force",
        action="store_true",
        help="Skip confirmation prompts"
    )

    parser.add_argument(
        "--list",
        action="store_true",
        help="List available apps and exit"
    )

    args = parser.parse_args()

    # Handle --list option
    if args.list:
        console.print("[bold]Available Apps:[/bold]")
        apps = get_available_apps()
        display_available_apps(apps)
        return

    # Determine app name
    app_name = args.app_name

    # Interactive mode or no app provided
    if args.interactive or not app_name:
        app_name = interactive_mode()
        if not app_name:
            sys.exit(1)

    # Start the container
    success = start_container(app_name, args.port, args.force)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n[yellow]Operation cancelled by user.[/yellow]")
        sys.exit(1)
    except Exception as e:
        console.print(f"[red]Unexpected error: {e}[/red]")
        sys.exit(1)