from __future__ import annotations

import importlib.util
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


PATCH_DIR = Path(os.environ.get("MODAL_PATCH_DIR", "/root/patches"))


@dataclass(frozen=True)
class PatchSpec:
    name: str
    module: str
    patch_file: str
    strip: int
    includes: tuple[str, ...] = ()


PATCHES = (
    PatchSpec(
        name="flashinfer-pr-3312",
        module="flashinfer",
        patch_file="flashinfer-pr-3312.patch",
        strip=1,
    ),
)


def _package_parent(module_name: str) -> Path:
    spec = importlib.util.find_spec(module_name)
    if spec is None or spec.submodule_search_locations is None:
        raise RuntimeError(f"Could not find installed package {module_name!r}.")
    locations = list(spec.submodule_search_locations)
    if not locations:
        raise RuntimeError(f"Installed package {module_name!r} has no package path.")
    return Path(locations[0]).resolve().parent


def _git_apply_command(spec: PatchSpec, patch_path: Path) -> list[str]:
    cmd = ["git", "apply", f"-p{spec.strip}"]
    for include in spec.includes:
        cmd.append(f"--include={include}")
    cmd.append(str(patch_path))
    return cmd


def _check(cmd: list[str], *, cwd: Path) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        cmd,
        cwd=cwd,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )


def _apply_patch(spec: PatchSpec) -> None:
    patch_path = PATCH_DIR / spec.patch_file
    if not patch_path.exists():
        raise RuntimeError(f"Missing patch file: {patch_path}")

    cwd = _package_parent(spec.module)
    base_cmd = _git_apply_command(spec, patch_path)
    reverse_cmd = [*base_cmd[:2], "--reverse", "--check", *base_cmd[2:]]
    check_cmd = [*base_cmd[:2], "--check", *base_cmd[2:]]

    reverse = _check(reverse_cmd, cwd=cwd)
    if reverse.returncode == 0:
        print(f"[patch] {spec.name} already applied under {cwd}")
        return

    check = _check(check_cmd, cwd=cwd)
    if check.returncode != 0:
        print(check.stdout, file=sys.stderr)
        raise RuntimeError(f"Patch {spec.name} does not apply under {cwd}.")

    print(f"[patch] applying {spec.name} under {cwd}")
    subprocess.run(base_cmd, cwd=cwd, check=True)


def main() -> None:
    for patch in PATCHES:
        _apply_patch(patch)


if __name__ == "__main__":
    main()