"""Tests for Ascend environment normalization in profile refresh."""

from __future__ import annotations

from pathlib import Path

from sagellm.cli.commands import profile as profile_cmd


def test_normalize_ascend_env_prefers_toolkit_home(monkeypatch, tmp_path: Path) -> None:
    toolkit_home = tmp_path / "ascend-toolkit" / "8.0.RC3"
    (toolkit_home / "runtime" / "lib64").mkdir(parents=True)
    (toolkit_home / "fwkacllib" / "lib64").mkdir(parents=True)
    (toolkit_home / "hccl" / "lib64").mkdir(parents=True)
    (toolkit_home / "opp").mkdir(parents=True)
    (toolkit_home / "aicpu_kernels").mkdir(parents=True)

    monkeypatch.setattr(profile_cmd, "_iter_ascend_toolkit_homes", lambda: [toolkit_home])
    monkeypatch.setenv("ASCEND_HOME", "/usr/local/Ascend")
    monkeypatch.setenv("ASCEND_HOME_PATH", "/usr/local/Ascend")
    monkeypatch.setenv("ASCEND_RUNTIME_PATH", "/usr/local/Ascend/runtime")
    monkeypatch.setenv("LD_LIBRARY_PATH", "/usr/lib64")

    profile_cmd._normalize_ascend_env_for_refresh()

    assert profile_cmd.os.environ["ASCEND_HOME"] == str(toolkit_home)
    assert profile_cmd.os.environ["ASCEND_HOME_PATH"] == str(toolkit_home)
    assert profile_cmd.os.environ["ASCEND_RUNTIME_PATH"] == str(toolkit_home / "runtime")
    assert profile_cmd.os.environ["ASCEND_OPP_PATH"] == str(toolkit_home / "opp")
    assert profile_cmd.os.environ["ASCEND_AICPU_PATH"] == str(toolkit_home / "aicpu_kernels")
    assert profile_cmd.os.environ["LD_LIBRARY_PATH"].startswith(
        f"{toolkit_home / 'fwkacllib' / 'lib64'}:"
    )


def test_normalize_ascend_env_unsets_invalid_runtime(monkeypatch, tmp_path: Path) -> None:
    toolkit_home = tmp_path / "ascend-toolkit" / "8.5.0"
    (toolkit_home / "fwkacllib" / "lib64").mkdir(parents=True)
    (toolkit_home / "hccl" / "lib64").mkdir(parents=True)
    (toolkit_home / "opp").mkdir(parents=True)

    monkeypatch.setattr(profile_cmd, "_iter_ascend_toolkit_homes", lambda: [toolkit_home])
    monkeypatch.setenv("ASCEND_RUNTIME_PATH", "/usr/local/Ascend/runtime")

    profile_cmd._normalize_ascend_env_for_refresh()

    assert "ASCEND_RUNTIME_PATH" not in profile_cmd.os.environ
    assert profile_cmd.os.environ["ASCEND_HOME"] == str(toolkit_home)
    assert profile_cmd.os.environ["ASCEND_HOME_PATH"] == str(toolkit_home)
