import pytest
import os
import tempfile
from synth_nmr.nef_io import read_nef_restraints, write_nef_file


def test_read_nef_restraints_valid():
    """Test reading restraints from a valid NEF file."""
    restraints = [
        {
            "chain_1": "A",
            "index_1": 1,
            "res_name_1": "ALA",
            "atom_name_1": "H",
            "chain_2": "A",
            "index_2": 5,
            "res_name_2": "GLY",
            "atom_name_2": "HA",
            "distance": 4.5,
            "upper_limit": 5.0,
            "lower_limit": 1.8,
        }
    ]
    sequence = "AAAAAA"

    with tempfile.NamedTemporaryFile(suffix=".nef", mode="w", delete=False) as tmp:
        tmp_path = tmp.name

    try:
        write_nef_file(tmp_path, sequence, restraints)

        # Now read it back
        parsed = read_nef_restraints(tmp_path)

        assert len(parsed) == 1
        assert parsed[0]["chain_1"] == "A"
        assert parsed[0]["seq_1"] == 1
        assert parsed[0]["atom_1"] == "H"
        assert parsed[0]["dist"] == pytest.approx(4.5)
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


def test_read_nef_restraints_file_not_found(caplog):
    """Test handling of missing NEF file."""
    result = read_nef_restraints("non_existent_file.nef")
    assert result == []
    assert "NEF file not found" in caplog.text


def test_read_nef_restraints_malformed():
    """Test handling of malformed NEF file (missing column)."""
    malformed_nef = """
data_test
save_synthetic_noes
   _nef_distance_restraint_list.sf_category nef_distance_restraint_list
   _nef_distance_restraint_list.sf_framecode synthetic_noes
   _nef_distance_restraint_list.restraint_origin synthetic

   loop_
      _nef_distance_restraint.index
      _nef_distance_restraint.chain_code_1
      _nef_distance_restraint.sequence_code_1
      # Missing other headers
      1 A 1
   stop_
save_
"""
    with tempfile.NamedTemporaryFile(suffix=".nef", mode="w", delete=False) as tmp:
        tmp.write(malformed_nef)
        tmp_path = tmp.name

    try:
        parsed = read_nef_restraints(tmp_path)
        # Should return empty list because required keys like 'chain_code_2' are missing in headers
        assert parsed == []
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


def test_read_nef_restraints_invalid_data():
    """Test handling of invalid data in NEF file (e.g. string where number expected)."""
    invalid_data_nef = """
data_test
save_synthetic_noes
   _nef_distance_restraint_list.sf_category nef_distance_restraint_list
   _nef_distance_restraint_list.sf_framecode synthetic_noes
   _nef_distance_restraint_list.restraint_origin synthetic

   loop_
      _nef_distance_restraint.index
      _nef_distance_restraint.chain_code_1
      _nef_distance_restraint.sequence_code_1
      _nef_distance_restraint.residue_name_1
      _nef_distance_restraint.atom_name_1
      _nef_distance_restraint.chain_code_2
      _nef_distance_restraint.sequence_code_2
      _nef_distance_restraint.residue_name_2
      _nef_distance_restraint.atom_name_2
      _nef_distance_restraint.target_value
      1 A INVALID_SEQ ALA H A 5 GLY HA 4.5
   stop_
save_
"""
    with tempfile.NamedTemporaryFile(suffix=".nef", mode="w", delete=False) as tmp:
        tmp.write(invalid_data_nef)
        tmp_path = tmp.name

    try:
        parsed = read_nef_restraints(tmp_path)
        # Should skip the line with INVALID_SEQ
        assert parsed == []
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
