import pytest
from folding_studio.query import BoltzQuery


@pytest.fixture
def inference_parameters():
    """
    Provides the inference parameters dictionary to tests.
    """
    return {
        "recycling_steps": 3,
        "sampling_steps": 200,
        "diffusion_samples": 1,
        "step_scale": 1.638,
        "output_format": "mmcif",
        "num_workers": 2,
        "msa_pairing_strategy": "greedy",
        "write_full_pae": False,
        "write_full_pde": False,
        "seed": 0,
    }


def test_process_fasta_file(tmp_files):
    """Test processing a valid FASTA file."""
    file_path = tmp_files["monomer_fasta"]
    fasta_dict, yaml_dict = BoltzQuery._process_file(file_path)
    expected_fasta = ">tag1|tag2\nABCDEGF"
    expected_fasta_dict = {file_path.stem: expected_fasta}

    assert fasta_dict == expected_fasta_dict
    assert yaml_dict == {}


def test_process_yaml_file(tmp_files):
    """Test processing a valid YAML file."""
    file_path = tmp_files["yaml_file_path"]
    fasta_dict, yaml_dict = BoltzQuery._process_file(file_path)

    expected_yaml = {
        "version": 1,
        "sequences": [{"protein": {"id": "A", "sequence": "QLEDSEVEAVAKGLEE"}}],
    }
    expected_yaml_dict = {file_path.stem: expected_yaml}

    assert fasta_dict == {}
    assert yaml_dict == expected_yaml_dict


def test_process_file_invalid_extension(tmp_files):
    """Test processing a file with an invalid extension."""
    invalid_source = tmp_files["invalid_source"]
    with pytest.raises(
        ValueError, match=f"Unsupported format: {invalid_source.suffix}"
    ):
        BoltzQuery._process_file(invalid_source)


def test_from_file_fasta(tmp_files, inference_parameters):
    """Test creating BoltzQuery from a FASTA file."""
    file_path = tmp_files["monomer_fasta"]
    query = BoltzQuery.from_file(file_path, **inference_parameters)
    assert file_path.stem in query.fasta_dict
    assert not query.yaml_dict


def test_from_file_yaml(tmp_files, inference_parameters):
    """Test creating BoltzQuery from a YAML file."""
    file_path = tmp_files["yaml_file_path"]
    query = BoltzQuery.from_file(file_path, **inference_parameters)
    assert not query.fasta_dict
    assert file_path.stem in query.yaml_dict


def test_from_directory(tmp_files, inference_parameters):
    """Test creating BoltzQuery from a directory containing both FASTA and YAML files."""
    query = BoltzQuery.from_directory(
        tmp_files["mixed_fasta_yaml_dir"], **inference_parameters
    )
    assert "protein_A" in query.fasta_dict
    assert "protein_B" in query.yaml_dict


def test_from_source_file(tmp_files, inference_parameters):
    """Test creating BoltzQuery from a single file."""
    file_path = tmp_files["monomer_fasta"]
    query = BoltzQuery.from_file(str(file_path), **inference_parameters)
    assert "monomer" in query.fasta_dict


def test_from_source_directory(tmp_files, inference_parameters):
    """Test creating BoltzQuery from a directory."""