#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# ///
"""Install uv and the pinned shared Rust toolchain for zccache development."""

from __future__ import annotations

import os
import platform
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request
from contextlib import contextmanager
from pathlib import Path

import tomllib

GREEN = "\033[1;32m"
YELLOW = "\033[1;33m"
RED = "\033[1;31m"
RESET = "\033[0m"
ROOT = Path(__file__).resolve().parent
TOOLCHAIN_FILE = ROOT / "rust-toolchain.toml"
LOCK_TIMEOUT_SECONDS = 15 * 60
LOCK_POLL_SECONDS = 0.2


def log(msg: str) -> None:
    print(f"{GREEN}[install]{RESET} {msg}")


def warn(msg: str) -> None:
    print(f"{YELLOW}[install]{RESET} {msg}")


def err(msg: str) -> None:
    print(f"{RED}[install]{RESET} {msg}", file=sys.stderr)


def die(msg: str) -> None:
    err(msg)
    raise SystemExit(1)


def cargo_home() -> Path:
    if os.environ.get("CARGO_HOME"):
        return Path(os.environ["CARGO_HOME"]).expanduser()
    return Path.home() / ".cargo"


def rustup_home() -> Path:
    if os.environ.get("RUSTUP_HOME"):
        return Path(os.environ["RUSTUP_HOME"]).expanduser()
    return Path.home() / ".rustup"


def cargo_bin_dir() -> Path:
    return cargo_home() / "bin"


def prepend_cargo_bin_to_path() -> None:
    bin_dir = cargo_bin_dir()
    current_path = os.environ.get("PATH", "")
    path_parts = current_path.split(os.pathsep) if current_path else []
    normalized_bin = os.path.normcase(os.path.normpath(str(bin_dir)))
    normalized_parts = {
        os.path.normcase(os.path.normpath(part))
        for part in path_parts
        if part
    }
    if normalized_bin not in normalized_parts:
        os.environ["PATH"] = str(bin_dir) + (os.pathsep + current_path if current_path else "")


def rustup_exe_name() -> str:
    return "rustup.exe" if os.name == "nt" else "rustup"


def homes_are_overridden() -> bool:
    return "CARGO_HOME" in os.environ or "RUSTUP_HOME" in os.environ


def rustup_executable() -> str | None:
    candidate = cargo_bin_dir() / rustup_exe_name()
    if candidate.exists():
        return str(candidate)
    if homes_are_overridden():
        return None
    return shutil.which("rustup")


def rustup_env() -> dict[str, str]:
    env = os.environ.copy()
    env["RUSTUP_INIT_SKIP_PATH_CHECK"] = "yes"
    env["CARGO_HOME"] = str(cargo_home())
    env["RUSTUP_HOME"] = str(rustup_home())
    return env


def host_target_triple() -> str:
    system = platform.system()
    machine = platform.machine().lower()
    arch = {
        "amd64": "x86_64",
        "x86_64": "x86_64",
        "arm64": "aarch64",
        "aarch64": "aarch64",
    }.get(machine)
    if arch is None:
        die(f"unsupported architecture: {machine}")
    if system == "Windows":
        return f"{arch}-pc-windows-msvc"
    if system == "Linux":
        return f"{arch}-unknown-linux-gnu"
    if system == "Darwin":
        return f"{arch}-apple-darwin"
    die(f"unsupported platform: {system}")


def rustup_init_url() -> str:
    suffix = ".exe" if os.name == "nt" else ""
    return f"https://static.rust-lang.org/rustup/dist/{host_target_triple()}/rustup-init{suffix}"


def load_toolchain_spec() -> dict[str, object]:
    with TOOLCHAIN_FILE.open("rb") as handle:
        data = tomllib.load(handle)
    toolchain = data.get("toolchain")
    if not isinstance(toolchain, dict):
        die(f"missing [toolchain] in {TOOLCHAIN_FILE}")

    channel = toolchain.get("channel")
    profile = toolchain.get("profile", "minimal")
    components = toolchain.get("components", [])
    targets = toolchain.get("targets", [])

    if not isinstance(channel, str) or not channel:
        die(f"missing toolchain.channel in {TOOLCHAIN_FILE}")
    if not isinstance(profile, str):
        die(f"toolchain.profile must be a string in {TOOLCHAIN_FILE}")
    if not isinstance(components, list) or not all(isinstance(item, str) for item in components):
        die(f"toolchain.components must be a list of strings in {TOOLCHAIN_FILE}")
    if not isinstance(targets, list) or not all(isinstance(item, str) for item in targets):
        die(f"toolchain.targets must be a list of strings in {TOOLCHAIN_FILE}")

    return {
        "channel": channel,
        "profile": profile,
        "components": components,
        "targets": targets,
    }


@contextmanager
def file_lock(name: str):
    lock_root = rustup_home() / "tmp" / "zccache-locks"
    lock_root.mkdir(parents=True, exist_ok=True)
    lock_path = lock_root / f"{name}.lock"
    deadline = time.time() + LOCK_TIMEOUT_SECONDS
    while True:
        try:
            fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            try:
                os.write(fd, f"pid={os.getpid()}\ntime={time.time()}\n".encode("utf-8"))
            finally:
                os.close(fd)
            break
        except FileExistsError:
            if time.time() >= deadline:
                raise TimeoutError(f"timed out waiting for lock {lock_path}")
            time.sleep(LOCK_POLL_SECONDS)
    try:
        yield
    finally:
        try:
            lock_path.unlink()
        except FileNotFoundError:
            pass


def download(url: str, destination: Path) -> None:
    log(f"Downloading {url} ...")
    destination.parent.mkdir(parents=True, exist_ok=True)
    with urllib.request.urlopen(url) as response, destination.open("wb") as handle:
        shutil.copyfileobj(response, handle)


def install_uv() -> None:
    if shutil.which("uv"):
        log(f"uv already installed: {shutil.which('uv')}")
        return

    log("Installing uv ...")
    if os.name == "nt":
        if not shutil.which("powershell"):
            die("Cannot install uv: powershell not found.")
        subprocess.run(
            ["powershell", "-NoProfile", "-Command", "irm https://astral.sh/uv/install.ps1 | iex"],
            check=True,
        )
    else:
        tmp = Path(tempfile.mkdtemp(prefix="zccache-uv-install-")) / "install.sh"
        try:
            download("https://astral.sh/uv/install.sh", tmp)
            tmp.chmod(0o755)
            subprocess.run(["sh", str(tmp)], check=True)
        finally:
            shutil.rmtree(tmp.parent, ignore_errors=True)

    for directory in [Path.home() / ".cargo" / "bin", Path.home() / ".local" / "bin"]:
        if directory.is_dir():
            os.environ["PATH"] = str(directory) + os.pathsep + os.environ.get("PATH", "")

    if not shutil.which("uv"):
        warn("uv installed but not on PATH. Restart your terminal.")


def uv_sync() -> None:
    if shutil.which("uv"):
        log("Running uv sync ...")
        subprocess.run(["uv", "sync"], check=False)


def bootstrap_rustup() -> None:
    if rustup_executable():
        prepend_cargo_bin_to_path()
        return

    with file_lock("rustup-bootstrap"):
        if rustup_executable():
            prepend_cargo_bin_to_path()
            return

        suffix = ".exe" if os.name == "nt" else ""
        temp_dir = Path(tempfile.mkdtemp(prefix="zccache-rustup-init-"))
        temp_path = temp_dir / f"rustup-init{suffix}"
        download(rustup_init_url(), temp_path)
        try:
            if os.name != "nt":
                temp_path.chmod(0o755)
            subprocess.run(
                [
                    str(temp_path),
                    "-y",
                    "--profile",
                    "minimal",
                    "--default-toolchain",
                    "none",
                    "--default-host",
                    host_target_triple(),
                    "--no-modify-path",
                ],
                check=True,
                env=rustup_env(),
            )
        finally:
            shutil.rmtree(temp_dir, ignore_errors=True)

    prepend_cargo_bin_to_path()


def ensure_shared_rust() -> None:
    spec = load_toolchain_spec()
    bootstrap_rustup()
    rustup = rustup_executable()
    if not rustup:
        die("rustup bootstrap completed but rustup is still unavailable")

    cmd = [
        rustup,
        "toolchain",
        "install",
        str(spec["channel"]),
        "--profile",
        str(spec["profile"]),
        "--no-self-update",
    ]
    for component in spec["components"]:
        cmd.extend(["-c", str(component)])
    for target in spec["targets"]:
        cmd.extend(["-t", str(target)])

    lock_name = f"toolchain-{str(spec['channel']).replace('/', '-').replace(':', '-')}"
    with file_lock(lock_name):
        subprocess.run(cmd, check=True, env=rustup_env())


def verify() -> bool:
    prepend_cargo_bin_to_path()
    if not shutil.which("rustc"):
        err("rustc not found on PATH after installation.")
        err("Restart your terminal, then run: rustc --version")
        return False

    log("")
    log(f"rustc:   {subprocess.check_output(['rustc', '--version'], text=True).strip()}")
    log(f"cargo:   {subprocess.check_output(['cargo', '--version'], text=True).strip()}")

    if shutil.which("clippy-driver"):
        log("clippy:  available")
    else:
        warn("clippy: not found")

    if shutil.which("rustfmt"):
        log("rustfmt: available")
    else:
        warn("rustfmt: not found")

    if shutil.which("uv"):
        log(f"uv:      {subprocess.check_output(['uv', '--version'], text=True).strip()}")
    else:
        warn("uv: not found on PATH")

    rustc_path = shutil.which("rustc") or ""
    if ".cargo" not in rustc_path and "cargo" not in rustc_path:
        warn("")
        warn(f"WARNING: rustc at '{rustc_path}' is NOT from rustup.")
        warn("An old system Rust may shadow the rustup toolchain.")
        warn(f'  export CARGO_HOME="{cargo_home()}"')
        warn(f'  export RUSTUP_HOME="{rustup_home()}"')
        warn('  export PATH="$CARGO_HOME/bin:$PATH"')

    log("")
    log("Done. Run: soldr cargo check --workspace")
    return True


def main() -> None:
    spec = load_toolchain_spec()
    log("zccache - toolchain installer (uv + Rust)")
    log(f"Pinned: Rust {spec['channel']}")
    log(f"Cargo home:  {cargo_home()}")
    log(f"Rustup home: {rustup_home()}")
    log(f"Platform:    {host_target_triple()}")
    log("")

    install_uv()
    uv_sync()
    ensure_shared_rust()

    if not verify():
        warn("")
        warn("Restart your terminal, then run:")
        warn("  rustc --version")
        warn("  soldr cargo check --workspace")


if __name__ == "__main__":
    main()
