import os
import sys
import subprocess
import tempfile
import pytest
import numpy as np
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
import taco_format


@pytest.fixture
def temp_dir():
    """Create a temporary directory for test files"""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield tmpdir


def create_test_trajectory(num_frames=10, num_atoms=3):
    """Create a test trajectory with multiple frames"""
    trajectory = []

    for i in range(num_frames):
        # Create water molecule with time-dependent positions
        positions = [
            [0.0 + 0.1 * i, 0.0, 0.0],  # O
            [0.95 + 0.05 * np.sin(i * 0.1), 0.0, 0.0],  # H1
            [0.45, 0.85 + 0.05 * np.cos(i * 0.1), 0.0],  # H2
        ]

        atoms = Atoms("H2O", positions=positions)
        atoms.set_cell([10.0, 10.0, 10.0])
        atoms.set_pbc([True, True, True])

        # Set velocities
        velocities = [
            [0.1 * np.sin(i * 0.2), 0.0, 0.1],
            [0.0, 0.1 * np.cos(i * 0.2), 0.0],
            [0.1, 0.0, 0.1 * np.sin(i * 0.3)],
        ]
        atoms.set_velocities(velocities)

        # Set forces and energy
        forces = [
            [0.1 + 0.01 * i, 0.2, 0.3],
            [0.4, 0.5 + 0.01 * i, 0.6],
            [0.7, 0.8, 0.9 + 0.01 * i],
        ]
        energy = -10.5 - 0.1 * i

        calc = SinglePointCalculator(atoms, energy=energy, forces=forces)
        atoms.calc = calc

        trajectory.append(atoms)

    return trajectory


def find_taco_binary():
    """Find the taco binary in the build output or PATH"""
    possible_paths = ["target/release/taco", "target/debug/taco", "taco"]  # In PATH

    for path in possible_paths:
        try:
            result = subprocess.run(
                [path, "--version"], capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                return path
        except (subprocess.TimeoutExpired, FileNotFoundError):
            continue

    pytest.skip(
        "taco binary not found. Please build it first with 'cargo build --release'"
    )


def run_taco_command(args, cwd=None, check=True):
    """Run a taco command and return the result"""
    taco_binary = find_taco_binary()
    cmd = [taco_binary] + args

    result = subprocess.run(cmd, capture_output=True, text=True, cwd=cwd, timeout=30)

    if check and result.returncode != 0:
        pytest.fail(
            f"Command failed: {' '.join(cmd)}\n"
            f"stdout: {result.stdout}\n"
            f"stderr: {result.stderr}\n"
            f"return code: {result.returncode}"
        )

    return result


class TestTacoBinary:
    """Test cases for the taco command line utility"""

    def test_binary_exists_and_shows_help(self):
        """Test that the binary exists and shows help"""
        result = run_taco_command(["--help"])
        assert "TACO trajectory file utility" in result.stdout
        assert "info" in result.stdout
        assert "extract" in result.stdout
        assert "convert" in result.stdout
        assert "check" in result.stdout
        assert "stats" in result.stdout

    def test_info_command(self, temp_dir):
        """Test the info command"""
        # Create a test trajectory
        trajectory = create_test_trajectory(5, 3)
        taco_file = os.path.join(temp_dir, "test_trajectory.taco")

        # Write trajectory using the Python API
        taco_format.write(taco_file, trajectory)

        # Test info command
        result = run_taco_command(["info", taco_file])

        # Check that important information is displayed
        assert "TACO File Information" in result.stdout
        assert "Number of atoms: 3" in result.stdout
        assert "Number of frames: 5" in result.stdout
        assert "Format version:" in result.stdout
        assert "Simulation Metadata" in result.stdout
        assert "Compression Settings" in result.stdout

    def test_extract_frame_range(self, temp_dir):
        """Test extracting a range of frames"""
        # Create a test trajectory
        trajectory = create_test_trajectory(10, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "extracted.taco")

        # Write trajectory
        taco_format.write(input_file, trajectory)

        # Extract frames 2-5
        result = run_taco_command(
            ["extract", input_file, output_file, "--start", "2", "--end", "6"]
        )

        assert "Successfully extracted 4 frames" in result.stdout

        # Verify the extracted file
        extracted_trajectory = taco_format.read(output_file)
        assert len(extracted_trajectory) == 4

        # Check that the content matches the original frames 2-5
        for i, atoms in enumerate(extracted_trajectory):
            original_atoms = trajectory[i + 2]
            np.testing.assert_allclose(
                atoms.get_positions(), original_atoms.get_positions(), atol=1e-3
            )

    def test_extract_specific_frames(self, temp_dir):
        """Test extracting specific frame indices"""
        # Create a test trajectory
        trajectory = create_test_trajectory(10, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "specific.taco")

        # Write trajectory with full precision to avoid compression artifacts
        taco_format.write(input_file, trajectory, lossless=True)

        # Extract specific frames: 0, 3, 7, 9
        result = run_taco_command(
            ["extract", input_file, output_file, "--frames", "0,3,7,9"]
        )

        assert "Successfully extracted 4 frames" in result.stdout

        # Verify the extracted file
        extracted_trajectory = taco_format.read(output_file)
        assert len(extracted_trajectory) == 4
        # Check specific frames
        frame_indices = [0, 3, 7, 9]
        for i, atoms in enumerate(extracted_trajectory):
            original_atoms = trajectory[frame_indices[i]]
            np.testing.assert_allclose(
                atoms.get_positions(),
                original_atoms.get_positions(),
                atol=1.0,  # More tolerant for compressed data with half precision
            )

    def test_extract_with_step(self, temp_dir):
        """Test extracting frames with a step size"""
        # Create a test trajectory
        trajectory = create_test_trajectory(10, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "stepped.taco")

        # Write trajectory with full precision to avoid compression artifacts
        taco_format.write(input_file, trajectory, lossless=True)

        # Extract every 3rd frame
        result = run_taco_command(["extract", input_file, output_file, "--step", "3"])

        # Should extract frames 0, 3, 6, 9 (4 frames total)
        assert "Successfully extracted 4 frames" in result.stdout

        # Verify the extracted file
        extracted_trajectory = taco_format.read(output_file)
        assert len(extracted_trajectory) == 4
        # Check that we got the right frames
        expected_indices = [0, 3, 6, 9]
        for i, atoms in enumerate(extracted_trajectory):
            original_atoms = trajectory[expected_indices[i]]
            np.testing.assert_allclose(
                atoms.get_positions(),
                original_atoms.get_positions(),
                atol=1.0,  # More tolerant for compressed data with half precision
            )

    def test_convert_command(self, temp_dir):
        """Test the convert command with different compression settings"""
        # Create a test trajectory
        trajectory = create_test_trajectory(5, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "converted.taco")

        # Write trajectory with default settings
        taco_format.write(input_file, trajectory)

        # Convert with higher compression
        result = run_taco_command(
            [
                "convert",
                input_file,
                output_file,
                "--compression",
                "10",
                "--full-frame-interval",
                "5",
            ]
        )

        assert "Successfully converted" in result.stdout

        # Verify the converted file can be read
        converted_trajectory = taco_format.read(output_file)
        assert len(converted_trajectory) == 5

        # Check that positions are preserved (within tolerance for lossy compression)
        for i, atoms in enumerate(converted_trajectory):
            original_atoms = trajectory[i]
            np.testing.assert_allclose(
                atoms.get_positions(),
                original_atoms.get_positions(),
                atol=1e-2,  # Larger tolerance for compressed data
            )

    def test_check_command(self, temp_dir):
        """Test the check command for file validation"""
        # Create a test trajectory
        trajectory = create_test_trajectory(5, 3)
        taco_file = os.path.join(temp_dir, "test.taco")

        # Write trajectory
        taco_format.write(taco_file, trajectory)

        # Check the file
        result = run_taco_command(["check", taco_file])

        assert "File validation successful" in result.stdout
        assert "no errors found" in result.stdout

    def test_check_command_verbose(self, temp_dir):
        """Test the check command with verbose output"""
        # Create a test trajectory
        trajectory = create_test_trajectory(5, 3)
        taco_file = os.path.join(temp_dir, "test.taco")

        # Write trajectory
        taco_format.write(taco_file, trajectory)

        # Check the file with verbose output
        result = run_taco_command(["check", taco_file, "--verbose"])

        assert "Validating 5 frames" in result.stdout
        assert "File validation successful" in result.stdout

    def test_stats_command(self, temp_dir):
        """Test the stats command"""
        # Create a test trajectory
        trajectory = create_test_trajectory(8, 3)
        taco_file = os.path.join(temp_dir, "test.taco")

        # Write trajectory
        taco_format.write(taco_file, trajectory)

        # Get stats
        result = run_taco_command(["stats", taco_file])

        assert "Trajectory Statistics" in result.stdout
        assert "Position ranges:" in result.stdout
        assert "Data availability:" in result.stdout
        assert "Frames with velocities:" in result.stdout
        assert "Frames with forces:" in result.stdout

    def test_stats_command_detailed(self, temp_dir):
        """Test the stats command with detailed output"""
        # Create a smaller trajectory for detailed output
        trajectory = create_test_trajectory(5, 3)
        taco_file = os.path.join(temp_dir, "test.taco")

        # Write trajectory
        taco_format.write(taco_file, trajectory)

        # Get detailed stats
        result = run_taco_command(["stats", taco_file, "--detailed"])

        assert "Detailed Frame Information" in result.stdout
        assert "Frame 0:" in result.stdout
        assert "positions=true" in result.stdout
        assert "velocities=true" in result.stdout
        assert "forces=true" in result.stdout

    def test_error_handling_nonexistent_file(self):
        """Test error handling for non-existent files"""
        result = run_taco_command(["info", "/nonexistent/file.taco"], check=False)
        assert result.returncode != 0
        assert (
            "No such file or directory" in result.stderr
            or "cannot find the file" in result.stderr.lower()
        )

    def test_error_handling_invalid_frame_range(self, temp_dir):
        """Test error handling for invalid frame ranges"""
        # Create a small trajectory
        trajectory = create_test_trajectory(5, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "output.taco")

        taco_format.write(input_file, trajectory)

        # Try to extract beyond the available frames
        result = run_taco_command(
            ["extract", input_file, output_file, "--start", "10", "--end", "15"],
            check=False,
        )

        assert result.returncode != 0
        assert "exceeds total frames" in result.stderr

    def test_error_handling_invalid_frame_indices(self, temp_dir):
        """Test error handling for invalid frame indices"""
        # Create a small trajectory
        trajectory = create_test_trajectory(5, 3)
        input_file = os.path.join(temp_dir, "input.taco")
        output_file = os.path.join(temp_dir, "output.taco")

        taco_format.write(input_file, trajectory)

        # Try to extract frames that don't exist
        result = run_taco_command(
            ["extract", input_file, output_file, "--frames", "0,3,10,15"], check=False
        )

        assert result.returncode != 0
        assert "exceeds total frames" in result.stderr

    def test_version_command(self):
        """Test the version command"""
        result = run_taco_command(["--version"])
        assert result.returncode == 0
        assert "taco" in result.stdout
        # Check that it follows semantic versioning pattern
        import re

        assert re.search(r"\d+\.\d+\.\d+", result.stdout)

    def test_python_cli_consistency(self, temp_dir):
        """Test that Python CLI and Rust binary produce consistent results"""
        # Create a test trajectory
        trajectory = create_test_trajectory(5, 3)
        taco_file = os.path.join(temp_dir, "test.taco")
        taco_format.write(taco_file, trajectory)

        # Test info command through binary
        binary_result = run_taco_command(["info", taco_file])

        # Test info command through Python CLI
        python_result = None
        try:
            import subprocess
            import sys

            python_result = subprocess.run(
                [
                    sys.executable,
                    "-c",
                    f"import taco_format; taco_format.run_cli(['info', '{taco_file}'])",
                ],
                capture_output=True,
                text=True,
                timeout=30,
            )
        except Exception:
            pytest.skip("Could not test Python CLI entry point")

        # Both should succeed and contain similar information
        assert binary_result.returncode == 0
        assert python_result.returncode == 0
        assert "Number of atoms: 3" in binary_result.stdout
        assert "Number of frames: 5" in binary_result.stdout

    def test_comprehensive_workflow_with_all_commands(self, temp_dir):
        """Test all CLI commands in a comprehensive workflow"""
        # Create a test trajectory with known properties
        trajectory = create_test_trajectory(15, 3)
        original_file = os.path.join(temp_dir, "workflow.taco")
        taco_format.write(original_file, trajectory, lossless=True)

        # 1. Check file integrity
        check_result = run_taco_command(["check", original_file, "--verbose"])
        assert "File validation successful" in check_result.stdout
        assert "Validating 15 frames" in check_result.stdout

        # 2. Get comprehensive file info
        info_result = run_taco_command(["info", original_file])
        assert "Number of atoms: 3" in info_result.stdout
        assert "Number of frames: 15" in info_result.stdout
        assert "Precision mode: Full" in info_result.stdout

        # 3. Get detailed statistics
        stats_result = run_taco_command(["stats", original_file, "--detailed"])
        assert "Position ranges:" in stats_result.stdout
        assert "Frames with velocities:" in stats_result.stdout
        assert "Frames with forces:" in stats_result.stdout
        assert "Detailed Frame Information" in stats_result.stdout

        # 4. Extract subset with step parameter
        subset_file = os.path.join(temp_dir, "workflow_subset.taco")
        extract_result = run_taco_command(
            [
                "extract",
                original_file,
                subset_file,
                "--start",
                "2",
                "--end",
                "12",
                "--step",
                "3",
            ]
        )
        assert "Successfully extracted" in extract_result.stdout

        # 5. Verify extracted file has correct number of frames
        subset_info = run_taco_command(["info", subset_file])
        assert "Number of frames: 4" in subset_info.stdout  # frames 2, 5, 8, 11

        # 6. Convert with high compression and full precision
        converted_file = os.path.join(temp_dir, "workflow_converted.taco")
        convert_result = run_taco_command(
            ["convert", subset_file, converted_file, "--compression", "9", "--lossless"]
        )
        assert "Successfully converted" in convert_result.stdout

        # 7. Final integrity check
        final_check = run_taco_command(["check", converted_file])
        assert "File validation successful" in final_check.stdout

        # 8. Verify data integrity through Python API
        final_trajectory = taco_format.read(converted_file)
        assert len(final_trajectory) == 4

        # Verify that the extracted frames have the expected properties
        for atoms in final_trajectory:
            assert len(atoms) == 3  # 3 atoms (H2O)
            assert atoms.get_velocities() is not None
            assert atoms.get_forces() is not None
