python/wayy_db/cli/deploy.py

"""Deployment commands for the WayyDB CLI.

Supports:
- Local: start uvicorn directly or via Docker
- HuggingFace Spaces: push to HF Docker space
- Docker: build and run container
"""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Optional

import typer

from wayy_db.cli.config import load_config, save_config
from wayy_db.cli.output import console, print_error, print_info, print_success

deploy_app = typer.Typer(
    name="deploy",
    help="Deploy WayyDB service",
    no_args_is_help=True,
)


def _find_project_root() -> Path:
    """Walk up from cwd looking for pyproject.toml with wayy-db."""
    cwd = Path.cwd()
    for parent in [cwd, *cwd.parents]:
        toml = parent / "pyproject.toml"
        if toml.exists() and "wayy-db" in toml.read_text():
            return parent
    raise FileNotFoundError(
        "Cannot find wayyDB project root (no pyproject.toml with wayy-db found). "
        "Run this command from within the wayyDB repo."
    )


def _run(cmd: list[str], cwd: Optional[Path] = None, check: bool = True) -> subprocess.CompletedProcess[str]:
    """Run a subprocess with live output."""
    console.print(f"[dim]$ {' '.join(cmd)}[/dim]")
    return subprocess.run(cmd, cwd=cwd, check=check, text=True)


# --- Local serve ---


@deploy_app.command("local")
def deploy_local(
    port: int = typer.Option(8080, "--port", "-p", help="Port to serve on"),
    host: str = typer.Option("0.0.0.0", "--host", help="Host to bind to"),
    data_path: str = typer.Option("./data/wayydb", "--data-path", "-d", help="Data directory"),
    workers: int = typer.Option(1, "--workers", "-w", help="Number of uvicorn workers"),
) -> None:
    """Start WayyDB server locally with uvicorn."""
    os.makedirs(data_path, exist_ok=True)
    os.environ["WAYY_DATA_PATH"] = str(Path(data_path).resolve())
    os.environ["PORT"] = str(port)
    os.environ["CORS_ORIGINS"] = "*"

    print_info("Data path", os.environ["WAYY_DATA_PATH"])
    print_info("Serving on", f"http://{host}:{port}")
    console.print("[dim]Press Ctrl+C to stop[/dim]\n")

    try:
        _find_project_root()
        # Running from source — use api.main:app
        api_module = "api.main:app"
    except FileNotFoundError:
        # Installed package — api module should be importable
        api_module = "api.main:app"

    cmd = [
        sys.executable, "-m", "uvicorn", api_module,
        "--host", host,
        "--port", str(port),
        "--workers", str(workers),
    ]

    try:
        _run(cmd)
    except KeyboardInterrupt:
        console.print("\n[dim]Server stopped.[/dim]")
    except subprocess.CalledProcessError:
        print_error("Failed to start server. Is uvicorn installed? (pip install wayy-db[api])")
        raise typer.Exit(1)


# --- Docker ---


@deploy_app.command("docker")
def deploy_docker(
    port: int = typer.Option(8080, "--port", "-p", help="Host port to expose"),
    tag: str = typer.Option("wayydb:latest", "--tag", "-t", help="Docker image tag"),
    data_volume: str = typer.Option("wayydb-data", "--volume", "-v", help="Docker volume for data persistence"),
    build: bool = typer.Option(True, "--build/--no-build", help="Build image before running"),
    detach: bool = typer.Option(True, "--detach/--foreground", help="Run in background"),
) -> None:
    """Build and run WayyDB in Docker."""
    if not shutil.which("docker"):
        print_error("Docker not found. Install Docker: https://docs.docker.com/get-docker/")
        raise typer.Exit(1)

    try:
        root = _find_project_root()
    except FileNotFoundError as e:
        print_error(str(e))
        raise typer.Exit(1)

    if build:
        console.print("[bold]Building Docker image...[/bold]")
        _run(["docker", "build", "-t", tag, "."], cwd=root)
        print_success(f"Built {tag}")

    # Create volume if needed
    _run(["docker", "volume", "create", data_volume], check=False)

    # Stop existing container if running
    _run(["docker", "rm", "-f", "wayydb"], check=False)

    cmd = [
        "docker", "run",
        "--name", "wayydb",
        "-p", f"{port}:8080",
        "-v", f"{data_volume}:/data/wayydb",
        "-e", "CORS_ORIGINS=*",
    ]

    if detach:
        cmd.append("-d")

    cmd.append(tag)

    _run(cmd)

    if detach:
        print_success(f"WayyDB running at http://localhost:{port}")
        print_info("Container", "wayydb")
        print_info("Volume", data_volume)
        console.print("[dim]Stop with: docker stop wayydb[/dim]")
    else:
        console.print("\n[dim]Container stopped.[/dim]")


# --- HuggingFace Spaces ---


@deploy_app.command("hf")
def deploy_hf(
    repo: str = typer.Option("", "--repo", "-r", help="HF Space repo (user/name). Uses git remote 'hf' if not set."),
    token: Optional[str] = typer.Option(None, "--token", help="HuggingFace token (or set HF_TOKEN env var)"),
) -> None:
    """Deploy WayyDB to HuggingFace Spaces (Docker).

    Pushes the current repo state to a HuggingFace Space configured as a Docker space.
    The Space must already exist. Create one at: https://huggingface.co/new-space?sdk=docker
    """
    if not shutil.which("git"):
        print_error("git not found")
        raise typer.Exit(1)

    try:
        root = _find_project_root()
    except FileNotFoundError as e:
        print_error(str(e))
        raise typer.Exit(1)

    # Check if hf remote exists
    result = subprocess.run(
        ["git", "remote", "get-url", "hf"], capture_output=True, text=True, cwd=root
    )
    hf_remote_exists = result.returncode == 0
    existing_url = result.stdout.strip() if hf_remote_exists else ""

    if repo:
        hf_token = token or os.environ.get("HF_TOKEN", "")
        if hf_token:
            remote_url = f"https://user:{hf_token}@huggingface.co/spaces/{repo}"
        else:
            remote_url = f"https://huggingface.co/spaces/{repo}"

        if hf_remote_exists:
            _run(["git", "remote", "set-url", "hf", remote_url], cwd=root)
        else:
            _run(["git", "remote", "add", "hf", remote_url], cwd=root)
    elif not hf_remote_exists:
        print_error(
            "No 'hf' git remote found. Either:\n"
            "  1. Run: wayy deploy hf --repo <user>/<space-name>\n"
            "  2. Add manually: git remote add hf https://huggingface.co/spaces/<user>/<name>"
        )
        raise typer.Exit(1)

    console.print("[bold]Pushing to HuggingFace Spaces...[/bold]")

    # HF Spaces rejects pushes containing large files in history (even deleted ones).
    # Create a clean orphan commit with only the current tree to avoid this.
    try:
        # Create a temporary orphan branch with just the current working tree
        _run(["git", "checkout", "--orphan", "_hf_deploy"], cwd=root)
        _run(["git", "add", "-A"], cwd=root)
        _run(
            ["git", "commit", "-m", "Deploy wayyDB to HuggingFace Spaces", "--allow-empty"],
            cwd=root,
        )
        _run(["git", "push", "hf", "_hf_deploy:main", "--force"], cwd=root)
    except subprocess.CalledProcessError:
        # Clean up temp branch before erroring
        subprocess.run(["git", "checkout", "main"], cwd=root, capture_output=True)
        subprocess.run(["git", "branch", "-D", "_hf_deploy"], cwd=root, capture_output=True)
        print_error("Push failed. Check your HF token and Space configuration.")
        raise typer.Exit(1)
    finally:
        # Always return to main branch and clean up
        subprocess.run(["git", "checkout", "main"], cwd=root, capture_output=True)
        subprocess.run(["git", "branch", "-D", "_hf_deploy"], cwd=root, capture_output=True)

    # Extract space URL from remote
    result = subprocess.run(
        ["git", "remote", "get-url", "hf"], capture_output=True, text=True, cwd=root
    )
    remote_url = result.stdout.strip()

    # Parse space name from URL
    space_name = ""
    if "huggingface.co/spaces/" in remote_url:
        space_name = remote_url.split("huggingface.co/spaces/")[-1].rstrip(".git")
    elif repo:
        space_name = repo

    if space_name:
        space_url = f"https://huggingface.co/spaces/{space_name}"
        # HF Spaces with Docker get a direct URL
        space_direct = f"https://{space_name.replace('/', '-')}.hf.space"
        print_success(f"Deployed to HuggingFace Spaces")
        print_info("Space", space_url)
        print_info("API", space_direct)
        console.print(f"\n[dim]Connect with: wayy connect {space_direct}[/dim]")
    else:
        print_success("Pushed to HuggingFace Spaces")


# --- Status / logs ---


@deploy_app.command("stop")
def deploy_stop(
    name: str = typer.Option("wayydb", "--name", "-n", help="Container name"),
) -> None:
    """Stop a running WayyDB Docker container."""
    if not shutil.which("docker"):
        print_error("Docker not found")
        raise typer.Exit(1)

    _run(["docker", "stop", name], check=False)
    _run(["docker", "rm", name], check=False)
    print_success(f"Stopped {name}")


@deploy_app.command("logs")
def deploy_logs(
    name: str = typer.Option("wayydb", "--name", "-n", help="Container name"),
    follow: bool = typer.Option(False, "--follow", "-f", help="Follow log output"),
    tail: int = typer.Option(100, "--tail", help="Number of lines to show"),
) -> None:
    """View logs from a running WayyDB Docker container."""
    if not shutil.which("docker"):
        print_error("Docker not found")
        raise typer.Exit(1)

    cmd = ["docker", "logs", "--tail", str(tail)]
    if follow:
        cmd.append("-f")
    cmd.append(name)

    try:
        _run(cmd, check=False)
    except KeyboardInterrupt:
        pass