import os
import platform
import re
import shutil
import subprocess
import tarfile
import urllib.request
from pathlib import Path

import gradio as gr
import spaces

# ============================================================
# ZeroGPU requirement: at least one module-level @spaces.GPU
# function, or the Space fails startup with
# "No @spaces.GPU function detected during startup".
# Never called; the app serving port 7860 is the bansos daemon.
# ============================================================

@spaces.GPU
def _zerogpu_startup_stub():
    return None


# A launched Gradio server is what triggers the ZeroGPU startup report
# (spaces hooks gr.Blocks.launch). It serves only this stub page on 7861;
# bansos-router itself serves port 7860.
demo = gr.Blocks()
with demo:
    gr.Markdown(
        "bansos-router runs on this Space's port 7860. "
        "Open the root URL once the Space is Running."
    )


def launch_gradio_stub():
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7861,
        prevent_thread_lock=True,
        show_error=True,
    )


# ============================================================
# bansos-router launcher (npm package, prebuilt — no build)
# ============================================================

NODE_ROOT = Path("/tmp/node22")
PORT = "7860"


def log(message):
    print(f"\n[bansos HF] {message}", flush=True)


def run(cmd, cwd=None, env=None):
    log("+ " + " ".join(str(x) for x in cmd))
    subprocess.run(cmd, cwd=str(cwd) if cwd else None, env=env, check=True)


def node_arch():
    arch = platform.machine().lower()
    if arch in ("x86_64", "amd64"):
        return "x64"
    if arch in ("aarch64", "arm64"):
        return "arm64"
    raise RuntimeError(f"Unsupported architecture: {arch}")


def install_node22():
    node_bin = NODE_ROOT / "bin" / "node"
    npm_bin = NODE_ROOT / "bin" / "npm"
    if node_bin.exists() and npm_bin.exists():
        log("Node.js already available.")
        return node_bin, npm_bin

    log("Finding latest Node.js 22...")
    arch = node_arch()
    with urllib.request.urlopen(
        "https://nodejs.org/dist/latest-v22.x/SHASUMS256.txt", timeout=60
    ) as response:
        shasums = response.read().decode("utf-8")
    match = re.search(rf"([a-f0-9]{{64}})\s+(node-v[\d.]+-linux-{arch}\.tar\.xz)", shasums)
    if not match:
        raise RuntimeError("Node.js 22 binary was not found.")
    expected_hash, filename = match.group(1), match.group(2)
    archive = Path("/tmp") / filename
    log(f"Downloading {filename}...")
    urllib.request.urlretrieve(f"https://nodejs.org/dist/latest-v22.x/{filename}", archive)

    import hashlib
    digest = hashlib.sha256()
    with open(archive, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    if digest.hexdigest() != expected_hash:
        raise RuntimeError("Node.js SHA256 verification failed.")

    extract_root = Path("/tmp/node22-extract")
    shutil.rmtree(extract_root, ignore_errors=True)
    extract_root.mkdir(parents=True, exist_ok=True)
    log("Extracting Node.js...")
    with tarfile.open(archive, "r:xz") as tar:
        tar.extractall(extract_root)
    extracted = next(p for p in extract_root.iterdir() if p.is_dir())
    shutil.rmtree(NODE_ROOT, ignore_errors=True)
    shutil.move(str(extracted), str(NODE_ROOT))
    log("Node.js 22 installed.")
    return node_bin, npm_bin


def main():
    node_bin, npm_bin = install_node22()

    env = os.environ.copy()
    env["PATH"] = str(NODE_ROOT / "bin") + os.pathsep + env.get("PATH", "")
    # State dir (~/.bansos). Free-tier containers reset on restart.
    env["HOME"] = "/tmp"
    env["PORT"] = PORT

    run([str(node_bin), "--version"], env=env)

    bansos_bin = NODE_ROOT / "bin" / "bansos"
    if not bansos_bin.exists():
        log("Installing bansos-router (prebuilt, npm)...")
        run([str(npm_bin), "install", "--global", "--no-audit", "--no-fund", "bansos-router"], env=env)
    else:
        log("bansos-router already installed.")

    run([str(bansos_bin), "--version"], env=env)

    log("Launching Gradio stub on 7861 (ZeroGPU startup report)...")
    launch_gradio_stub()

    log(f"Starting bansos daemon on 0.0.0.0:{PORT}")
    # Foreground (no --bg): this process IS the Space. HF proxies
    # https://<space>.hf.space -> :7860, so allow non-loopback clients.
    os.execvpe(
        str(bansos_bin),
        [str(bansos_bin), "start",
         "--bind", "0.0.0.0", "--port", PORT,
         "--unsafe-allow-non-loopback"],
        env,
    )


if __name__ == "__main__":
    print("\n========================================\n  bansos-router - Hugging Face Space\n========================================\n", flush=True)
    main()