"""Tests for configuration loading."""

from __future__ import annotations


import yaml
import pytest

from obliteratus.config import StudyConfig


SAMPLE_CONFIG = {
    "model": {
        "name": "gpt2",
        "task": "causal_lm",
        "dtype": "float32",
        "device": "cpu",
    },
    "dataset": {
        "name": "wikitext",
        "subset": "wikitext-2-raw-v1",
        "split": "test",
        "text_column": "text",
        "max_samples": 50,
    },
    "strategies": [
        {"name": "layer_removal", "params": {}},
        {"name": "ffn_ablation", "params": {}},
    ],
    "metrics": ["perplexity"],
    "batch_size": 4,
    "max_length": 256,
    "output_dir": "results/test",
}


class TestStudyConfig:
    def test_from_dict(self):
        config = StudyConfig.from_dict(SAMPLE_CONFIG)
        assert config.model.name == "gpt2"
        assert config.model.task == "causal_lm"
        assert config.dataset.name == "wikitext"
        assert len(config.strategies) == 2
        assert config.strategies[0].name == "layer_removal"

    def test_from_yaml(self, tmp_path):
        yaml_path = tmp_path / "test_config.yaml"
        yaml_path.write_text(yaml.dump(SAMPLE_CONFIG))

        config = StudyConfig.from_yaml(yaml_path)
        assert config.model.name == "gpt2"
        assert config.batch_size == 4

    def test_roundtrip(self):
        config = StudyConfig.from_dict(SAMPLE_CONFIG)
        d = config.to_dict()
        config2 = StudyConfig.from_dict(d)
        assert config2.model.name == config.model.name
        assert config2.dataset.name == config.dataset.name
        assert len(config2.strategies) == len(config.strategies)

    def test_model_quantization_and_label_count_roundtrip(self):
        raw = {
            **SAMPLE_CONFIG,
            "model": {
                **SAMPLE_CONFIG["model"],
                "num_labels": 7,
                "quantization": "4bit",
            },
        }
        config = StudyConfig.from_dict(raw)
        assert config.model.quantization == "4bit"
        assert config.model.num_labels == 7
        assert StudyConfig.from_dict(config.to_dict()).model == config.model

    def test_remote_config_roundtrip_preserves_security_and_execution_settings(self):
        raw = {
            **SAMPLE_CONFIG,
            "remote": {
                "host": "compute.example",
                "user": "runner",
                "port": 2222,
                "ssh_key": "/keys/id",
                "known_hosts_file": "/keys/known_hosts",
                "remote_dir": "/srv/obliteratus",
                "install_timeout": 120,
                "python": "/opt/python",
                "sync_results": False,
                "gpus": "02, 0",
                "install_source": "obliteratus==0.1.2",
            },
        }
        config = StudyConfig.from_dict(raw)
        assert config.remote is not None
        assert config.remote.gpus == "2,0"
        assert StudyConfig.from_dict(config.to_dict()) == config

    @pytest.mark.parametrize(
        "remote",
        [
            {"host": ""},
            {"host": "host", "port": 0},
            {"host": "host", "remote_dir": "relative"},
            {"host": "host", "gpus": "0; injected"},
            {"host": "host", "install_timeout": 0},
        ],
    )
    def test_remote_config_rejects_invalid_public_values(self, remote):
        with pytest.raises(ValueError):
            StudyConfig.from_dict({**SAMPLE_CONFIG, "remote": remote})
