#!/usr/bin/env python3

from __future__ import annotations

import os
import shutil
import subprocess
import sys
import textwrap
from pathlib import Path


EXPORT = Path(os.environ.get(
    "EXPORT",
    "/home/metris/hf-releases/Multiscreen-8.7M-Base",
)).resolve()

RUNTIME = EXPORT / "runtime"
OUTPUT = EXPORT / "runtime-minimal"
TRACE_DIR = EXPORT / ".runtime-trace"

CPU_RUNNER = EXPORT / "run_cpu_without_triton.py"
GPU_RUNNER = EXPORT / "run_gpu_triton.py"


def fail(message: str) -> None:
    raise SystemExit(f"ERROR: {message}")


if not RUNTIME.is_dir():
    fail(f"Missing runtime directory: {RUNTIME}")

if not CPU_RUNNER.is_file():
    fail(f"Missing CPU runner: {CPU_RUNNER}")

if not GPU_RUNNER.is_file():
    fail(f"Missing GPU runner: {GPU_RUNNER}")


shutil.rmtree(TRACE_DIR, ignore_errors=True)
shutil.rmtree(OUTPUT, ignore_errors=True)

TRACE_DIR.mkdir(parents=True)
OUTPUT.mkdir(parents=True)

sitecustomize = TRACE_DIR / "sitecustomize.py"

sitecustomize.write_text(
    textwrap.dedent(
        r'''
        from __future__ import annotations

        import atexit
        import os
        import sys
        from pathlib import Path


        ROOT = Path(os.environ["TRACE_RUNTIME_ROOT"]).resolve()
        OUTPUT = Path(os.environ["TRACE_OUTPUT"]).resolve()


        def source_path(module):
            candidates = []

            spec = getattr(module, "__spec__", None)
            if spec is not None:
                origin = getattr(spec, "origin", None)
                if origin:
                    candidates.append(origin)

            module_file = getattr(module, "__file__", None)
            if module_file:
                candidates.append(module_file)

            for candidate in candidates:
                try:
                    path = Path(candidate).resolve()
                except Exception:
                    continue

                if path.suffix == ".pyc":
                    cache_parts = list(path.parts)
                    if "__pycache__" in cache_parts:
                        cache_index = cache_parts.index("__pycache__")
                        parent = Path(*cache_parts[:cache_index])
                        stem = path.name.split(".cpython-")[0]
                        possible_source = parent / f"{stem}.py"
                        if possible_source.exists():
                            path = possible_source.resolve()

                if path.suffix != ".py":
                    continue

                try:
                    path.relative_to(ROOT)
                except ValueError:
                    continue

                return path

            return None


        @atexit.register
        def write_trace():
            found = set()

            for module in list(sys.modules.values()):
                if module is None:
                    continue

                path = source_path(module)
                if path is None:
                    continue

                found.add(path.relative_to(ROOT).as_posix())

            OUTPUT.parent.mkdir(parents=True, exist_ok=True)
            OUTPUT.write_text(
                "".join(f"{item}\n" for item in sorted(found)),
                encoding="utf-8",
            )
        '''
    ).strip()
    + "\n",
    encoding="utf-8",
)


def run_trace(name: str, runner: Path, extra_env: dict[str, str]) -> Path:
    trace_file = TRACE_DIR / f"{name}.txt"

    env = os.environ.copy()
    env.update(extra_env)
    env["TRACE_RUNTIME_ROOT"] = str(RUNTIME)
    env["TRACE_OUTPUT"] = str(trace_file)

    existing_pythonpath = env.get("PYTHONPATH", "")
    pythonpath_parts = [
        str(TRACE_DIR),
        str(RUNTIME),
    ]

    if existing_pythonpath:
        pythonpath_parts.append(existing_pythonpath)

    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)

    print(f"\n===== TRACING {name.upper()} =====")
    print(f"Runner: {runner}")

    subprocess.run(
        [sys.executable, str(runner)],
        cwd=EXPORT,
        env=env,
        check=True,
    )

    if not trace_file.is_file():
        fail(f"Trace was not created: {trace_file}")

    return trace_file


cpu_trace = run_trace(
    "cpu",
    CPU_RUNNER,
    {
        "CUDA_VISIBLE_DEVICES": "",
        "MULTISCREEN_BACKEND": "torch",
    },
)

gpu_trace = run_trace(
    "gpu",
    GPU_RUNNER,
    {
        "MULTISCREEN_BACKEND": "triton",
    },
)


required: set[Path] = set()

for trace_file in (cpu_trace, gpu_trace):
    for line in trace_file.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        required.add(Path(line))


# Package __init__.py files may not always appear independently in traces.
for relative_path in list(required):
    parent = relative_path.parent

    while str(parent) not in ("", "."):
        init_file = parent / "__init__.py"

        if (RUNTIME / init_file).is_file():
            required.add(init_file)

        parent = parent.parent


# Explicitly preserve the custom backend modules even if a lazy path was not
# exercised by one particular smoke prompt.
for relative_path in (
    Path("litgpt/__init__.py"),
    Path("litgpt/multiscreen_triton.py"),
    Path("litgpt/multiscreen_projection_triton.py"),
):
    if (RUNTIME / relative_path).is_file():
        required.add(relative_path)


for relative_path in sorted(required):
    source = RUNTIME / relative_path

    if not source.is_file():
        continue

    destination = OUTPUT / relative_path
    destination.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, destination)


manifest = OUTPUT / "MINIMAL_RUNTIME_FILES.txt"
manifest.write_text(
    "".join(f"{path.as_posix()}\n" for path in sorted(required)),
    encoding="utf-8",
)

print("\n===== MINIMAL RUNTIME BUILT =====")
print(f"Location: {OUTPUT}")
print(f"Python files: {sum(1 for p in OUTPUT.rglob('*.py'))}")

subprocess.run(["du", "-sh", str(RUNTIME), str(OUTPUT)], check=False)

print("\nDo not delete the original runtime yet.")
print("Rename it, test the minimal runtime, and only then remove the backup.")