import argparse
import contextlib
import shutil
import subprocess
import sys
import zipfile
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class Package:
    # name of the package on PyPI.
    pypi_name: str
    # type of the package, one of "dev", "skinny", "tracing", "release"
    type: str
    # path to the package relative to the root of the repository
    build_path: str


DEV = Package("mlflow", "dev", ".")
RELEASE = Package("mlflow", "release", ".")
SKINNY = Package("mlflow-skinny", "skinny", "libs/skinny")
TRACING = Package("mlflow-tracing", "tracing", "libs/tracing")

PACKAGES = [
    DEV,
    SKINNY,
    RELEASE,
    TRACING,
]

JS_BUILD_DIR = Path("mlflow/server/js/build")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build MLflow package.")
    parser.add_argument(
        "--package-type",
        help="Package type to build. Default is 'dev'.",
        choices=[p.type for p in PACKAGES],
        default="dev",
    )
    parser.add_argument(
        "--sha",
        help="If specified, include the SHA in the wheel name as a build tag.",
    )
    return parser.parse_args()


@contextlib.contextmanager
def restore_changes() -> Generator[None, None, None]:
    try:
        yield
    finally:
        subprocess.check_call([
            "git",
            "restore",
            "README.md",
            "pyproject.toml",
        ])


def validate_ui_assets_pre_build(package: Package) -> None:
    if package != RELEASE:
        return
    if not JS_BUILD_DIR.exists() or not any(JS_BUILD_DIR.iterdir()):
        raise RuntimeError("Build the UI first before building the release package.")


def validate_ui_assets_post_build(wheel_path: Path, package: Package) -> None:
    ui_asset_prefix = f"{JS_BUILD_DIR.as_posix()}/"
    with zipfile.ZipFile(wheel_path) as zf:
        has_ui_assets = any(name.startswith(ui_asset_prefix) for name in zf.namelist())
    if package == RELEASE:
        if not has_ui_assets:
            raise RuntimeError(
                f"UI assets are missing from the release wheel: {wheel_path}. "
                "Build the UI first before building the release package."
            )
    elif package in (SKINNY, TRACING):
        if has_ui_assets:
            raise RuntimeError(
                f"UI assets should not be included in the {package.type} wheel: {wheel_path}."
            )


def main() -> None:
    args = parse_args()

    # Initialize submodules (e.g., mlflow/assistant/skills)
    subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"])

    # Clean up build artifacts generated by previous builds
    paths_to_clean_up = ["build"]
    for pkg in PACKAGES:
        paths_to_clean_up += [
            f"{pkg.build_path}/dist",
            f"{pkg.build_path}/{pkg.pypi_name}.egg_info",
        ]
    for path in map(Path, paths_to_clean_up):
        if not path.exists():
            continue
        if path.is_file():
            path.unlink()
        else:
            shutil.rmtree(path)

    package = next(p for p in PACKAGES if p.type == args.package_type)

    validate_ui_assets_pre_build(package)

    with restore_changes():
        pyproject = Path("pyproject.toml")
        if package == RELEASE:
            pyproject.write_text(Path("pyproject.release.toml").read_text())

        DIST_DIR = Path("dist").resolve()
        DIST_DIR.mkdir(exist_ok=True)
        subprocess.check_call([
            sys.executable,
            "-m",
            "build",
            package.build_path,
            "--outdir",
            DIST_DIR,
        ])

    wheel = next(DIST_DIR.glob("mlflow*.whl"))
    validate_ui_assets_post_build(wheel, package)

    if args.sha:
        name, version, rest = wheel.name.split("-", 2)
        build_tag = f"0.sha.{args.sha}"  # build tag must start with a digit
        wheel.rename(wheel.with_name(f"{name}-{version}-{build_tag}-{rest}"))


if __name__ == "__main__":
    main()