#!/usr/bin/env python3

"""Validate the Hygon DCU Python environment."""

from __future__ import annotations

import sys
from importlib import metadata


EXACT_VERSIONS = {
    "numpy": "1.26.4",
}

DAS_VERSIONS = {
    "torch": "2.9.0",
    "triton": "3.3.0",
    "flash_attn": "2.6.1",
    "flash_mla": "1.2.0",
    "lightop": "0.6.0",
    "lmslim": "0.3.1",
}


def main() -> int:
    errors: list[str] = []

    print("Checking Hygon DCU Python environment...")

    for pkg_name, expected_version in EXACT_VERSIONS.items():
        try:
            installed_ver = metadata.version(pkg_name)
            print(f"  {pkg_name}: {installed_ver}")
        except metadata.PackageNotFoundError:
            errors.append(f"- Missing `{pkg_name}` (expected `{expected_version}`).")
            continue
        if installed_ver != expected_version:
            errors.append(
                f"- `{pkg_name}` mismatch: expected `{expected_version}`, "
                f"got `{installed_ver}`."
            )

    for pkg_name, expected_base in DAS_VERSIONS.items():
        try:
            installed_ver = metadata.version(pkg_name)
            print(f"  {pkg_name}: {installed_ver}")
        except metadata.PackageNotFoundError:
            errors.append(
                f"- Missing `{pkg_name}` (expected Hygon build based on "
                f"`{expected_base}`)."
            )
            continue
        if not installed_ver.startswith(expected_base):
            errors.append(
                f"- `{pkg_name}` mismatch: expected base `{expected_base}`, "
                f"got `{installed_ver}`."
            )
        if "+das" not in installed_ver:
            errors.append(
                f"- `{pkg_name}` is not a Hygon custom build: "
                f"`{installed_ver}`."
            )

    if errors:
        print("\nEnvironment check failed:")
        for error in errors:
            print(error)
        return 1

    print("\nEnvironment check passed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())