"""
Tests for guide extraction module.
"""

import os
import tempfile
import pytest

from crisprdesigner2 import extract_guides, reverse_complement


class TestReverseComplement:
    """Tests for reverse_complement function."""

    def test_basic_reverse_complement(self):
        assert reverse_complement('ATGC') == 'GCAT'

    def test_palindrome(self):
        assert reverse_complement('ATAT') == 'ATAT'

    def test_single_base(self):
        assert reverse_complement('A') == 'T'
        assert reverse_complement('G') == 'C'

    def test_with_n(self):
        assert reverse_complement('ATNGC') == 'GCNAT'


class TestExtractGuides:
    """Tests for extract_guides function."""

    @pytest.fixture
    def temp_fasta(self):
        """Create a temporary FASTA file with known sequence."""
        # Create a sequence with known NGG PAM sites
        # Positions marked:
        # Forward strand NGG at position 25 (protospacer 5-25)
        # Reverse strand (CCN at position 50, protospacer 53-73 in forward coords)
        #
        # Sequence: 100 bp
        # AAAAA ACGTACGTACGTACGTACGT NGG ACGTACGTACGTACGTACGT CC A ACGTACGTACGTACGTACGT AAAAA
        # 0     5                    25   28                   48    51                   71
        #
        # Forward guide: protospacer at 5-25, PAM at 25-28
        # Reverse guide: CCN at 48-51, protospacer at 51-71

        sequence = (
            "AAAAA"                      # 0-5
            "ACGTACGTACGTACGTACGT"       # 5-25 (protospacer for forward)
            "AGG"                        # 25-28 (PAM for forward)
            "ACGTACGTACGTACGTACGT"       # 28-48
            "CCA"                        # 48-51 (PAM for reverse = TGG on reverse strand)
            "ACGTACGTACGTACGTACGT"       # 51-71 (protospacer for reverse)
            "AAAAAAAAAAAAAAAAAAAAAAAAAAAA"  # 71-99 padding
        )

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            f.write(sequence + "\n")
            fasta_path = f.name

        yield fasta_path, sequence

        # Cleanup
        os.unlink(fasta_path)
        if os.path.exists(fasta_path + '.fai'):
            os.unlink(fasta_path + '.fai')

    def test_forward_strand_extraction(self, temp_fasta):
        """Test extraction of forward strand guides."""
        fasta_path, sequence = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)

        # Find forward strand guides
        forward_guides = [g for g in guides if g['strand'] == '+']

        # Should find the AGG PAM
        assert len(forward_guides) >= 1

        # Check one of the forward guides
        agg_guide = None
        for g in forward_guides:
            if g['GuideSequenceWithPAM'].endswith('AGG'):
                if g['start'] == 5:
                    agg_guide = g
                    break

        assert agg_guide is not None, f"Expected forward guide not found. Found: {forward_guides}"
        assert agg_guide['protospacer'] == 'ACGTACGTACGTACGTACGT'
        assert len(agg_guide['GuideSequenceWithPAM']) == 23
        assert agg_guide['start'] == 5
        assert agg_guide['end'] == 25
        assert agg_guide['guideSet'] == 'test_region'

    def test_reverse_strand_extraction(self, temp_fasta):
        """Test extraction of reverse strand guides."""
        fasta_path, sequence = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)

        # Find reverse strand guides
        reverse_guides = [g for g in guides if g['strand'] == '-']

        # Should find guides from CCN patterns
        assert len(reverse_guides) >= 1

        # Check that all reverse guides have correct format
        for g in reverse_guides:
            assert len(g['GuideSequenceWithPAM']) == 23
            assert g['GuideSequenceWithPAM'].endswith('GG')  # Should end with GG (NGG PAM)
            assert len(g['protospacer']) == 20

    def test_no_pam_sites(self):
        """Test that a sequence with no NGG PAMs returns empty list."""
        # Create a sequence with no NGG or CCN patterns
        sequence = "ATATATAT" * 20  # 160 bp of AT repeats - no NGG or CCN

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            f.write(sequence + "\n")
            fasta_path = f.name

        try:
            guides = extract_guides('chr1', 0, 160, 'test_region', fasta_path)
            assert guides == []
        finally:
            os.unlink(fasta_path)
            if os.path.exists(fasta_path + '.fai'):
                os.unlink(fasta_path + '.fai')

    def test_guide_sequence_length(self, temp_fasta):
        """Test that GuideSequenceWithPAM is always 23 nt."""
        fasta_path, _ = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)

        for guide in guides:
            assert len(guide['GuideSequenceWithPAM']) == 23, \
                f"GuideSequenceWithPAM should be 23 nt, got {len(guide['GuideSequenceWithPAM'])}"

    def test_protospacer_length(self, temp_fasta):
        """Test that protospacer is always 20 nt."""
        fasta_path, _ = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)

        for guide in guides:
            assert len(guide['protospacer']) == 20, \
                f"Protospacer should be 20 nt, got {len(guide['protospacer'])}"

    def test_boundary_conditions_start(self):
        """Test guide extraction at the very start of a region."""
        # Put an NGG PAM so protospacer starts at position 0
        sequence = "ACGTACGTACGTACGTACGT" + "AGG" + "A" * 77

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            f.write(sequence + "\n")
            fasta_path = f.name

        try:
            guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)
            forward_guides = [g for g in guides if g['strand'] == '+']

            # Should find guide at position 0
            start_guide = [g for g in forward_guides if g['start'] == 0]
            assert len(start_guide) == 1
            assert start_guide[0]['protospacer'] == 'ACGTACGTACGTACGTACGT'
        finally:
            os.unlink(fasta_path)
            if os.path.exists(fasta_path + '.fai'):
                os.unlink(fasta_path + '.fai')

    def test_boundary_conditions_end(self):
        """Test guide extraction at the very end of a region."""
        # Put an NGG PAM at the very end (PAM ends at position 100)
        sequence = "A" * 77 + "ACGTACGTACGTACGTACGT" + "AGG"

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            f.write(sequence + "\n")
            fasta_path = f.name

        try:
            guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)
            forward_guides = [g for g in guides if g['strand'] == '+']

            # Should find guide with protospacer at 77-97
            end_guide = [g for g in forward_guides if g['start'] == 77]
            assert len(end_guide) == 1
            assert end_guide[0]['end'] == 97
        finally:
            os.unlink(fasta_path)
            if os.path.exists(fasta_path + '.fai'):
                os.unlink(fasta_path + '.fai')

    def test_reverse_complement_coordinate_conversion(self):
        """Test that reverse strand coordinates are correctly converted."""
        # Create a simple sequence with a known CCN pattern
        # CCN at position 10-13, protospacer at 13-33
        sequence = "AAAAAAAAAA" + "CCA" + "GCTAGCTAGCTAGCTAGCTA" + "A" * 67

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            f.write(sequence + "\n")
            fasta_path = f.name

        try:
            guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)
            reverse_guides = [g for g in guides if g['strand'] == '-']

            # Find the guide from our CCA pattern
            cca_guide = [g for g in reverse_guides if g['start'] == 13]
            assert len(cca_guide) == 1

            guide = cca_guide[0]
            # Protospacer should be at positions 13-33 in forward coordinates
            assert guide['start'] == 13
            assert guide['end'] == 33

            # The protospacer sequence should be reverse complement of the forward seq
            forward_seq = "GCTAGCTAGCTAGCTAGCTA"
            expected_protospacer = reverse_complement(forward_seq)
            assert guide['protospacer'] == expected_protospacer

            # GuideSequenceWithPAM should include the PAM (TGG on reverse strand)
            # CCA on forward is TGG on reverse
            assert guide['GuideSequenceWithPAM'] == expected_protospacer + "TGG"
        finally:
            os.unlink(fasta_path)
            if os.path.exists(fasta_path + '.fai'):
                os.unlink(fasta_path + '.fai')

    def test_locus_format(self, temp_fasta):
        """Test that locus is correctly formatted."""
        fasta_path, _ = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)

        for guide in guides:
            locus = guide['locus']
            # Format should be chr:start-end:strand
            parts = locus.split(':')
            assert len(parts) == 3
            assert parts[0] == guide['chr']
            assert parts[2] in ['+', '-']

            coords = parts[1].split('-')
            assert len(coords) == 2
            assert int(coords[0]) == guide['start']
            assert int(coords[1]) == guide['end']

    def test_multiple_pam_sites(self):
        """Test extraction with multiple PAM sites in close proximity."""
        # Create sequence with multiple NGG sites
        sequence = "ACGTACGTACGTACGTACGTNGGNGGACGTACGTACGTACGTACGTNGGA" + "A" * 50

        with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
            f.write(">chr1\n")
            # Replace N with A for valid sequence
            f.write(sequence.replace('N', 'A') + "\n")
            fasta_path = f.name

        try:
            guides = extract_guides('chr1', 0, 100, 'test_region', fasta_path)
            # Should find multiple guides
            assert len(guides) > 0
        finally:
            os.unlink(fasta_path)
            if os.path.exists(fasta_path + '.fai'):
                os.unlink(fasta_path + '.fai')

    def test_guide_set_assignment(self, temp_fasta):
        """Test that guideSet is correctly assigned from region_name."""
        fasta_path, _ = temp_fasta

        guides = extract_guides('chr1', 0, 100, 'my_custom_region', fasta_path)

        for guide in guides:
            assert guide['guideSet'] == 'my_custom_region'


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
