from __future__ import annotations

from sagellm.cuda_repair import install_flash_attn


def test_install_flash_attn_uses_binary_only_by_default(monkeypatch, capsys) -> None:
    captured: dict[str, object] = {}

    class _Result:
        ok = True
        output = "ok"
        returncode = 0

    class _Verify:
        ok = True
        output = "2.8.3"
        returncode = 0

    def _fake_install_probe(
        python_executable: str, packages: list[str], *, extra_args: list[str] | None = None
    ):
        captured["python_executable"] = python_executable
        captured["packages"] = packages
        captured["extra_args"] = extra_args
        return _Result()

    monkeypatch.setattr("sagellm.cuda_repair.pip_install_probe", _fake_install_probe)
    monkeypatch.setattr(
        "sagellm.cuda_repair.run_python_probe", lambda python_executable, script: _Verify()
    )

    exit_code = install_flash_attn("python", allow_source_build=False)

    assert exit_code == 0
    assert captured["packages"] == ["flash-attn>=2.5.0"]
    assert captured["extra_args"] == ["--only-binary=:all:"]
    assert "FlashAttention 安装成功" in capsys.readouterr().out


def test_install_flash_attn_fails_fast_without_source_build(monkeypatch, capsys) -> None:
    class _Result:
        ok = False
        output = "ERROR: No matching distribution found for flash-attn"
        returncode = 1

    monkeypatch.setattr("sagellm.cuda_repair.pip_install_probe", lambda *args, **kwargs: _Result())

    exit_code = install_flash_attn("python", allow_source_build=False)

    assert exit_code == 1
    output = capsys.readouterr().out
    assert "禁止默认源码编译" in output
    assert "SAGELLM_ALLOW_FLASH_ATTN_SOURCE_BUILD=1" in output


def test_install_flash_attn_allows_explicit_source_build(monkeypatch) -> None:
    captured: dict[str, object] = {}

    class _Result:
        ok = False
        output = "build failed"
        returncode = 2

    def _fake_install_probe(
        python_executable: str, packages: list[str], *, extra_args: list[str] | None = None
    ):
        captured["extra_args"] = extra_args
        return _Result()

    monkeypatch.setattr("sagellm.cuda_repair.pip_install_probe", _fake_install_probe)

    exit_code = install_flash_attn("python", allow_source_build=True)

    assert exit_code == 2
    assert captured["extra_args"] == ["--no-build-isolation"]
