import os
import unittest
from io import StringIO
from unittest.mock import patch

import biotite.structure as struc
import biotite.structure.io.pdb as pdb

from synth_nmr import synth_nmr_cli as cli
from synth_nmr.synth_nmr_cli import interactive_mode, process_commands


class TestCLICoverage(unittest.TestCase):
    def setUp(self):
        # Reset global state to avoid cross-test interference
        cli.structure = None
        cli.ensemble = None

        # Create a minimal PDB file for testing (3 residues for backbone dihedrals)
        self.pdb_path = "test_cli.pdb"
        n_res = 3
        structure = struc.AtomArray(n_res * 10)
        atom_names = ["N", "CA", "C", "O", "CB", "H", "HA", "HB1", "HB2", "CG"]
        elements = ["N", "C", "C", "O", "C", "H", "H", "H", "H", "C"]

        for r in range(n_res):
            start = r * 10
            for i in range(10):
                idx = start + i
                structure.res_id[idx] = r + 1
                structure.res_name[idx] = "ALA"
                structure.chain_id[idx] = "A"
                structure.atom_name[idx] = atom_names[i]
                structure.element[idx] = elements[i]
                # Provide some spacing so dihedrals aren't completely zero/invalid
                structure.coord[idx] = [r * 3.8, i * 1.5, (i % 2) * 0.5]

        pdb_file = pdb.PDBFile()
        pdb_file.set_structure(structure)
        pdb_file.write(self.pdb_path)

    def tearDown(self):
        if os.path.exists(self.pdb_path):
            os.remove(self.pdb_path)

    def test_cli_simple_commands(self):
        """Test basic non-interactive commands."""
        args = [
            "read",
            "pdb",
            self.pdb_path,
            "calculate",
            "j-coupling",
            "predict",
            "shifts",
            "calculate",
            "rdc",
            "10.0",
            "0.5",
        ]
        with patch("sys.stdout", new=StringIO()) as fake_out:
            process_commands(args)
            output = fake_out.getvalue()
            assert "Read PDB file" in output
            assert "3J_HNHa" in output
            assert "Chain: A" in output
            assert "RDC:" in output

    def test_cli_ensemble_commands(self):
        """Test ensemble-related commands."""
        # load trajectory requires multiple frames
        args = [
            "load",
            "trajectory",
            self.pdb_path,
            self.pdb_path,
            "ensemble",
            "shifts",
            "ensemble",
            "noes",
            "5.0",
            "ensemble",
            "rdcs",
            "10.0",
            "0.5",
            "ensemble",
            "j-coupling",
            "ensemble",
            "s2",
        ]
        with patch("sys.stdout", new=StringIO()) as fake_out:
            process_commands(args)
            output = fake_out.getvalue()
            assert "Loaded trajectory ensemble with 2 frames." in output
            assert "ppm" in output
            assert "r_eff" in output
            assert "D_NH" in output
            assert "3J_HNHa" in output
            assert "S² =" in output

    def test_interactive_mode_coverage(self):
        """Test interactive mode with various commands."""
        commands = [
            f"read pdb {self.pdb_path}",
            f"load trajectory {self.pdb_path} {self.pdb_path}",
            "ensemble shifts",
            "ensemble noes 6.0",
            "ensemble rdcs 12.0 0.6",
            "ensemble s2",
            "calculate rdc 11.0 0.4",
            "predict shifts",
            "calculate j-coupling",
            "help",
            "unknown_command",
            "exit",
        ]
        input_str = "\n".join(commands) + "\n"

        with patch("sys.stdin", StringIO(input_str)):
            # Create a dummy PDB file for loading
            with open("test_traj.pdb", "w") as f:
                f.write(
                    "ATOM      1  N   ALA A   1       0.000   0.000   0.000  1.00  0.00           N\n"
                )
                f.write(
                    "ATOM      2  CA  ALA A   1       1.458   0.000   0.000  1.00  0.00           C\n"
                )
                f.write(
                    "ATOM      3  C   ALA A   1       2.009   1.359   0.000  1.00  0.00           C\n"
                )
                f.write(
                    "ATOM      4  O   ALA A   1       1.327   2.378   0.000  1.00  0.00           O\n"
                )
                f.write(
                    "ATOM      5  CB  ALA A   1       2.009  -0.679   1.211  1.00  0.00           C\n"
                )
                f.write(
                    "ATOM      6  H   ALA A   1      -0.500  -0.866   0.000  1.00  0.00           H\n"
                )
                f.write(
                    "ATOM      7  HA  ALA A   1       1.800  -0.500  -0.800  1.00  0.00           H\n"
                )
                f.write("TER\n")
                f.write("END\n")

            cli.handle_interactive_command("load trajectory test_traj.pdb")
            assert cli.ensemble is not None

            # Now run ensemble subcommands via handle_interactive_command
            with patch(
                "synth_nmr.nmr.calculate_synthetic_noes",
                return_value=[{"index_1": 1, "index_2": 2, "distance": 3.0}],
            ):
                cli.handle_interactive_command("ensemble noes")

            with patch(
                "synth_nmr.j_coupling.calculate_hn_ha_coupling", return_value={"A": {1: 7.5}}
            ):
                cli.handle_interactive_command("ensemble j-coupling")

            with patch("synth_nmr.rdc.calculate_rdcs", return_value={1: 10.0}):
                cli.handle_interactive_command("ensemble rdcs")

            with patch(
                "synth_nmr.chemical_shifts.predict_chemical_shifts",
                return_value={"method": {1: {"N": 120.0}}},
            ):
                cli.handle_interactive_command("ensemble shifts")

            # Test error cases
            cli.handle_interactive_command("ensemble unknown")

            with patch("sys.stdout", new=StringIO()) as fake_out:
                interactive_mode()
                output = fake_out.getvalue()
                assert "Welcome to the synth-nmr CLI!" in output
                assert "Read PDB file" in output
                assert "Loaded trajectory ensemble" in output
                assert "Commands:" in output
                assert "Unknown command" in output

    def test_cli_error_handling(self):
        """Test CLI error handling paths."""
        bad_args = [
            "read",
            "pdb",
            "non_existent.pdb",
            "ensemble",
            "shifts",  # Should fail because no trajectory loaded
            "calculate",
            "rdc",
            "invalid_val",
            "load",
            "trajectory",
            "non_existent.pdb",
        ]
        with patch("sys.stdout", new=StringIO()) as fake_out:
            process_commands(bad_args)
            output = fake_out.getvalue()
            assert "Error: File not found" in output
            assert "Error: No trajectory loaded" in output
            assert "Warning: Could not read" in output


if __name__ == "__main__":
    unittest.main()
