konduktor/pyserve/ctl/main.py

"""
PyServeCTL - Main entry point

Usage:
    pyservectl [OPTIONS] COMMAND [ARGS]...
"""

import sys
from pathlib import Path
from typing import TYPE_CHECKING, Optional

import click

from .commands import (
    config_cmd,
    down_cmd,
    health_cmd,
    init_cmd,
    logs_cmd,
    ps_cmd,
    restart_cmd,
    scale_cmd,
    start_cmd,
    stop_cmd,
    top_cmd,
    up_cmd,
)

from .. import __version__

if TYPE_CHECKING:
    from ..config import Config
    from .state import StateManager

DEFAULT_CONFIG = "config.yaml"
DEFAULT_STATE_DIR = ".pyserve"


class Context:
    def __init__(self) -> None:
        self.config_file: str = DEFAULT_CONFIG
        self.state_dir: Path = Path(DEFAULT_STATE_DIR)
        self.verbose: bool = False
        self.debug: bool = False
        self.project: Optional[str] = None
        self._config: Optional["Config"] = None
        self._state: Optional["StateManager"] = None

    @property
    def config(self) -> "Config":
        if self._config is None:
            from ..config import Config

            if Path(self.config_file).exists():
                self._config = Config.from_yaml(self.config_file)
            else:
                self._config = Config()
        return self._config

    @property
    def state(self) -> "StateManager":
        if self._state is None:
            from .state import StateManager

            self._state = StateManager(self.state_dir, self.project)
        return self._state


pass_context = click.make_pass_decorator(Context, ensure=True)


@click.group(invoke_without_command=True)
@click.option(
    "-c",
    "--config",
    "config_file",
    default=DEFAULT_CONFIG,
    envvar="PYSERVE_CONFIG",
    help=f"Path to configuration file (default: {DEFAULT_CONFIG})",
    type=click.Path(),
)
@click.option(
    "-p",
    "--project",
    "project",
    default=None,
    envvar="PYSERVE_PROJECT",
    help="Project name for isolation",
)
@click.option(
    "-v",
    "--verbose",
    is_flag=True,
    help="Enable verbose output",
)
@click.option(
    "--debug",
    is_flag=True,
    help="Enable debug mode",
)
@click.version_option(version=__version__, prog_name="pyservectl")
@click.pass_context
def cli(ctx: click.Context, config_file: str, project: Optional[str], verbose: bool, debug: bool) -> None:
    """
    PyServeCTL - Service management CLI for PyServe.

    Docker-compose-like tool for managing PyServe services.

    \b
    Quick Start:
      pyservectl init          # Initialize a new project
      pyservectl up            # Start all services
      pyservectl ps            # Show service status
      pyservectl logs -f       # Follow logs
      pyservectl down          # Stop all services

    \b
    Examples:
      pyservectl up -d                    # Start in background
      pyservectl up -c prod.yaml          # Use custom config
      pyservectl logs api -f --tail 100   # Follow api logs
      pyservectl restart api admin        # Restart specific services
      pyservectl scale api=4              # Scale api to 4 workers
    """
    ctx.ensure_object(Context)
    ctx.obj.config_file = config_file
    ctx.obj.verbose = verbose
    ctx.obj.debug = debug
    ctx.obj.project = project

    if ctx.invoked_subcommand is None:
        click.echo(ctx.get_help())


cli.add_command(init_cmd, name="init")
cli.add_command(config_cmd, name="config")
cli.add_command(up_cmd, name="up")
cli.add_command(down_cmd, name="down")
cli.add_command(start_cmd, name="start")
cli.add_command(stop_cmd, name="stop")
cli.add_command(restart_cmd, name="restart")
cli.add_command(ps_cmd, name="ps")
cli.add_command(logs_cmd, name="logs")
cli.add_command(top_cmd, name="top")
cli.add_command(health_cmd, name="health")
cli.add_command(scale_cmd, name="scale")

# Alias 'status' -> 'ps'
cli.add_command(ps_cmd, name="status")


def main() -> None:
    try:
        cli(standalone_mode=False)
    except click.ClickException as e:
        e.show()
        sys.exit(e.exit_code)
    except KeyboardInterrupt:
        click.echo("\nInterrupted by user")
        sys.exit(130)
    except Exception as e:
        if "--debug" in sys.argv:
            raise
        click.secho(f"Error: {e}", fg="red", err=True)
        sys.exit(1)


if __name__ == "__main__":
    main()