from unittest import mock

import pytest
from folding_studio.cli import app
from folding_studio.query import BoltzQuery
from typer.testing import CliRunner

runner = CliRunner()


@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,
    }


@pytest.fixture()
def mock_save_inference_parameters():
    """Monkeypatch Response.download_results so that no real file download occurs."""
    with mock.patch("folding_studio.query.Query.save_parameters") as m:
        yield m


def test_boltz_fold_with_valid_fasta_file(
    mock_send_request,
    mock_download_results,
    mock_save_inference_parameters,
    tmp_files,
    inference_parameters,
):
    """
    Running the command with a valid FASTA file should
    process successfully and generate the correct payload.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "boltz",
            str(tmp_files["monomer_fasta"]),
            "--project-code",
            "test_project",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    # Create the expected query using the same parameters as defaults.
    expected_query = BoltzQuery.from_file(
        path=str(tmp_files["monomer_fasta"]), **inference_parameters
    )
    mock_send_request.assert_called_once()
    actual_query = mock_send_request.call_args[0][0]
    # Compare payload keys
    payload = actual_query.payload
    expected = expected_query.payload
    for key in ["fasta_files", "yaml_files", "parameters"]:
        assert payload[key] == expected[key]

    mock_save_inference_parameters.assert_called_once()
    mock_download_results.assert_called_once()


def test_boltz_fold_with_valid_fasta_directory(
    mock_send_request,
    mock_download_results,
    mock_save_inference_parameters,
    tmp_files,
    inference_parameters,
):
    """
    Running the command with a directory containing valid FASTA files should
    process successfully and combine all FASTA files into the payload.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "boltz",
            str(tmp_files["valid_dir"]),
            "--project-code",
            "test_project",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code == 0, result.output

    expected_query = BoltzQuery.from_directory(
        path=tmp_files["valid_dir"], **inference_parameters
    )
    mock_send_request.assert_called_once()
    actual_query = mock_send_request.call_args[0][0]
    payload = actual_query.payload
    expected = expected_query.payload
    assert payload["fasta_files"] == expected["fasta_files"]

    mock_save_inference_parameters.assert_called_once()
    mock_download_results.assert_called_once()


def test_boltz_fold_with_empty_directory(tmp_files):
    """
    Running the command with an empty directory should return an error,
    since no FASTA files are available.
    """
    result = runner.invoke(
        app,
        [
            "predict",
            "boltz",
            str(tmp_files["empty_dir"]),
            "--project-code",
            "test_project",
            "--output",
            tmp_files["output_dir"],
        ],
    )
    assert result.exit_code != 0, result.output