"""Tests for Ascend env parity check: check_ascend_env_parity() and env check-ascend CLI.

Covers:
- No-toolkit path (ASCEND machines absent)
- Canonical layout (opp + runtime subdirs)
- Arch-specific layout (aarch64-linux/lib64, no top-level opp)
- Versioned toolkit directories (no 'latest' symlink)
- Plugin home discovery: latest symlink vs versioned fallback
- Per-variable status: ok / missing / abnormal / optional-absent
- CLI command integration (non-Ascend machine silently exits 0)
"""

from __future__ import annotations

from pathlib import Path

from sagellm.cli.utils import env_detector
from sagellm.cli.utils.env_detector import (
    ASCEND_CRITICAL_VARS,
    AscendVarCheck,
    check_ascend_env_parity,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _results_by_name(checks: list[AscendVarCheck]) -> dict[str, AscendVarCheck]:
    return {c.name: c for c in checks}


# ---------------------------------------------------------------------------
# _iter_ascend_toolkit_homes — layout compatibility
# ---------------------------------------------------------------------------


class TestIterAscendToolkitHomes:
    """_iter_ascend_toolkit_homes must discover homes across multiple layouts."""

    def test_returns_empty_when_toolkit_root_missing(self, tmp_path, monkeypatch):
        monkeypatch.setattr(
            env_detector,
            "_iter_ascend_toolkit_homes",
            lambda: [],
        )
        # Direct test: patch FS root
        import sagellm.cli.utils.env_detector as mod

        original = env_detector._iter_ascend_toolkit_homes
        # Reset to real function and patch Path
        monkeypatch.setattr(mod.Path, "__new__", Path.__new__)  # noop
        result = original()
        # On a non-Ascend machine the toolkit root is absent → empty
        # (may return [] or real results depending on machine; just check it's a list)
        assert isinstance(result, list)

    def test_canonical_layout(self, tmp_path, monkeypatch):
        """Layout: toolkit_root/ver/opp + runtime."""
        toolkit_root = tmp_path / "ascend-toolkit"
        ver = toolkit_root / "7.0.0.beta001"
        (ver / "opp").mkdir(parents=True)
        (ver / "runtime").mkdir()

        monkeypatch.setenv("ASCEND_HOME", "")  # clear env hints
        monkeypatch.setenv("ASCEND_HOME_PATH", "")
        monkeypatch.setenv("ASCEND_RUNTIME_PATH", "")

        import sagellm.cli.utils.env_detector as mod

        class _PatchedPath(type(tmp_path)):
            pass

        # Monkey-patch the toolkit root lookup directly
        def _patched():
            # Temporarily swap the hardcoded root
            root = toolkit_root
            if not root.exists():
                return []
            candidates = []
            seen: set[str] = set()

            def _append(p):
                if p is None:
                    return
                s = str(p)
                if s in seen or not p.is_dir():
                    return
                seen.add(s)
                candidates.append(p)

            _append(root / "latest")
            for child in sorted(root.iterdir(), reverse=True):
                if child.name == "latest":
                    continue
                _append(child)

            _arch_names = ("aarch64-linux", "x86_64-linux", "arm64-linux")
            return [
                h
                for h in candidates
                if (
                    (h / "opp").is_dir()
                    and any((h / d).is_dir() for d in ("runtime", "fwkacllib", "hccl", "lib64"))
                )
                or any((h / arch / "lib64").is_dir() for arch in _arch_names)
            ]

        monkeypatch.setattr(mod, "_iter_ascend_toolkit_homes", _patched)
        result = mod._iter_ascend_toolkit_homes()
        assert ver in result

    def test_arch_specific_layout(self, tmp_path, monkeypatch):
        """Layout: toolkit_root/ver/aarch64-linux/lib64 (no top-level opp)."""
        toolkit_root = tmp_path / "ascend-toolkit"
        ver = toolkit_root / "7.1.0.alpha001"
        (ver / "aarch64-linux" / "lib64").mkdir(parents=True)

        import sagellm.cli.utils.env_detector as mod

        def _patched():
            root = toolkit_root
            if not root.exists():
                return []
            candidates = []
            seen: set[str] = set()

            def _append(p):
                s = str(p)
                if s in seen or not p.is_dir():
                    return
                seen.add(s)
                candidates.append(p)

            for child in sorted(root.iterdir(), reverse=True):
                if child.name == "latest":
                    continue
                _append(child)

            _arch_names = ("aarch64-linux", "x86_64-linux", "arm64-linux")
            return [
                h
                for h in candidates
                if (
                    (h / "opp").is_dir()
                    and any((h / d).is_dir() for d in ("runtime", "fwkacllib", "hccl", "lib64"))
                )
                or any((h / arch / "lib64").is_dir() for arch in _arch_names)
            ]

        monkeypatch.setattr(mod, "_iter_ascend_toolkit_homes", _patched)
        result = mod._iter_ascend_toolkit_homes()
        assert ver in result

    def test_versioned_without_latest_symlink(self, tmp_path, monkeypatch):
        """Layout: multiple versioned dirs, no 'latest' symlink."""
        toolkit_root = tmp_path / "ascend-toolkit"
        ver_old = toolkit_root / "6.3.0"
        ver_new = toolkit_root / "7.0.0"
        for ver in (ver_old, ver_new):
            (ver / "opp").mkdir(parents=True)
            (ver / "runtime").mkdir()

        import sagellm.cli.utils.env_detector as mod

        def _patched():
            root = toolkit_root
            if not root.exists():
                return []
            candidates = []
            seen: set[str] = set()

            def _append(p):
                s = str(p)
                if s in seen or not p.is_dir():
                    return
                seen.add(s)
                candidates.append(p)

            _append(root / "latest")  # won't exist
            for child in sorted(root.iterdir(), reverse=True):
                if child.name == "latest":
                    continue
                _append(child)

            _arch_names = ("aarch64-linux", "x86_64-linux", "arm64-linux")
            return [
                h
                for h in candidates
                if (
                    (h / "opp").is_dir()
                    and any((h / d).is_dir() for d in ("runtime", "fwkacllib", "hccl", "lib64"))
                )
                or any((h / arch / "lib64").is_dir() for arch in _arch_names)
            ]

        monkeypatch.setattr(mod, "_iter_ascend_toolkit_homes", _patched)
        result = mod._iter_ascend_toolkit_homes()
        # Both are valid; newest should come first (sorted reverse)
        assert ver_new in result
        assert ver_old in result
        assert result.index(ver_new) < result.index(ver_old)


# ---------------------------------------------------------------------------
# _discover_ascend_plugin_home — versioned fallback
# ---------------------------------------------------------------------------


class TestDiscoverAscendPluginHome:
    def test_returns_none_when_root_missing(self, tmp_path, monkeypatch):
        import sagellm.cli.utils.env_detector as mod

        def _patched():
            root = tmp_path / "nonexistent"
            if not root.is_dir():
                return None
            latest = root / "latest"
            if latest.is_dir():
                return latest
            try:
                versioned = sorted([p for p in root.iterdir() if p.is_dir()], reverse=True)
            except OSError:
                return None
            return versioned[0] if versioned else None

        monkeypatch.setattr(mod, "_discover_ascend_plugin_home", _patched)
        assert mod._discover_ascend_plugin_home() is None

    def test_prefers_latest_symlink(self, tmp_path, monkeypatch):
        plugin_root = tmp_path / "tfplugin"
        latest = plugin_root / "latest"
        latest.mkdir(parents=True)
        versioned = plugin_root / "6.3.0"
        versioned.mkdir()

        import sagellm.cli.utils.env_detector as mod

        def _patched():
            if not plugin_root.is_dir():
                return None
            lat = plugin_root / "latest"
            if lat.is_dir():
                return lat
            try:
                v = sorted([p for p in plugin_root.iterdir() if p.is_dir()], reverse=True)
            except OSError:
                return None
            return v[0] if v else None

        monkeypatch.setattr(mod, "_discover_ascend_plugin_home", _patched)
        assert mod._discover_ascend_plugin_home() == latest

    def test_falls_back_to_versioned_dir(self, tmp_path, monkeypatch):
        plugin_root = tmp_path / "tfplugin"
        old = plugin_root / "5.0.0"
        old.mkdir(parents=True)
        new = plugin_root / "7.0.0"
        new.mkdir()

        import sagellm.cli.utils.env_detector as mod

        def _patched():
            if not plugin_root.is_dir():
                return None
            lat = plugin_root / "latest"
            if lat.is_dir():
                return lat
            try:
                v = sorted([p for p in plugin_root.iterdir() if p.is_dir()], reverse=True)
            except OSError:
                return None
            return v[0] if v else None

        monkeypatch.setattr(mod, "_discover_ascend_plugin_home", _patched)
        result = mod._discover_ascend_plugin_home()
        assert result == new  # lexicographically latest


# ---------------------------------------------------------------------------
# check_ascend_env_parity — core logic
# ---------------------------------------------------------------------------


class TestCheckAscendEnvParity:
    """Unit-tests for AscendVarCheck statuses."""

    def test_no_toolkit_returns_no_toolkit_for_all(self):
        """Passing toolkit_home=None → all vars get status='no-toolkit'."""
        results = check_ascend_env_parity(toolkit_home=None)
        assert len(results) == len(ASCEND_CRITICAL_VARS)
        for chk in results:
            assert chk.status == "no-toolkit"
            assert chk.repair is None

    def test_all_critical_vars_covered(self, tmp_path):
        """Result includes exactly the ASCEND_CRITICAL_VARS."""
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()
        results = check_ascend_env_parity(toolkit_home=tk)
        names = {c.name for c in results}
        assert names == set(ASCEND_CRITICAL_VARS)

    def test_missing_status_when_var_unset(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        # Ensure the vars are absent
        for v in ("ASCEND_HOME", "ASCEND_TOOLKIT_HOME", "ASCEND_RUNTIME_PATH"):
            monkeypatch.delenv(v, raising=False)

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        assert results["ASCEND_HOME"].status == "missing"
        assert results["ASCEND_TOOLKIT_HOME"].status == "missing"
        assert results["ASCEND_RUNTIME_PATH"].status == "missing"

    def test_ok_status_when_vars_match(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        monkeypatch.setenv("ASCEND_HOME", str(tk))
        monkeypatch.setenv("ASCEND_TOOLKIT_HOME", str(tk))
        monkeypatch.setenv("ASCEND_RUNTIME_PATH", str(tk / "runtime"))
        monkeypatch.setenv("ASCEND_OPP_PATH", str(tk / "opp"))
        monkeypatch.setenv("ASCEND_AICPU_PATH", str(tk))  # fallback value

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        assert results["ASCEND_HOME"].status == "ok"
        assert results["ASCEND_TOOLKIT_HOME"].status == "ok"
        assert results["ASCEND_RUNTIME_PATH"].status == "ok"
        assert results["ASCEND_OPP_PATH"].status == "ok"

    def test_abnormal_status_when_var_mismatches(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()
        wrong = tmp_path / "wrong_toolkit"
        wrong.mkdir()

        monkeypatch.setenv("ASCEND_HOME", str(wrong))

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_HOME"]
        assert chk.status == "abnormal"
        assert chk.repair is not None
        assert str(tk) in chk.repair

    def test_repair_suggestion_contains_export(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        monkeypatch.delenv("ASCEND_HOME", raising=False)

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_HOME"]
        assert chk.repair is not None
        assert chk.repair.startswith("export ASCEND_HOME=")

    def test_optional_absent_for_plugin_home_when_no_plugin(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        monkeypatch.delenv("ASCEND_PLUGIN_HOME", raising=False)

        # No plugin root exists → expected=None → optional-absent
        import sagellm.cli.utils.env_detector as mod

        monkeypatch.setattr(mod, "_discover_ascend_plugin_home", lambda: None)

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_PLUGIN_HOME"]
        assert chk.status == "optional-absent"

    def test_ld_library_path_abnormal_when_missing_segments(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "lib64").mkdir()
        (tk / "runtime" / "lib64").mkdir(parents=True)

        # Set LD_LIBRARY_PATH to something that does NOT include tk/lib64
        monkeypatch.setenv("LD_LIBRARY_PATH", "/some/other/path")

        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["LD_LIBRARY_PATH"]
        assert chk.status == "abnormal"
        assert chk.repair is not None
        # Repair should make paths appear before existing value
        assert "/some/other/path" in chk.repair

    def test_ld_library_path_ok_when_all_segments_present(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "lib64").mkdir()

        # Pre-build the expected segments
        from sagellm.cli.utils.env_detector import check_ascend_env_parity as _fn

        # First call with nothing set to know what's expected
        monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
        results_before = _results_by_name(_fn(toolkit_home=tk))
        expected = results_before["LD_LIBRARY_PATH"].expected or ""

        # Now set LD_LIBRARY_PATH to exactly those segments
        monkeypatch.setenv("LD_LIBRARY_PATH", expected)
        results_after = _results_by_name(_fn(toolkit_home=tk))
        assert results_after["LD_LIBRARY_PATH"].status == "ok"

    def test_aicpu_fallback_to_toolkit_home(self, tmp_path, monkeypatch):
        """If no aicpu_kernels dir exists, expected falls back to toolkit_home."""
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        monkeypatch.delenv("ASCEND_AICPU_PATH", raising=False)
        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_AICPU_PATH"]
        assert chk.expected == str(tk)

    def test_aicpu_uses_aicpu_kernels_when_present(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()
        (tk / "aicpu_kernels").mkdir()

        monkeypatch.delenv("ASCEND_AICPU_PATH", raising=False)
        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_AICPU_PATH"]
        assert chk.expected == str(tk / "aicpu_kernels")

    def test_aicpu_uses_arch_specific_opskernel(self, tmp_path, monkeypatch):
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()
        opskernel = tk / "aarch64-linux" / "lib64" / "plugin" / "opskernel"
        opskernel.mkdir(parents=True)

        monkeypatch.delenv("ASCEND_AICPU_PATH", raising=False)
        results = _results_by_name(check_ascend_env_parity(toolkit_home=tk))
        chk = results["ASCEND_AICPU_PATH"]
        assert chk.expected == str(opskernel)


# ---------------------------------------------------------------------------
# CLI integration: env check-ascend
# ---------------------------------------------------------------------------


class TestEnvCheckAscendCLI:
    def test_skips_silently_on_non_ascend_machine(self, monkeypatch):
        """On non-Ascend machine (no npu-smi), command exits 0 with a warning."""
        from click.testing import CliRunner

        import sagellm.cli.commands.env as env_cmd
        from sagellm.cli.commands.env import env

        monkeypatch.setattr(env_cmd.shutil, "which", lambda _cmd: None)

        runner = CliRunner()
        result = runner.invoke(env, ["check-ascend"])
        assert result.exit_code == 0
        assert "npu-smi not found" in result.output

    def test_exits_1_when_vars_missing(self, tmp_path, monkeypatch):
        """When npu-smi is found but vars are missing, exits 1."""
        from click.testing import CliRunner

        import sagellm.cli.commands.env as env_cmd
        import sagellm.cli.utils.env_detector as mod
        from sagellm.cli.commands.env import env

        # Fake Ascend machine
        monkeypatch.setattr(
            env_cmd.shutil, "which", lambda cmd: "/usr/bin/npu-smi" if cmd == "npu-smi" else None
        )

        # Canonical toolkit layout
        tk = tmp_path / "toolkit"
        (tk / "opp").mkdir(parents=True)
        (tk / "runtime").mkdir()

        # Clear all critical vars
        for v in ASCEND_CRITICAL_VARS:
            monkeypatch.delenv(v, raising=False)

        # Wire check_ascend_env_parity to use our fake toolkit (in the env command module)
        monkeypatch.setattr(mod, "_discover_ascend_plugin_home", lambda: None)
        monkeypatch.setattr(
            env_cmd,
            "check_ascend_env_parity",
            lambda: mod.check_ascend_env_parity(toolkit_home=tk),
        )

        runner = CliRunner()
        result = runner.invoke(env, ["check-ascend"])
        # Should exit 1 due to missing vars
        assert result.exit_code == 1

    def test_exits_0_when_all_ok(self, tmp_path, monkeypatch):
        """When all vars are correctly set, exits 0."""
        from click.testing import CliRunner

        import sagellm.cli.commands.env as env_cmd
        import sagellm.cli.utils.env_detector as mod
        from sagellm.cli.commands.env import env

        # Patch shutil.which inside the env command module
        monkeypatch.setattr(
            env_cmd.shutil, "which", lambda cmd: "/usr/bin/npu-smi" if cmd == "npu-smi" else None
        )

        # Patch check_ascend_env_parity in the env command module (already imported there)
        tk = tmp_path / "toolkit"

        def _all_ok():
            return [
                mod.AscendVarCheck(
                    name=v,
                    current=str(tk),
                    expected=str(tk),
                    status="ok",
                    repair=None,
                )
                for v in ASCEND_CRITICAL_VARS
            ]

        monkeypatch.setattr(env_cmd, "check_ascend_env_parity", _all_ok)

        runner = CliRunner()
        result = runner.invoke(env, ["check-ascend"])
        assert result.exit_code == 0
        assert "healthy" in result.output
