#!/usr/bin/env python3
# /// script
# requires-python = ">=3.9"
# dependencies = []
# ///
"""
Build Ray Docker images locally using raymake.
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

from ci.build.build_common import (
    BuildError,
    detect_host_arch,
    find_ray_root,
    get_git_commit,
    log,
    parse_file,
)

DEFAULT_ARCHITECTURE = "x86_64"
RAYCI_REGISTRY = "cr.ray.io/rayproject"

# Maps image type to its base name (YAML config key and Docker Hub repo).
_BASE_TYPE_MAP: dict[str, str] = {
    "ray": "ray",
    "ray-extra": "ray",
    "ray-llm": "ray-llm",
    "ray-llm-extra": "ray-llm",
}


def _load_ray_images() -> dict:
    path = find_ray_root() / "ray-images.json"
    if not path.exists():
        raise BuildError(f"Missing {path}")
    return json.loads(path.read_text())


def _build_image_type_config() -> dict[str, dict]:
    """Build IMAGE_TYPE_CONFIG from ray-images.yaml."""
    raw = _load_ray_images()
    config: dict[str, dict] = {}
    for image_type, yaml_key in _BASE_TYPE_MAP.items():
        yaml_cfg = raw[yaml_key]
        defaults = yaml_cfg["defaults"]
        platforms = yaml_cfg["platforms"]
        config[image_type] = {
            "python_versions": yaml_cfg["python"],
            "platforms": platforms,
            "architectures": yaml_cfg["architectures"],
            "default_python": defaults["python"],
            "default_platform": "cpu"
            if "cpu" in platforms
            else defaults["gpu_platform"],
        }
    return config


IMAGE_TYPE_CONFIG: dict[str, dict] = _build_image_type_config()


class RayImageError(Exception):
    """Raised when a RayImage field combination is invalid."""


@dataclass(frozen=True)
class RayImage:
    """Immutable identity of a Ray Docker image variant."""

    image_type: str
    python_version: str
    platform: str
    architecture: str = DEFAULT_ARCHITECTURE

    @property
    def wanda_image_name(self) -> str:
        """Wanda output image name (without registry prefix)."""
        if self.platform == "cpu":
            return f"{self.image_type}-py{self.python_version}-cpu{self.arch_suffix}"
        return f"{self.image_type}-py{self.python_version}-{self.platform}{self.arch_suffix}"

    @property
    def arch_suffix(self) -> str:
        if self.architecture == DEFAULT_ARCHITECTURE:
            return ""
        return f"-{self.architecture}"

    @property
    def repo(self) -> str:
        return _BASE_TYPE_MAP[self.image_type]

    @property
    def variation_suffix(self) -> str:
        if self.image_type.endswith("-extra"):
            return "-extra"
        return ""

    @property
    def platform_class(self) -> str:
        """Returns 'cpu', 'cuda', 'tpu', etc."""
        if self.platform == "cpu":
            return "cpu"
        if self.platform == "tpu":
            return "tpu"
        if self.platform.startswith("cu"):
            return "cuda"
        raise RayImageError(f"Could not determine platform class for {self.platform!r}")

    def get_wanda_spec_path(self) -> str:
        """Returns the path to the wanda spec file for this image."""
        key: WandaSpecKey = (self.image_type, self.platform_class)
        try:
            return WANDA_SPEC_PATHS[key]
        except KeyError:
            raise RayImageError(
                f"No wanda spec for image_type={self.image_type!r}, "
                f"platform={self.platform!r}"
            ) from None

    def validate(self) -> None:
        if self.image_type not in IMAGE_TYPE_CONFIG:
            valid = ", ".join(IMAGE_TYPE_CONFIG)
            raise RayImageError(
                f"Unknown image type {self.image_type!r}. Valid types: {valid}"
            )
        cfg = IMAGE_TYPE_CONFIG[self.image_type]
        if self.python_version not in cfg["python_versions"]:
            raise RayImageError(
                f"Invalid python version {self.python_version} "
                f"for {self.image_type}. "
                f"Valid versions: {', '.join(cfg['python_versions'])}"
            )
        if self.platform not in cfg["platforms"]:
            raise RayImageError(
                f"Invalid platform {self.platform} "
                f"for {self.image_type}. "
                f"Valid platforms: {', '.join(cfg['platforms'])}"
            )
        if self.architecture not in cfg["architectures"]:
            raise RayImageError(
                f"Invalid architecture {self.architecture} "
                f"for {self.image_type}. "
                f"Valid architectures: {', '.join(cfg['architectures'])}"
            )


# (image_type, platform_class) where platform_class is "cpu" or "cuda".
WandaSpecKey = tuple[str, str]

WANDA_SPEC_PATHS: dict[WandaSpecKey, str] = {
    ("ray", "cpu"): "ci/docker/ray-image-cpu.wanda.yaml",
    ("ray", "cuda"): "ci/docker/ray-image-cuda.wanda.yaml",
    ("ray", "tpu"): "ci/docker/ray-image-tpu.wanda.yaml",
    ("ray-extra", "cpu"): "ci/docker/ray-extra-image-cpu.wanda.yaml",
    ("ray-extra", "cuda"): "ci/docker/ray-extra-image-cuda.wanda.yaml",
    ("ray-llm", "cuda"): "ci/docker/ray-llm-image-cuda.wanda.yaml",
    ("ray-llm-extra", "cuda"): "ci/docker/ray-llm-extra-image-cuda.wanda.yaml",
}

SUPPORTED_IMAGE_TYPES: list[str] = list(dict.fromkeys(t for t, _ in WANDA_SPEC_PATHS))
REGISTRY_PREFIX = "cr.ray.io/rayproject/"


@dataclass(frozen=True)
class ImageBuildConfig:
    ray_image: RayImage
    ray_root: Path
    raymake_version: str
    manylinux_version: str
    ray_version: str
    commit: str
    wanda_spec_path: str
    wanda_image_tag: str
    nightly_alias: str | None
    build_env: dict[str, str]

    @classmethod
    def from_args(
        cls,
        image_type: str,
        python_version: str,
        image_platform: str,
        *,
        no_java: bool = False,
        no_dashboard: bool = False,
    ) -> ImageBuildConfig:
        ray_image = RayImage(
            image_type=image_type,
            python_version=python_version,
            platform=image_platform,
            architecture=cls._detect_host_arch(),
        )
        try:
            ray_image.validate()
            ray_image.get_wanda_spec_path()
        except RayImageError as e:
            raise BuildError(str(e)) from e

        root = find_ray_root()
        manylinux_version = parse_file(
            root / "rayci.env", r'MANYLINUX_VERSION=["\']?([^"\'\s]+)'
        )
        ray_version = parse_file(root / "rayci.env", r'RAY_VERSION=["\']?([^"\'\s]+)')
        commit = get_git_commit(root)

        cfg = IMAGE_TYPE_CONFIG[image_type]
        nightly_alias = (
            f"{REGISTRY_PREFIX}{ray_image.repo}:nightly{ray_image.variation_suffix}"
            if (
                ray_image.python_version == cfg["default_python"]
                and ray_image.platform == cfg["default_platform"]
                and ray_image.architecture == DEFAULT_ARCHITECTURE
            )
            else None
        )

        java_image = (
            "scratch"
            if no_java
            else f"{RAYCI_REGISTRY}/ray-java-build{ray_image.arch_suffix}"
        )
        dashboard_image = (
            "scratch"
            if no_dashboard
            else f"{RAYCI_REGISTRY}/ray-dashboard{ray_image.arch_suffix}"
        )

        build_env: dict[str, str] = {
            "PYTHON_VERSION": ray_image.python_version,
            "MANYLINUX_VERSION": manylinux_version,
            "HOSTTYPE": ray_image.architecture,
            "ARCH_SUFFIX": ray_image.arch_suffix,
            "JDK_SUFFIX": "" if no_java else "-jdk",
            "BUILDKITE_COMMIT": commit,
            "RAY_VERSION": ray_version,
            "IS_LOCAL_BUILD": "true",
            "IMAGE_TYPE": ray_image.repo,
            "RAY_JAVA_IMAGE": java_image,
            "RAY_DASHBOARD_IMAGE": dashboard_image,
        }
        if ray_image.platform.startswith("cu"):
            build_env["CUDA_VERSION"] = ray_image.platform.removeprefix("cu")

        return cls(
            ray_image=ray_image,
            ray_root=root,
            raymake_version=(root / ".rayciversion").read_text().strip(),
            manylinux_version=manylinux_version,
            ray_version=ray_version,
            commit=commit,
            wanda_spec_path=ray_image.get_wanda_spec_path(),
            wanda_image_tag=f"{REGISTRY_PREFIX}{ray_image.wanda_image_name}",
            nightly_alias=nightly_alias,
            build_env=build_env,
        )

    @staticmethod
    def _detect_host_arch() -> str:
        return detect_host_arch()


class ImageBuilder:
    def __init__(self, config: ImageBuildConfig):
        self.config = config

    def build(self) -> str:
        """Build the image and return the primary image tag."""
        if not shutil.which("raymake"):
            raise BuildError("raymake not found. Run via ./build-image.sh")

        log.info("Build configuration:")
        summary = {
            "Image Type": self.config.ray_image.image_type,
            "Python": self.config.ray_image.python_version,
            "Platform": self.config.ray_image.platform,
            "Arch": self.config.ray_image.architecture,
            "Commit": self.config.commit,
            "Raymake": self.config.raymake_version,
            "Ray Version": self.config.ray_version,
            "Wanda Spec": self.config.wanda_spec_path,
        }
        print("-" * 50)
        for k, v in summary.items():
            print(f"{k:<12}: {v}")
        print("-" * 50)

        cmd = [
            "raymake",
            str(self.config.wanda_spec_path),
        ]

        log.info(f"Running raymake: {self.config.wanda_spec_path}")
        log.info(f"Build environment: {self.config.build_env}")
        proc = subprocess.Popen(
            cmd,
            cwd=self.config.ray_root,
            env={**os.environ, **self.config.build_env},
        )
        try:
            if proc.wait() != 0:
                raise BuildError("raymake failed")
        except KeyboardInterrupt:
            try:
                proc.wait(timeout=2)
            except subprocess.TimeoutExpired:
                proc.kill()
            raise

        image_tag = self.config.wanda_image_tag
        log.info(f"Built image: {image_tag}")

        alias = self.config.nightly_alias
        if alias:
            self._docker_tag(image_tag, alias)
            log.info(f"Tagged alias: {alias}")

        return image_tag

    @staticmethod
    def _docker_tag(source: str, dest: str) -> None:
        result = subprocess.run(
            ["docker", "tag", source, dest],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            raise BuildError(f"docker tag failed: {result.stderr.strip()}")


def _image_type_table() -> str:
    """Build an ASCII table of all image types from IMAGE_TYPE_CONFIG."""

    def _tag_default(items, default):
        return [f"{v} (default)" if v == default else v for v in items]

    def _wrap(items, width):
        """Join items with ', ' and wrap into lines that fit in *width*."""
        lines: list[str] = []
        cur = ""
        for item in items:
            entry = item if not cur else f", {item}"
            if cur and len(cur) + len(entry) > width:
                lines.append(cur)
                cur = item
            else:
                cur += entry
        if cur:
            lines.append(cur)
        return lines

    headers = ("IMAGE TYPE", "PYTHON VERSIONS (-p)", "PLATFORMS (--platform)")
    plat_width = 52  # wrap platforms to this width

    def _build_rows(type_list):
        rows: list[tuple[str, str, list[str]]] = []
        for name in type_list:
            cfg = IMAGE_TYPE_CONFIG[name]
            py_cell = ", ".join(
                _tag_default(cfg["python_versions"], cfg["default_python"])
            )
            plat_lines = _wrap(
                _tag_default(cfg["platforms"], cfg["default_platform"]),
                plat_width,
            )
            rows.append((name, py_cell, plat_lines))
        return rows

    rows = _build_rows(SUPPORTED_IMAGE_TYPES)

    # Compute column widths.
    widths = [len(h) for h in headers]
    for label, py_cell, plat_lines in rows:
        widths[0] = max(widths[0], len(label))
        widths[1] = max(widths[1], len(py_cell))
        widths[2] = max(widths[2], *(len(l) for l in plat_lines))

    sep = "  "

    def fmt(cells):
        return sep.join(c.ljust(widths[i]) for i, c in enumerate(cells))

    lines = [fmt(headers), sep.join("-" * w for w in widths)]
    for label, py_cell, plat_lines in rows:
        lines.append(fmt((label, py_cell, plat_lines[0])))
        for cont in plat_lines[1:]:
            lines.append(fmt(("", "", cont)))

    return "\n".join(lines)


def _nightly_table() -> str:
    """Build a table mapping build commands to nightly docker run commands."""
    headers = ("BUILD", "RUN")
    rows = []
    for name in SUPPORTED_IMAGE_TYPES:
        cfg = IMAGE_TYPE_CONFIG[name]
        img = RayImage(name, cfg["default_python"], cfg["default_platform"])
        tag = f"{REGISTRY_PREFIX}{img.repo}:nightly{img.variation_suffix}"
        rows.append((f"./build-image.sh {name}", f"docker run -it {tag}"))

    widths = [len(h) for h in headers]
    for row in rows:
        for i, cell in enumerate(row):
            widths[i] = max(widths[i], len(cell))

    sep = "  "

    def fmt(cells):
        return sep.join(c.ljust(widths[i]) for i, c in enumerate(cells))

    lines = [fmt(headers), sep.join("-" * w for w in widths)]
    lines.extend(fmt(row) for row in rows)
    return "\n".join(lines)


def _build_examples() -> str:
    """Generate non-default example commands from IMAGE_TYPE_CONFIG."""
    ray_cfg = IMAGE_TYPE_CONFIG["ray"]
    ray_py = ray_cfg["default_python"]
    ray_plat = ray_cfg["default_platform"]
    alt_py = next((v for v in ray_cfg["python_versions"] if v != ray_py), None)
    alt_plat = next((p for p in ray_cfg["platforms"] if p != ray_plat), None)

    examples: list[tuple[str, str]] = []
    if alt_py:
        examples.append(
            (f"ray -p {alt_py}", f"ray {ray_plat} py{alt_py}"),
        )
    if alt_plat:
        examples.append(
            (f"ray --platform {alt_plat}", f"ray {alt_plat} py{ray_py}"),
        )
    if alt_py and alt_plat:
        examples.append(
            (f"ray -p {alt_py} --platform {alt_plat}", f"ray {alt_plat} py{alt_py}"),
        )

    if not examples:
        return ""

    cmd_prefix = "  ./build-image.sh "
    cmds = [cmd_prefix + args for args, _ in examples]
    max_cmd = max(len(c) for c in cmds)
    lines = []
    for cmd, (_, comment) in zip(cmds, examples):
        lines.append(f"{cmd:<{max_cmd}}  # {comment}")
    return "\n".join(lines)


def main():
    epilog_parts = [
        "Supported image types:",
        _image_type_table(),
        "",
        "Nightly images (built with defaults):",
        _nightly_table(),
        "",
        "Examples:",
        _build_examples(),
    ]

    parser = argparse.ArgumentParser(
        prog="./build-image.sh",
        description="Build Ray Docker images locally using raymake.",
        epilog="\n".join(epilog_parts),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "image_type",
        choices=SUPPORTED_IMAGE_TYPES,
        metavar="IMAGE_TYPE",
        help=f"Image type: {', '.join(SUPPORTED_IMAGE_TYPES)}",
    )
    parser.add_argument(
        "-p",
        "--python-version",
        default=None,
        metavar="VERSION",
        help="Python version (default depends on image type; see table below)",
    )
    parser.add_argument(
        "--platform",
        default=None,
        metavar="PLATFORM",
        help="Target platform (default depends on image type; see table below)",
    )
    parser.add_argument(
        "--list-platforms",
        action="store_true",
        help="List valid platforms for the given image type and exit",
    )
    parser.add_argument(
        "--no-java",
        action="store_true",
        default=False,
        help="Build wheel without Java JARs",
    )
    parser.add_argument(
        "--no-dashboard",
        action="store_true",
        default=False,
        help="Build wheel without the dashboard",
    )

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)

    args = parser.parse_args()

    cfg = IMAGE_TYPE_CONFIG[args.image_type]

    if args.list_platforms:
        print(f"Valid platforms for {args.image_type}:")
        for p in cfg["platforms"]:
            marker = " (default)" if p == cfg["default_platform"] else ""
            print(f"  {p}{marker}")
        sys.exit(0)

    python_version = args.python_version or cfg["default_python"]
    platform_choice = args.platform or cfg["default_platform"]

    try:
        config = ImageBuildConfig.from_args(
            args.image_type,
            python_version,
            platform_choice,
            no_java=args.no_java,
            no_dashboard=args.no_dashboard,
        )
        builder = ImageBuilder(config)
        image_tag = builder.build()

        print()
        log.info("Success!")
        print(f"  docker run -it {image_tag}")
        alias = config.nightly_alias
        if alias:
            print(f"  docker run -it {alias}")
    except BuildError as e:
        log.error(e)
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)


if __name__ == "__main__":
    main()