#!/usr/bin/env python3
"""Build, launch, and automatically restart the lightweight Blacksite UI gallery."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import tempfile
import time
from typing import Any, Iterable, Sequence


ROOT = Path(__file__).resolve().parents[2]
CARGO_LANE = ROOT / "scripts/codex/cargo_lane.py"
BUILD_STORAGE = ROOT / "scripts/codex/build_storage.py"
LOG_DIR = ROOT / ".codex/logs"
BUILD_GRAPH_CACHE = LOG_DIR / "ui-gallery-build-units.json"
METADATA_LIMIT_BYTES = 32 * 1024 * 1024
POLL_SECONDS = 0.25
DEBOUNCE_SECONDS = 0.35
WATCH_ROOTS = (
    ROOT / "crates/editor_ui/src",
    ROOT / "crates/material_schema/src",
    ROOT / "crates/ui_gallery/src",
)
WATCH_FILES = (
    ROOT / "Cargo.toml",
    ROOT / "Cargo.lock",
    ROOT / "crates/editor_ui/Cargo.toml",
    ROOT / "crates/material_schema/Cargo.toml",
    ROOT / "crates/ui_gallery/Cargo.toml",
)


def managed_command(*command: str) -> list[str]:
    return [sys.executable, str(CARGO_LANE), "exec", "dev", *command]


def storage_enforce(phase: str) -> int:
    return subprocess.run(
        [sys.executable, str(BUILD_STORAGE), "enforce", "--phase", phase],
        cwd=ROOT,
        check=False,
    ).returncode


def host_target() -> str:
    completed = subprocess.run(
        ["rustc", "-vV"],
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    for line in completed.stdout.splitlines():
        if line.startswith("host:"):
            return line.partition(":")[2].strip()
    return ""


def bounded_command_output(command: Sequence[str], limit: int) -> bytes:
    """Capture a known-small machine document without ever buffering unbounded stdout."""
    with tempfile.TemporaryFile() as output, tempfile.TemporaryFile() as errors:
        with subprocess.Popen(
            command,
            cwd=ROOT,
            stdout=subprocess.PIPE,
            stderr=errors,
        ) as process:
            assert process.stdout is not None
            size = 0
            while chunk := process.stdout.read(64 * 1024):
                size += len(chunk)
                if size > limit:
                    process.kill()
                    process.wait()
                    raise RuntimeError(
                        f"machine output exceeded the {limit // (1024 * 1024)} MiB safety cap"
                    )
                output.write(chunk)
            return_code = process.wait()
        if return_code != 0:
            errors.seek(0)
            detail = errors.read()[-4000:].decode("utf-8", errors="replace")
            raise RuntimeError(f"command failed ({return_code}): {detail}")
        output.seek(0)
        return output.read()


def build_graph_key() -> str:
    digest = hashlib.sha256()
    for path in WATCH_FILES:
        if path.is_file():
            digest.update(path.relative_to(ROOT).as_posix().encode("utf-8"))
            digest.update(path.read_bytes())
    digest.update(host_target().encode("utf-8"))
    return digest.hexdigest()


def cached_build_graph(key: str) -> set[str] | None:
    try:
        document = json.loads(BUILD_GRAPH_CACHE.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    package_ids = document.get("package_ids")
    if document.get("key") != key or not isinstance(package_ids, list):
        return None
    if not all(isinstance(package_id, str) for package_id in package_ids):
        return None
    return set(package_ids)


def save_build_graph(key: str, package_ids: set[str]) -> None:
    BUILD_GRAPH_CACHE.parent.mkdir(parents=True, exist_ok=True)
    BUILD_GRAPH_CACHE.write_text(
        json.dumps(
            {"key": key, "package_ids": sorted(package_ids)},
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )


def compiler_artifact(message: object) -> tuple[str, str, bool] | None:
    if not isinstance(message, dict) or message.get("reason") != "compiler-artifact":
        return None
    package_id = message.get("package_id")
    if not isinstance(package_id, str):
        return None
    target = message.get("target")
    if isinstance(target, dict) and isinstance(target.get("name"), str):
        target_name = target["name"]
    else:
        target_name = package_id.rsplit("#", 1)[-1]
    return package_id, target_name, bool(message.get("fresh", False))


def diagnostic_lines(message: object, raw_line: str) -> Iterable[str]:
    if not isinstance(message, dict):
        yield raw_line
        return
    if message.get("reason") == "compiler-message":
        diagnostic = message.get("message")
        if isinstance(diagnostic, dict):
            rendered = diagnostic.get("rendered")
            if isinstance(rendered, str) and rendered:
                yield rendered.rstrip("\n")


def build_gallery() -> int:
    graph_key = build_graph_key()
    graph = cached_build_graph(graph_key)
    total = len(graph) if graph else None
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    log_path = LOG_DIR / "ui-gallery-build.log"
    command = managed_command(
        "cargo",
        "build",
        "-p",
        "ui_gallery",
        "--message-format=json-render-diagnostics",
        "--color",
        "never",
    )
    process = subprocess.Popen(
        command,
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
    )
    assert process.stdout is not None
    completed: set[str] = set()
    rebuilt: set[str] = set()
    last_progress_bucket = -1
    with log_path.open("w", encoding="utf-8") as log:
        for raw_line in process.stdout:
            stripped = raw_line.rstrip("\n")
            try:
                message: object = json.loads(stripped)
            except json.JSONDecodeError:
                message = stripped
            for diagnostic in diagnostic_lines(message, stripped):
                log.write(diagnostic + "\n")
                log.flush()
            artifact = compiler_artifact(message)
            if artifact is None:
                continue
            package_id, target_name, fresh = artifact
            if not fresh:
                rebuilt.add(package_id)
            if package_id in completed:
                continue
            completed.add(package_id)
            if graph is None:
                in_graph = len(completed)
                if in_graph != 1 and in_graph % 5 != 0:
                    continue
                counter = f"{in_graph}/? discovering · {len(rebuilt)} rebuilt"
            else:
                in_graph = len(completed.intersection(graph))
                percent = min(100, (in_graph * 100) // max(1, total))
                progress_bucket = percent // 5
                if progress_bucket == last_progress_bucket and in_graph != total:
                    continue
                last_progress_bucket = progress_bucket
                counter = f"{min(in_graph, total)}/{total} resolved · {len(rebuilt)} rebuilt"
            print(
                f"\rResolving ui_gallery ({counter}) {target_name:<36}",
                end="",
                flush=True,
            )
    return_code = process.wait()
    print()
    if return_code == 0:
        save_build_graph(graph_key, completed)
        expected = len(graph) if graph else len(completed)
        resolved = len(completed.intersection(graph)) if graph else len(completed)
        print(
            f"ui_gallery ready ({resolved}/{expected} resolved · "
            f"{len(rebuilt)} rebuilt)"
        )
    else:
        print(f"ui_gallery build failed; see {log_path}", file=sys.stderr)
    return return_code


def lane_runtime_environment() -> dict[str, str]:
    completed = subprocess.run(
        [sys.executable, str(CARGO_LANE), "env", "--json", "dev"],
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        check=True,
    )
    document = json.loads(completed.stdout)
    environment = os.environ.copy()
    environment.update(document.get("environment", {}))
    runtime_deps = document.get("runtime_deps")
    if isinstance(runtime_deps, str) and runtime_deps:
        previous = environment.get("LD_LIBRARY_PATH")
        environment["LD_LIBRARY_PATH"] = (
            runtime_deps if not previous else runtime_deps + os.pathsep + previous
        )
    return environment


def gallery_binary() -> Path:
    completed = subprocess.run(
        [sys.executable, str(CARGO_LANE), "env", "--json", "dev"],
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        check=True,
    )
    document = json.loads(completed.stdout)
    return Path(document["target_dir"]) / "debug/ui_gallery"


def launch_gallery() -> subprocess.Popen[bytes]:
    binary = gallery_binary()
    if not binary.is_file():
        raise RuntimeError(f"gallery binary is missing after build: {binary}")
    return subprocess.Popen([str(binary)], cwd=ROOT, env=lane_runtime_environment())


def stop_gallery(process: subprocess.Popen[bytes] | None) -> None:
    if process is None or process.poll() is not None:
        return
    process.send_signal(signal.SIGTERM)
    try:
        process.wait(timeout=3.0)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait()


def rebuild_and_replace_gallery(
    current: subprocess.Popen[bytes] | None,
) -> subprocess.Popen[bytes] | None:
    """Keep the last-good gallery alive until its replacement has built and launched."""
    if build_gallery() != 0:
        return current
    replacement = launch_gallery()
    stop_gallery(current)
    return replacement


def source_snapshot() -> dict[Path, int]:
    files = list(WATCH_FILES)
    for root in WATCH_ROOTS:
        if root.is_dir():
            files.extend(root.rglob("*.rs"))
    return {
        path: path.stat().st_mtime_ns
        for path in files
        if path.is_file()
    }


def watch() -> int:
    if storage_enforce("pre") != 0:
        return 1
    child: subprocess.Popen[bytes] | None = None
    try:
        if build_gallery() != 0:
            return 1
        child = launch_gallery()
        print("Watching material_schema, editor_ui, and ui_gallery Rust sources (Ctrl+C to stop)")
        snapshot = source_snapshot()
        changed_at: float | None = None
        while True:
            time.sleep(POLL_SECONDS)
            current = source_snapshot()
            if current != snapshot:
                snapshot = current
                changed_at = time.monotonic()
            if changed_at is None or time.monotonic() - changed_at < DEBOUNCE_SECONDS:
                continue
            changed_at = None
            child = rebuild_and_replace_gallery(child)
    except KeyboardInterrupt:
        print("\nStopping ui_gallery")
        return 0
    finally:
        stop_gallery(child)
        storage_enforce("post")


def run_once(*, build_only: bool) -> int:
    if storage_enforce("pre") != 0:
        return 1
    try:
        result = build_gallery()
        if result != 0 or build_only:
            return result
        return launch_gallery().wait()
    finally:
        storage_enforce("post")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "command",
        choices=("watch", "run", "build"),
        nargs="?",
        default="watch",
    )
    arguments = parser.parse_args()
    if arguments.command == "watch":
        return watch()
    return run_once(build_only=arguments.command == "build")


if __name__ == "__main__":
    raise SystemExit(main())