#!/usr/bin/env python3
"""
Flask vs Next.js — end-to-end Lighthouse benchmark
----------------------------------------------------
1. Clones both repos
2. Installs dependencies & starts both servers
3. Runs Lighthouse against each
4. Pretty-prints results with colorama

Usage:
    python benchmark.py

Prerequisites: git, node, npm, psql, createdb, Python 3.9+
Lighthouse is installed automatically if missing.
"""

import subprocess, sys, os, json, time, shutil, socket, getpass
from pathlib import Path

# ── CONFIGURE THESE ──────────────────────────────────────────────────────────
FLASK_REPO    = "https://github.com/callmefredcom/flaskvibe-benchmark-flask"
NEXTJS_REPO   = "https://github.com/callmefredcom/flaskvibe-benchmark-nextjs"
WORK_DIR      = Path("./benchmark_run")
FLASK_PORT    = 5000
NEXTJS_PORT   = 3000
FLASK_PATH    = "/dashboard"   # path Lighthouse will audit on each app
NEXTJS_PATH   = "/dashboard"
_PG_USER      = getpass.getuser()
FLASK_DB_URL  = f"postgresql://{_PG_USER}@localhost/benchmark_flask"
NEXTJS_DB_URL = f"postgresql://{_PG_USER}@localhost/benchmark_nextjs"
# ─────────────────────────────────────────────────────────────────────────────


# ── Bootstrap colorama ────────────────────────────────────────────────────────
def _ensure_colorama():
    try:
        import colorama
    except ImportError:
        print("Installing colorama …")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "colorama", "-q"])

_ensure_colorama()
from colorama import init, Fore, Back, Style
init(autoreset=True)


# ── Helpers ───────────────────────────────────────────────────────────────────
def header(text):
    width = 62
    print()
    print(Fore.CYAN + Style.BRIGHT + "┌" + "─" * width + "┐")
    print(Fore.CYAN + Style.BRIGHT + "│" + f"  {text}".ljust(width) + "│")
    print(Fore.CYAN + Style.BRIGHT + "└" + "─" * width + "┘")

def step(text):
    print(Fore.YELLOW + "  ▸  " + Style.RESET_ALL + text)

def ok(text):
    print(Fore.GREEN + "  ✓  " + Style.RESET_ALL + text)

def err(text):
    print(Fore.RED + "  ✗  " + Style.RESET_ALL + text)

def info(text):
    print(Fore.BLUE + "  ℹ  " + Style.RESET_ALL + text)

def run(cmd, cwd=None, env=None, check=True):
    """Run a command, stream output, raise on failure."""
    result = subprocess.run(
        cmd, cwd=cwd, env=env,
        capture_output=True, text=True
    )
    if check and result.returncode != 0:
        err(f"Command failed: {' '.join(str(c) for c in cmd)}")
        print(Fore.RED + result.stderr[-2000:])
        sys.exit(1)
    return result

def wait_for_port(port, timeout=120, label="server"):
    """Block until a TCP port is open or timeout expires."""
    step(f"Waiting for {label} on port {port} …")
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=1):
                ok(f"{label} is up on :{port}")
                return True
        except OSError:
            time.sleep(1)
    err(f"Timed out waiting for {label} on port {port}")
    return False

def check_prereqs():
    header("Checking prerequisites")
    missing = []
    for tool in ("git", "node", "npm", "psql", "createdb"):
        if shutil.which(tool) is None:
            missing.append(tool)
            err(f"{tool} not found in PATH")
        else:
            ok(f"{tool} found")
    if missing:
        print(Fore.RED + f"\nInstall missing tools then re-run: {', '.join(missing)}")
        print(Fore.RED + "  psql/createdb come with PostgreSQL: https://www.postgresql.org/download/")
        sys.exit(1)

def ensure_lighthouse():
    header("Checking Lighthouse CLI")
    if shutil.which("lighthouse"):
        ok("lighthouse already installed")
        return
    step("Installing lighthouse globally …")
    run(["npm", "install", "-g", "lighthouse"])
    ok("lighthouse installed")

def clone_repos():
    header("Cloning repositories")
    WORK_DIR.mkdir(parents=True, exist_ok=True)

    for name, url in [("flask", FLASK_REPO), ("nextjs", NEXTJS_REPO)]:
        dest = WORK_DIR / name
        if dest.exists():
            info(f"{name} repo already cloned — skipping")
        else:
            step(f"Cloning {name} …")
            run(["git", "clone", url, str(dest)])
            ok(f"{name} cloned → {dest}")

def find_venv_python(venv_dir):
    """Return the first python binary that exists inside the venv."""
    ver = f"python{sys.version_info.major}.{sys.version_info.minor}"
    for name in ("python", "python3", ver):
        candidate = venv_dir / "bin" / name
        if candidate.exists():
            return candidate
    bin_dir = venv_dir / "bin"
    contents = list(bin_dir.iterdir()) if bin_dir.exists() else []
    err(f"No python binary found in {bin_dir}. Contents: {[p.name for p in contents]}")
    sys.exit(1)

def setup_flask():
    header("Setting up Flask app")
    app_dir = (WORK_DIR / "flask").resolve()
    venv_dir = app_dir / ".venv"

    step("Creating virtualenv …")
    run([sys.executable, "-m", "venv", str(venv_dir)])

    venv_python = find_venv_python(venv_dir)
    info(f"Using venv python: {venv_python}")

    step("Installing Python requirements …")
    run([str(venv_python), "-m", "pip", "install", "-r", "requirements.txt", "-q"], cwd=app_dir)
    ok("Flask deps installed")

    step("Creating Postgres database benchmark_flask …")
    r = run(["createdb", "benchmark_flask"], check=False)
    if r.returncode != 0 and "already exists" in r.stderr:
        info("benchmark_flask already exists — skipping createdb")
    elif r.returncode != 0:
        err(f"createdb failed: {r.stderr.strip()}")
        sys.exit(1)

    step("Applying schema …")
    schema = (app_dir / "schema.sql").read_text()
    r = subprocess.run(["psql", "benchmark_flask"], input=schema,
                       capture_output=True, text=True)
    if r.returncode != 0:
        err(f"psql schema failed: {r.stderr.strip()}")
        sys.exit(1)

    step("Seeding Flask database …")
    db_env = {**os.environ, "DATABASE_URL": FLASK_DB_URL}
    run([str(venv_python), "seed.py"], cwd=app_dir, env=db_env)
    ok("Flask database ready")

    return app_dir, venv_dir

def setup_nextjs():
    header("Setting up Next.js app")
    app_dir = (WORK_DIR / "nextjs").resolve()

    step("Running npm install …")
    run(["npm", "install", "--legacy-peer-deps"], cwd=app_dir)
    ok("Next.js deps installed")

    env_file = app_dir / ".env.local"
    step("Writing .env.local …")
    env_file.write_text(
        f"DATABASE_URL={NEXTJS_DB_URL}\n"
        "NEXTAUTH_SECRET=benchmark-secret-not-for-prod\n"
        f"NEXTAUTH_URL=http://localhost:{NEXTJS_PORT}\n"
    )
    ok(f".env.local written → {env_file}")

    step("Creating Postgres database benchmark_nextjs …")
    r = run(["createdb", "benchmark_nextjs"], check=False)
    if r.returncode != 0 and "already exists" in r.stderr:
        info("benchmark_nextjs already exists — skipping createdb")
    elif r.returncode != 0:
        err(f"createdb failed: {r.stderr.strip()}")
        sys.exit(1)

    step("Pushing Prisma schema …")
    db_env = {**os.environ, "DATABASE_URL": NEXTJS_DB_URL}
    run(["npx", "prisma", "db", "push", "--skip-generate", "--accept-data-loss"],
        cwd=app_dir, env=db_env)

    step("Seeding Next.js database …")
    run(["npm", "run", "db:seed"], cwd=app_dir, env=db_env)
    ok("Next.js database ready")

    step("Building Next.js for production (this takes ~60s) …")
    run(["npm", "run", "build"], cwd=app_dir, env=db_env)
    ok("Next.js production build complete")

    return app_dir

def start_flask(app_dir, venv_dir):
    header("Starting Flask server")
    venv_python = find_venv_python(venv_dir)
    env = {**os.environ, "FLASK_APP": "app", "FLASK_ENV": "development",
           "FLASK_RUN_PORT": str(FLASK_PORT), "DATABASE_URL": FLASK_DB_URL}

    for entry in ("app:app", "wsgi:app", "app"):
        env["FLASK_APP"] = entry
        proc = subprocess.Popen(
            [str(venv_python), "-m", "flask", "run", "--port", str(FLASK_PORT)],
            cwd=app_dir, env=env,
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        time.sleep(2)
        if proc.poll() is None:
            step(f"Flask process started (PID {proc.pid})")
            return proc
        proc.kill()

    err("Could not start Flask — check FLASK_APP in your repo")
    sys.exit(1)

def start_nextjs(app_dir):
    header("Starting Next.js server (production)")
    proc = subprocess.Popen(
        ["npm", "start"],
        cwd=app_dir,
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        env={**os.environ, "PORT": str(NEXTJS_PORT), "DATABASE_URL": NEXTJS_DB_URL}
    )
    step(f"Next.js process started (PID {proc.pid})")
    return proc

def run_lighthouse(url, output_path):
    step(f"Running Lighthouse → {url}")
    run([
        "lighthouse", url,
        "--output=json",
        "--throttling-method=simulate",
        f"--output-path={output_path}",
        "--chrome-flags=--headless --no-sandbox",
        "--quiet",
    ])
    ok(f"Results saved → {output_path}")

def score_color(score):
    if score >= 90:
        return Fore.GREEN + Style.BRIGHT
    if score >= 50:
        return Fore.YELLOW + Style.BRIGHT
    return Fore.RED + Style.BRIGHT

def ms(val):
    if val is None:
        return "n/a"
    return f"{val/1000:.2f}s" if val >= 1000 else f"{val:.0f}ms"

def parse_and_display(label, result_path):
    with open(result_path) as f:
        data = json.load(f)

    cats = data.get("categories", {})
    audits = data.get("audits", {})

    scores = {
        "Performance":     int((cats.get("performance",     {}).get("score", 0) or 0) * 100),
        "Accessibility":   int((cats.get("accessibility",   {}).get("score", 0) or 0) * 100),
        "Best Practices":  int((cats.get("best-practices",  {}).get("score", 0) or 0) * 100),
        "SEO":             int((cats.get("seo",             {}).get("score", 0) or 0) * 100),
    }

    metrics = {
        "First Contentful Paint":   audits.get("first-contentful-paint",  {}).get("numericValue"),
        "Largest Contentful Paint": audits.get("largest-contentful-paint", {}).get("numericValue"),
        "Time to Interactive":      audits.get("interactive",              {}).get("numericValue"),
        "Total Blocking Time":      audits.get("total-blocking-time",      {}).get("numericValue"),
        "Cumulative Layout Shift":  audits.get("cumulative-layout-shift",  {}).get("numericValue"),
        "Speed Index":              audits.get("speed-index",              {}).get("numericValue"),
    }

    col_w = 28
    print()
    banner = f"  {label} Results  "
    print(Fore.WHITE + Back.BLUE + Style.BRIGHT + banner.center(col_w + 16))
    print()

    print(Fore.CYAN + Style.BRIGHT + "  Lighthouse Scores")
    print(Fore.CYAN + "  " + "─" * (col_w + 10))
    for name, val in scores.items():
        bar_len = val // 5
        bar = "█" * bar_len + "░" * (20 - bar_len)
        color = score_color(val)
        print(f"  {name:<22} {color}{val:>3}/100  {bar}")

    print()

    print(Fore.CYAN + Style.BRIGHT + "  Core Web Vitals & Metrics")
    print(Fore.CYAN + "  " + "─" * (col_w + 10))
    for name, val in metrics.items():
        if name == "Cumulative Layout Shift":
            display = f"{val:.3f}" if val is not None else "n/a"
        else:
            display = ms(val)
        print(f"  {name:<30} {Fore.WHITE + Style.BRIGHT}{display}")

    return scores, metrics

def compare(flask_scores, flask_metrics, next_scores, next_metrics):
    header("Side-by-side Comparison")

    print(f"\n  {'Category':<26} {'Flask':>10}   {'Next.js':>10}   {'Winner'}")
    print("  " + "─" * 62)

    for cat in flask_scores:
        fs = flask_scores[cat]
        ns = next_scores[cat]
        if fs > ns:
            winner = Fore.GREEN + "Flask" + Style.RESET_ALL
        elif ns > fs:
            winner = Fore.BLUE + "Next.js" + Style.RESET_ALL
        else:
            winner = Fore.YELLOW + "Tie"
        fc = score_color(fs) + str(fs) + Style.RESET_ALL
        nc = score_color(ns) + str(ns) + Style.RESET_ALL
        print(f"  {cat:<26} {fc:>20}   {nc:>20}   {winner}")

    print()
    print(f"\n  {'Metric':<32} {'Flask':>10}   {'Next.js':>10}")
    print("  " + "─" * 62)

    for m in flask_metrics:
        fv = flask_metrics[m]
        nv = next_metrics.get(m)
        if m == "Cumulative Layout Shift":
            fd = f"{fv:.3f}" if fv is not None else "n/a"
            nd = f"{nv:.3f}" if nv is not None else "n/a"
        else:
            fd = ms(fv)
            nd = ms(nv)
        if fv is not None and nv is not None:
            if fv < nv:
                fd = Fore.GREEN + fd + Style.RESET_ALL
            elif nv < fv:
                nd = Fore.GREEN + nd + Style.RESET_ALL
        print(f"  {m:<32} {fd:>20}   {nd:>20}")

    print()


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    flask_proc = nextjs_proc = None

    print(Fore.CYAN + Style.BRIGHT + """
┌──────────────────────────────────────────────────────────────┐
│          Flask vs Next.js — Lighthouse Benchmark             │
└──────────────────────────────────────────────────────────────┘
""")

    try:
        check_prereqs()
        ensure_lighthouse()
        clone_repos()

        app_dir, venv_dir = setup_flask()
        nextjs_dir = setup_nextjs()

        flask_proc  = start_flask(app_dir, venv_dir)
        nextjs_proc = start_nextjs(nextjs_dir)

        if not wait_for_port(FLASK_PORT,  label="Flask"):
            sys.exit(1)
        if not wait_for_port(NEXTJS_PORT, label="Next.js"):
            sys.exit(1)

        time.sleep(3)

        flask_out  = WORK_DIR / "flask-results.json"
        nextjs_out = WORK_DIR / "nextjs-results.json"

        header("Running Lighthouse audits")
        run_lighthouse(f"http://localhost:{FLASK_PORT}{FLASK_PATH}",   str(flask_out))
        run_lighthouse(f"http://localhost:{NEXTJS_PORT}{NEXTJS_PATH}", str(nextjs_out))

        header("Flask Results")
        flask_scores, flask_metrics = parse_and_display("Flask", flask_out)

        header("Next.js Results")
        next_scores, next_metrics   = parse_and_display("Next.js", nextjs_out)

        compare(flask_scores, flask_metrics, next_scores, next_metrics)

        ok("Benchmark complete!")
        info(f"Raw JSON results saved in: {WORK_DIR.resolve()}")

    finally:
        header("Shutting down servers")
        for proc, name in [(flask_proc, "Flask"), (nextjs_proc, "Next.js")]:
            if proc and proc.poll() is None:
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    proc.kill()
                ok(f"{name} server stopped")


if __name__ == "__main__":
    main()