#!/usr/bin/env python3
import os
import sys
import shutil
import signal
import subprocess
import threading
import time
from pathlib import Path
import argparse


def which(cmd: str):
    return shutil.which(cmd)


def is_executable(p: Path):
    return p.exists() and os.access(str(p), os.X_OK)


def build_cmd_display(cmd):
    return " ".join(cmd) if isinstance(cmd, (list, tuple)) else str(cmd)


def pick_python_and_uvicorn(repo_root: Path):
    venv_bin = repo_root / ".venv" / "bin"
    candidates = {}

    # Pick python
    venv_python = venv_bin / "python"
    if is_executable(venv_python):
        candidates["python"] = str(venv_python)
    else:
        candidates["python"] = which("python3") or which("python") or sys.executable

    # Pick uvicorn
    venv_uvicorn = venv_bin / "uvicorn"
    if is_executable(venv_uvicorn):
        candidates["uvicorn_mode"] = "binary"
        candidates["uvicorn"] = str(venv_uvicorn)
    else:
        sys_uvicorn = which("uvicorn")
        if sys_uvicorn:
            candidates["uvicorn_mode"] = "binary"
            candidates["uvicorn"] = sys_uvicorn
        else:
            candidates["uvicorn_mode"] = "module"
            candidates["uvicorn"] = candidates["python"]

    return candidates


def stream_output(proc, prefix):
    # ANSI colors: backend -> green, frontend -> cyan, others -> yellow
    RESET = "\033[0m"
    COLOR = "\033[36m" if prefix == "frontend" else ("\033[32m" if prefix == "backend" else "\033[33m")

    for line in iter(proc.stdout.readline, b""):
        if not line:
            break
        try:
            decoded = line.decode(errors="replace")
            sys.stdout.write(f"{COLOR}[{prefix}]{RESET} {decoded}")
        except Exception:
            sys.stdout.write(f"{COLOR}[{prefix}]{RESET} {line!r}\n")
        sys.stdout.flush()


def start_process(cmd, cwd, prefix, env=None):
    proc = subprocess.Popen(
        cmd,
        cwd=cwd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        env=env or os.environ.copy(),
        text=False,
        bufsize=1,
    )
    t = threading.Thread(target=stream_output, args=(proc, prefix), daemon=True)
    t.start()
    return proc


def terminate_process(proc, name, timeout=8):
    if proc.poll() is not None:
        return
    try:
        proc.terminate()
    except Exception:
        pass
    for _ in range(timeout * 10):
        if proc.poll() is not None:
            return
        time.sleep(0.1)
    try:
        proc.kill()
    except Exception:
        pass


def main():
    parser = argparse.ArgumentParser(description="Dev runner: backend (FastAPI/uvicorn) + frontend (Next.js)")
    parser.add_argument("--backend-host", default=os.getenv("BACKEND_HOST", "127.0.0.1"))
    parser.add_argument("--backend-port", default=os.getenv("BACKEND_PORT", "8000"))
    parser.add_argument("--no-frontend", action="store_true", help="Start backend only")
    parser.add_argument("--no-backend", action="store_true", help="Start frontend only")
    parser.add_argument("--frontend-cmd", default=os.getenv("FRONTEND_CMD", "npm run dev"))
    parser.add_argument("--backend-app", default=os.getenv("BACKEND_APP", "main:app"), help="Uvicorn app path, e.g. main:app")
    args = parser.parse_args()

    # scripts/dev.py -> 仓库根目录
    repo_root = Path(__file__).resolve().parents[1]
    backend_dir = repo_root / "backend"
    frontend_dir = repo_root / "frontend"

    picks = pick_python_and_uvicorn(repo_root)
    py_path = picks["python"]
    uvicorn_mode = picks["uvicorn_mode"]
    uvicorn_ref = picks["uvicorn"]

    print("======== dev runner ========")
    print(f"Repo root: {repo_root}")
    print(f"Python:    {py_path}")
    if uvicorn_mode == "binary":
        print(f"Uvicorn:   {uvicorn_ref} (binary)")
    else:
        print(f"Uvicorn:   {uvicorn_ref} -m uvicorn (module)")
    print("============================")

    procs = []

    # Backend
    if not args.no_backend:
        if uvicorn_mode == "binary":
            backend_cmd = [
                uvicorn_ref,
                args.backend_app,
                "--reload",
                "--host", args.backend_host,
                "--port", args.backend_port,
            ]
        else:
            backend_cmd = [
                py_path, "-m", "uvicorn",
                args.backend_app,
                "--reload",
                "--host", args.backend_host,
                "--port", args.backend_port,
            ]
        print(f"Starting backend: {build_cmd_display(backend_cmd)} (cwd={backend_dir})")
        procs.append(("backend", start_process(backend_cmd, backend_dir, "backend")))

    # Frontend
    if not args.no_frontend:
        frontend_env = os.environ.copy()
        frontend_cmd = args.frontend_cmd
        print(f"Starting frontend: {frontend_cmd} (cwd={frontend_dir})")
        proc_fe = subprocess.Popen(
            frontend_cmd,
            cwd=frontend_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=frontend_env,
            shell=True,
            text=False,
            bufsize=1,
        )
        t = threading.Thread(target=stream_output, args=(proc_fe, "frontend"), daemon=True)
        t.start()
        procs.append(("frontend", proc_fe))

    def handle_signal(signum, frame):
        print("\nCaught signal, terminating children...")
        for name, p in procs:
            terminate_process(p, name)
        sys.exit(0)

    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            signal.signal(sig, handle_signal)
        except Exception:
            pass

    exit_code = 0
    try:
        while True:
            all_done = True
            for name, p in procs:
                ret = p.poll()
                if ret is None:
                    all_done = False
                else:
                    if ret != 0:
                        exit_code = ret
            if all_done:
                break
            time.sleep(0.3)
    except KeyboardInterrupt:
        handle_signal(signal.SIGINT, None)

    for name, p in procs:
        terminate_process(p, name)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()