mytoolkit/bin/commands/server.py

"""Server management for workspace projects."""

import os
import signal
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Optional

import click


Mode = Literal["dev", "serve"]


@dataclass(frozen=True)
class Service:
    name: str
    serve_port: int
    dev_port: Optional[int] = None
    description: str = ""
    project_dir: Optional[Path] = None
    serve_cmd: Optional[list[str]] = None
    dev_cmd: Optional[list[str]] = None

    @property
    def manageable(self) -> bool:
        return self.project_dir is not None and self.serve_cmd is not None

    @property
    def has_dev(self) -> bool:
        return self.dev_port is not None and self.dev_cmd is not None

    def port_for(self, dev: bool) -> Optional[int]:
        return self.dev_port if dev else self.serve_port

    def cmd_for(self, dev: bool) -> Optional[list[str]]:
        return self.dev_cmd if dev else self.serve_cmd


KNOWN_SERVICES: list[Service] = [
    Service("metabot-api", 9100, description="MetaBot HTTP API"),
    Service("metabot-memory", 8100, description="MetaBot MetaMemory"),
    Service(
        "mywebpage",
        serve_port=3010,
        dev_port=3000,
        description="APAAM Lab website (Docusaurus)",
        project_dir=Path.home() / "Documents" / "myWork" / "sysu" / "04-学术积累" / "mywebpage",
        serve_cmd=["npm", "run", "serve"],
        dev_cmd=["npm", "run", "dev"],
    ),
    Service(
        "apaam-webpage",
        serve_port=3011,
        dev_port=3001,
        description="Phynexis docs (Docusaurus)",
        project_dir=Path.home() / "Documents" / "myResearch" / "myProjects" / "apaam" / "repo" / "webpage",
        serve_cmd=["npm", "run", "serve"],
        dev_cmd=["npm", "run", "dev"],
    ),
    Service(
        "xiaohe-webpage",
        serve_port=3012,
        dev_port=3002,
        description="Xiaohe company website (Docusaurus)",
        project_dir=Path.home() / "Documents" / "myResearch" / "myProjects" / "apaam" / "repo" / "xiaohe-webpage",
        serve_cmd=["npm", "run", "serve"],
        dev_cmd=["npm", "run", "dev"],
    ),
    Service("myslides", 3456, description="Slidev dev server"),
]

_BY_NAME: dict[str, Service] = {s.name: s for s in KNOWN_SERVICES}
_ALL_NAMES = list(_BY_NAME.keys())
_MANAGEABLE_NAMES = [s.name for s in KNOWN_SERVICES if s.manageable]


def manageable_services() -> list[Service]:
    return [s for s in KNOWN_SERVICES if s.manageable]


def get_service(name: str) -> Service:
    return _BY_NAME[name]


def no_color_env() -> dict:
    env = os.environ.copy()
    env["NO_COLOR"] = "1"
    return env


def _listening_map() -> dict[int, dict]:
    """Single batched lsof query for all TCP listening sockets, keyed by port."""
    try:
        result = subprocess.run(
            ["lsof", "-nP", "-iTCP", "-sTCP:LISTEN"],
            capture_output=True,
            text=True,
            errors="replace",
            timeout=5,
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return {}
    if result.returncode != 0 or not result.stdout.strip():
        return {}

    by_port: dict[int, dict] = {}
    for line in result.stdout.strip().split("\n")[1:]:
        parts = line.split()
        if len(parts) < 9:
            continue
        listen = next((t for t in reversed(parts) if ":" in t and not t.startswith("(")), None)
        if not listen:
            continue
        try:
            port = int(listen.rsplit(":", 1)[1])
        except (ValueError, IndexError):
            continue
        by_port[port] = {
            "command": parts[0],
            "pid": parts[1],
            "user": parts[2],
            "listen": listen,
        }
    return by_port


def _port_status(port: int) -> Optional[dict]:
    return _listening_map().get(port)


def _pid_command(pid: str) -> str:
    try:
        result = subprocess.run(
            ["ps", "-p", pid, "-o", "comm="],
            capture_output=True,
            text=True,
            timeout=2,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass
    return ""


def _resolve(svc: Service, dev: bool) -> tuple[Optional[int], Optional[list[str]], Mode]:
    return svc.port_for(dev), svc.cmd_for(dev), "dev" if dev else "serve"


@click.group(name="server")
def server_cmd():
    """Manage workspace dev/serve servers."""
    pass


def _print_row(name: str, mode: str, port: int, desc: str, info: Optional[dict]) -> bool:
    if info:
        pid = info.get("pid", "?")
        click.secho(
            f"{name:<16} {mode:<6} {port:<6} {'running':<10} {pid:<8} {desc}",
            fg="green",
        )
        return True
    click.secho(
        f"{name:<16} {mode:<6} {port:<6} {'stopped':<10} {'-':<8} {desc}",
        fg="black",
    )
    return False


@server_cmd.command("list")
def server_list():
    """List all known services (serve + dev modes) and their status."""
    click.echo("")
    click.secho(
        f"{'Service':<16} {'Mode':<6} {'Port':<6} {'Status':<10} {'PID':<8} {'Description'}",
        fg="bright_blue",
        bold=True,
    )
    click.echo("-" * 90)

    listening = _listening_map()
    running = 0
    total = 0
    for svc in KNOWN_SERVICES:
        primary_label = "serve" if svc.serve_cmd is not None else "-"
        if _print_row(svc.name, primary_label, svc.serve_port, svc.description, listening.get(svc.serve_port)):
            running += 1
        total += 1

        if svc.has_dev:
            assert svc.dev_port is not None
            if _print_row("", "dev", svc.dev_port, "(hot reload)", listening.get(svc.dev_port)):
                running += 1
            total += 1

    click.echo("-" * 90)
    click.echo(f"Total: {total} ports, {running} running, {total - running} stopped")
    click.echo("")


@server_cmd.command("check")
@click.argument("port", type=int)
def server_check(port):
    """Check if a specific port is open."""
    info = _port_status(port)
    if info:
        cmd = _pid_command(info.get("pid", ""))
        click.secho(
            f"Port {port} is open (PID {info.get('pid', '?')}, {cmd or info.get('command', '')})",
            fg="green",
        )
    else:
        click.secho(f"Port {port} is not listening", fg="red")


@server_cmd.command("scan")
@click.argument("start", type=int, default=3000)
@click.argument("end", type=int, default=4000)
def server_scan(start, end):
    """Scan a port range for listening services."""
    click.echo(f"Scanning ports {start}-{end} ...")
    listening = _listening_map()
    found = 0
    for port in range(start, end + 1):
        info = listening.get(port)
        if info:
            cmd = _pid_command(info.get("pid", ""))
            click.secho(
                f"  {port:<5} {cmd or info.get('command', '')} (PID {info.get('pid', '?')})",
                fg="green",
            )
            found += 1
    if found == 0:
        click.echo("No listening ports found.")
    else:
        click.echo(f"\nFound {found} open port(s).")


@server_cmd.command("start")
@click.argument("name", type=click.Choice(_MANAGEABLE_NAMES, case_sensitive=False))
@click.option("--dev", is_flag=True, help="Start hot-reload dev server on the dev port (default: serve)")
@click.option("--rebuild", is_flag=True, help="Force npm run build before serve (no effect with --dev)")
def server_start(name, dev, rebuild):
    """Start a server in the background.

    By default, serve mode is used (production build) on the serve port.
    Use --dev to start the hot-reload dev server on the dev port instead.
    """
    svc = get_service(name)
    port, cmd, mode = _resolve(svc, dev)

    if port is None or cmd is None:
        click.secho(f"{name} has no {mode} mode configured", fg="red")
        return

    assert svc.project_dir is not None  # _MANAGEABLE_NAMES choice guarantees this

    if _port_status(port) is not None:
        click.secho(f"{name} ({mode}) is already running on port {port}", fg="yellow")
        return

    if not svc.project_dir.exists():
        click.secho(f"Project directory not found: {svc.project_dir}", fg="red")
        return

    env = no_color_env()

    if not dev:
        build_dir = svc.project_dir / "build"
        if rebuild or not build_dir.exists():
            reason = "forced rebuild" if rebuild else "build directory missing"
            click.secho(f"Running npm run build ({reason}) ...", fg="cyan")
            build_proc = subprocess.run(
                ["npm", "run", "build"],
                cwd=svc.project_dir,
                env=env,
                capture_output=True,
                text=True,
            )
            if build_proc.returncode != 0:
                click.secho(f"Build failed for {name}", fg="red")
                click.echo(build_proc.stderr[-500:] if build_proc.stderr else "")
                return

    log_file = svc.project_dir / "logs" / f"{svc.name}-{mode}.log"
    log_file.parent.mkdir(exist_ok=True)
    full_cmd = ["nohup"] + cmd
    click.secho(f"Starting {name} ({svc.description}) on port {port} ({mode}) ...", fg="cyan")

    with open(log_file, "a") as log:
        proc = subprocess.Popen(
            full_cmd,
            cwd=svc.project_dir,
            stdout=log,
            stderr=subprocess.STDOUT,
            start_new_session=True,
            env=env,
        )

    max_wait = 20 if dev else 30
    for _ in range(max_wait * 2):
        time.sleep(0.5)
        if _port_status(port) is not None:
            click.secho(f"{name} ({mode}) started successfully (PID {proc.pid})", fg="green")
            return

    click.secho(f"{name} may have failed to start (check {log_file})", fg="yellow")


def _kill_port(name: str, mode: Mode, port: int):
    info = _port_status(port)
    if info is None:
        click.secho(f"{name} ({mode}) is not running on port {port}", fg="yellow")
        return

    pid = info["pid"]
    try:
        os.kill(int(pid), signal.SIGTERM)
        click.secho(f"Sent SIGTERM to {name} ({mode}, PID {pid})", fg="green")
    except ProcessLookupError:
        click.secho(f"Process {pid} not found", fg="yellow")
        return

    for _ in range(10):
        time.sleep(0.5)
        if _port_status(port) is None:
            click.secho(f"{name} ({mode}) stopped", fg="green")
            return

    try:
        os.kill(int(pid), signal.SIGKILL)
        click.secho(f"Sent SIGKILL to {name} ({mode}, PID {pid})", fg="red")
    except ProcessLookupError:
        pass


@server_cmd.command("stop")
@click.argument("name", type=click.Choice(_MANAGEABLE_NAMES, case_sensitive=False))
@click.option("--dev", is_flag=True, help="Stop the dev-port instance instead of serve")
@click.option("--all", "all_modes", is_flag=True, help="Stop both dev and serve instances")
def server_stop(name, dev, all_modes):
    """Stop a running server."""
    svc = get_service(name)

    if all_modes:
        targets: list[tuple[Mode, int]] = [("serve", svc.serve_port)]
        if svc.has_dev:
            assert svc.dev_port is not None
            targets.append(("dev", svc.dev_port))
    else:
        port, _, mode = _resolve(svc, dev)
        if port is None:
            click.secho(f"{name} has no {mode} mode configured", fg="red")
            return
        targets = [(mode, port)]

    for target_mode, target_port in targets:
        _kill_port(name, target_mode, target_port)


@server_cmd.command("restart")
@click.argument("name", type=click.Choice(_MANAGEABLE_NAMES, case_sensitive=False))
@click.option("--dev", is_flag=True, help="Restart the dev instance (default: serve)")
@click.option("--rebuild", is_flag=True, help="Force npm run build before serve (no effect with --dev)")
def server_restart(name, dev, rebuild):
    """Restart a server."""
    ctx = click.get_current_context()
    ctx.invoke(server_stop, name=name, dev=dev, all_modes=False)
    ctx.invoke(server_start, name=name, dev=dev, rebuild=rebuild)


@server_cmd.command("status")
@click.argument("name", type=click.Choice(_ALL_NAMES, case_sensitive=False))
def server_status(name):
    """Show detailed status of a service (both serve and dev ports)."""
    svc = get_service(name)

    click.echo("")
    click.secho(f"Service:  {svc.name}", bold=True)
    click.echo(f"Desc:     {svc.description}")
    if svc.project_dir:
        click.echo(f"Project:  {svc.project_dir}")

    def _show(label: str, port: Optional[int]):
        if port is None:
            return
        info = _port_status(port)
        if info:
            cmd = _pid_command(info.get("pid", ""))
            click.secho(
                f"{label:<8} {port}    running (PID {info.get('pid', '?')}, {cmd or info.get('command', '')})",
                fg="green",
            )
        else:
            click.secho(f"{label:<8} {port}    stopped", fg="black")

    _show("Serve:", svc.serve_port)
    _show("Dev:", svc.dev_port)
    click.echo("")