#!/usr/bin/env python3
"""Bootstrap the Needle 3 structured data lab.

  1. Verifies Python >= 3.9
  2. Creates a virtual environment in .venv
  3. Pre-fetches the cactus engine and needle3 weights into the package
     cache (~/.cache/cactus, idempotent)
  4. Starts the web server on http://127.0.0.1:8001 (unless --no-serve)

Package installation (cactus-compute, FastAPI, uvicorn) is handled by
bootstrap.sh via uv before this script runs.

Usage:
  python bootstrap.py              # fetch (if needed), then serve
  python bootstrap.py --no-serve   # fetch only, then exit
  python bootstrap.py --port 9000  # serve on a different port

Uses only the Python standard library — safe to run with any bare 3.9+
interpreter on Windows or Linux.
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

MIN_PYTHON = (3, 9)
ROOT = Path(__file__).resolve().parent
VENV = ROOT / ".venv"


def die(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(1)


def venv_python() -> Path:
    if sys.platform == "win32":
        return VENV / "Scripts" / "python.exe"
    return VENV / "bin" / "python"


def venv_cactus() -> Path:
    if sys.platform == "win32":
        return VENV / "Scripts" / "cactus.exe"
    return VENV / "bin" / "cactus"


def run(cmd: list, env_extra: dict | None = None) -> None:
    """Run a command in the foreground, streaming its output."""
    env = dict(os.environ)
    env["QWEN_STREAM_MAX_LIFESPAN"] = "0"
    if env_extra:
        env.update(env_extra)
    print(f"$ {' '.join(str(c) for c in cmd)}", flush=True)
    proc = subprocess.run(cmd, cwd=ROOT, env=env)
    if proc.returncode != 0:
        die(f"command failed (exit {proc.returncode}): {' '.join(str(c) for c in cmd)}")


def check_python() -> None:
    if sys.version_info < MIN_PYTHON:
        die(
            f"Python >= {'.'.join(map(str, MIN_PYTHON))} is required; "
            f"found {sys.version.split()[0]}"
        )


def ensure_venv() -> None:
    py = venv_python()
    if py.exists():
        print(f"[1/2] Virtual environment already present at {VENV}")
        return
    print(f"[1/2] Creating virtual environment at {VENV}")
    run([sys.executable, "-m", "venv", str(VENV)])


def fetch_runtime() -> None:
    cli = venv_cactus()
    if not cli.exists():
        die("cactus CLI not found in .venv — install cactus-compute first (see bootstrap.sh)")
    print(f"[2/2] Ensuring cactus engine + needle3 weights are cached ...")
    run([str(cli), "fetch"])
    run([str(cli), "download", "needle3"])


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--no-serve", action="store_true", help="fetch only, do not start the server")
    parser.add_argument("--host", default="127.0.0.1", help="bind address (default: 127.0.0.1)")
    parser.add_argument("--port", type=int, default=8001, help="port (default: 8001)")
    args = parser.parse_args()

    check_python()
    ensure_venv()
    fetch_runtime()

    if args.no_serve:
        print("Done. Start the server later with:")
        print(f"  {venv_python()} -m uvicorn server:app --app-dir {ROOT} --host {args.host} --port {args.port}")
        return

    print(f"\nStarting server at http://{args.host}:{args.port}")
    run(
        [
            str(venv_python()),
            "-m",
            "uvicorn",
            "server:app",
            "--app-dir",
            str(ROOT),
            "--host",
            args.host,
            "--port",
            str(args.port),
        ]
    )


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(0)