from __future__ import annotations

from pathlib import Path
import json
import os

import pytest

from ins_pricing.cli.utils import notebook_utils


def test_validate_entry_runner_config_raises_when_nproc_exceeds_visible_cuda(
    monkeypatch,
):
    monkeypatch.setattr(
        notebook_utils,
        "_visible_cuda_device_count",
        lambda: 1,
    )
    raw_cfg = {"use_gpu": True}
    runner_cfg = {"nproc_per_node": 2}

    with pytest.raises(ValueError, match="runner.nproc_per_node=2"):
        notebook_utils._validate_entry_runner_config(
            raw_cfg,
            runner_cfg=runner_cfg,
            config_path=Path("examples/config_resn_direct.json"),
        )


def test_validate_entry_runner_config_warns_memory_saving_ddp_workers(
    monkeypatch,
    capsys,
):
    monkeypatch.setattr(
        notebook_utils,
        "_visible_cuda_device_count",
        lambda: 2,
    )
    raw_cfg = {
        "use_gpu": True,
        "dataloader_workers": 2,
        "env": {
            "BAYESOPT_RESOURCE_PROFILE": "memory_saving",
            "BAYESOPT_DDP_ALLOW_WORKERS_IN_MEMORY_SAVING": "1",
        },
    }
    runner_cfg = {"nproc_per_node": 2}

    notebook_utils._validate_entry_runner_config(
        raw_cfg,
        runner_cfg=runner_cfg,
        config_path=Path("examples/config_resn_direct.json"),
    )

    captured = capsys.readouterr()
    assert "DDP + memory_saving" in captured.out


def test_validate_entry_runner_config_skips_single_process(monkeypatch):
    monkeypatch.setattr(
        notebook_utils,
        "_visible_cuda_device_count",
        lambda: 0,
    )
    raw_cfg = {"use_gpu": True}
    runner_cfg = {"nproc_per_node": 1}

    notebook_utils._validate_entry_runner_config(
        raw_cfg,
        runner_cfg=runner_cfg,
        config_path=Path("examples/config_resn_direct.json"),
    )


def test_build_run_spec_from_config_returns_env_without_mutating_process(
    tmp_path,
    monkeypatch,
):
    config_path = tmp_path / "config.json"
    config_path.write_text(
        json.dumps(
            {
                "data_dir": "./Data",
                "model_list": ["od"],
                "model_categories": ["bc"],
                "target": "y",
                "weight": "w",
                "feature_list": ["x1"],
                "categorical_features": [],
                "env": {"OMP_NUM_THREADS": "7"},
                "runner": {"mode": "entry", "model_keys": ["ft"], "nproc_per_node": 1},
            }
        ),
        encoding="utf-8",
    )
    monkeypatch.setenv("OMP_NUM_THREADS", "3")

    cmd, mode, env = notebook_utils.build_run_spec_from_config(config_path)

    assert mode == "entry"
    assert "--config-json" in cmd
    assert env is not None
    assert env["OMP_NUM_THREADS"] == "7"
    assert os.environ["OMP_NUM_THREADS"] == "3"
