"""Build every demo crate (or the named ones) into pkg/<demo>/index.js + index_bg.wasm.

Run from anywhere:  python3 scripts/build.py [demo ...] [--dev]

Each crate under crates/<category>/<demo>/ is its own wasm-pack build, so pages only
download the code they use and special build flags stay local to one crate. A crate can
declare them in its Cargo.toml:

    [package.metadata.webtt]
    rustflags = "-C target-feature=+simd128"   # extra RUSTFLAGS for this crate only
    toolchain = "nightly"                      # run wasm-pack via `rustup run <toolchain>`
    wasm-pack-args = ["--", "-Z", "build-std=panic_abort,std"]
    raw = true                                 # no wasm-bindgen: plain `cargo build`, copy the .wasm

    [[package.metadata.webtt.variants]]         # extra builds of the same crate → pkg/<name>/
    name = "pathtracer-mt"
    toolchain = "nightly"
    needs-rust-src = true                      # -Z build-std: skipped with a hint if missing
    rustflags = "..."
    wasm-pack-args = ["--", "--features", "threads", "-Z", "build-std=panic_abort,std"]

Builds with their own toolchain or rustflags use target/<name>/, so they don't invalidate
the shared target/ cache of the others.
"""
import os
import pathlib
import shutil
import subprocess
import sys
import time
import tomllib

ROOT = pathlib.Path(__file__).resolve().parent.parent
CATEGORIES = ['graphics', 'simulation', 'algorithms', 'media', 'tools', 'games']


def discover():
    """{build name: (crate dir, metadata dict)} for every demo crate and its variants."""
    crates = {}
    for cat in CATEGORIES:
        for toml in sorted((ROOT / 'crates' / cat).glob('*/Cargo.toml')):
            data = tomllib.loads(toml.read_text(encoding='utf-8'))
            meta = dict(data.get('package', {}).get('metadata', {}).get('webtt', {}))
            variants = meta.pop('variants', [])
            crates[toml.parent.name] = (toml.parent, meta)
            for v in variants:
                crates[v['name']] = (toml.parent, {**v, 'variant': True})
    return crates


def missing_toolchain(meta):
    """Why this build can't run here, or None. Never triggers rustup's auto-install:
    only asks about a toolchain after `rustup toolchain list` shows it exists."""
    tc = meta.get('toolchain')
    if not tc:
        return None
    listed = subprocess.run(['rustup', 'toolchain', 'list'], capture_output=True, text=True).stdout
    if not any(line.split()[0] == tc or line.startswith(tc + '-') for line in listed.splitlines() if line.strip()):
        return f'toolchain "{tc}" is not installed (rustup toolchain install {tc})'
    if meta.get('needs-rust-src'):
        sysroot = subprocess.run(['rustup', 'run', tc, 'rustc', '--print', 'sysroot'], capture_output=True, text=True).stdout.strip()
        if not (pathlib.Path(sysroot) / 'lib' / 'rustlib' / 'src' / 'rust' / 'library').is_dir():
            return f'rust-src is missing for {tc} (rustup component add rust-src --toolchain {tc})'
    return None


def build(name, crate, meta, dev):
    out = ROOT / 'pkg' / name
    env = dict(os.environ)
    if meta.get('rustflags'):
        env['RUSTFLAGS'] = (env.get('RUSTFLAGS', '') + ' ' + meta['rustflags']).strip()
    if meta.get('rustflags') or meta.get('toolchain'):
        env['CARGO_TARGET_DIR'] = str(ROOT / 'target' / name)
    prefix = ['rustup', 'run', meta['toolchain']] if meta.get('toolchain') else []

    if meta.get('raw'):
        # Plain cdylib without wasm-bindgen (e.g. loaded inside an AudioWorklet)
        profile = 'debug' if dev else 'release'
        cmd = prefix + ['cargo', 'build', '-p', name, '--target', 'wasm32-unknown-unknown'] + ([] if dev else ['--release'])
        subprocess.run(cmd, cwd=ROOT, env=env, check=True)
        out.mkdir(parents=True, exist_ok=True)
        target_dir = pathlib.Path(env.get('CARGO_TARGET_DIR', ROOT / 'target'))
        artifact = target_dir / 'wasm32-unknown-unknown' / profile / f"{name.replace('-', '_')}.wasm"
        shutil.copy(artifact, out / 'index.wasm')
        # wasm-pack runs wasm-opt for the other crates; do the same here when it is installed
        if not dev and shutil.which('wasm-opt'):
            subprocess.run(['wasm-opt', '-Os', str(out / 'index.wasm'), '-o', str(out / 'index.wasm')], check=True)
        return out / 'index.wasm'

    cmd = prefix + ['wasm-pack', 'build', str(crate), '--target', 'web', '--no-pack',
                    '--out-dir', str(out), '--out-name', 'index'] + (['--dev'] if dev else [])
    cmd += meta.get('wasm-pack-args', [])
    subprocess.run(cmd, cwd=ROOT, env=env, check=True)
    return out / 'index_bg.wasm'


def main():
    args = [a for a in sys.argv[1:] if not a.startswith('--')]
    dev = '--dev' in sys.argv
    crates = discover()
    unknown = [a for a in args if a not in crates]
    if unknown:
        sys.exit(f"unknown demo(s): {', '.join(unknown)}\navailable: {', '.join(crates)}")
    names = args or list(crates)
    for name in names:
        crate, meta = crates[name]
        reason = missing_toolchain(meta)
        if reason:
            if args:  # asked for by name: fail loudly
                sys.exit(f'✗ {name}: {reason}')
            print(f'– {name:14} skipped: {reason}', flush=True)
            continue
        t0 = time.time()
        wasm = build(name, crate, meta, dev)
        print(f'✓ {name:14} {wasm.stat().st_size / 1024:6.1f} KB  ({time.time() - t0:.1f}s)', flush=True)


if __name__ == '__main__':
    main()