"""
Guide extraction module for CRISPR gRNA design.

Extracts candidate 20nt gRNA spacers at NGG PAM sites from genomic regions.
"""

import re

import pyfastx

# Padding needed for 30nt context extraction (4bp upstream + 3bp downstream of PAM)
CONTEXT_PADDING = 7


def open_fasta(fasta_path):
    """
    Open a FASTA file. pyfastx handles gzip natively without decompression.

    Args:
        fasta_path: Path to FASTA file (.fa, .fa.gz, or .fasta)

    Returns:
        pyfastx.Fasta object
    """
    return pyfastx.Fasta(fasta_path)


def reverse_complement(seq):
    """Return the reverse complement of a DNA sequence."""
    complement = {'A': 'T', 'T': 'A', 'G': 'C', 'C': 'G',
                  'a': 't', 't': 'a', 'g': 'c', 'c': 'g',
                  'N': 'N', 'n': 'n'}
    return ''.join(complement.get(base, 'N') for base in reversed(seq))


def extract_guides(chrom, start, end, region_name, fasta_path):
    """
    Given a genomic region, extract all candidate 20nt gRNA spacers
    at NGG PAM sites on both strands.

    Args:
        chrom: Chromosome name (e.g., 'chr1')
        start: 0-based start position
        end: End position (half-open interval)
        region_name: Name identifier for the region (used as guideSet)
        fasta_path: Path to genome FASTA file

    Returns:
        A list of dicts, each with keys:
            chr, start, end, locus, strand,
            GuideSequenceWithPAM, guideSet, protospacer, context_sequence

    Notes:
        - BED coordinates are 0-based half-open
        - The start and end in output rows refer to the protospacer only (not the PAM)
        - GuideSequenceWithPAM is 23 nt (20 nt protospacer + 3 nt PAM)
        - context_sequence is 30 nt (4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream)
          for use with RS3 scoring
    """
    fasta = open_fasta(fasta_path)

    # Fetch the sequence for the region with padding for context extraction
    # We need 4bp upstream and 3bp downstream of the PAM for 30nt context
    padded_start = max(0, start - CONTEXT_PADDING)
    chrom_length = len(fasta[chrom])
    padded_end = min(chrom_length, end + CONTEXT_PADDING)

    # Fetch padded sequence
    padded_sequence = str(fasta[chrom][padded_start:padded_end]).upper()

    # Calculate offset from padded coordinates to original coordinates
    offset = start - padded_start

    # Extract the original region sequence for guide finding
    sequence = padded_sequence[offset:offset + (end - start)]
    region_length = end - start

    guides = []

    # Scan forward strand for NGG PAM sites
    # Pattern: 20bp protospacer followed by NGG
    # We need at least 23 bases (20 + 3) to have a valid guide
    for match in re.finditer(r'(?=(.{20}[ACGT]GG))', sequence):
        guide_with_pam = match.group(1)
        pos_in_seq = match.start()  # Position of first base of protospacer in sequence

        # Protospacer is positions 0-19, PAM is positions 20-22
        protospacer = guide_with_pam[:20]

        # Calculate genomic coordinates for protospacer only
        guide_start = start + pos_in_seq
        guide_end = guide_start + 20  # Half-open, so end is exclusive

        locus = f"{chrom}:{guide_start}-{guide_end}:+"

        # Extract 30nt context for RS3 scoring
        # Format: 4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream
        context_start_in_padded = offset + pos_in_seq - 4
        context_end_in_padded = offset + pos_in_seq + 26

        if context_start_in_padded >= 0 and context_end_in_padded <= len(padded_sequence):
            context_sequence = padded_sequence[context_start_in_padded:context_end_in_padded]
        else:
            # Pad with N if context extends beyond available sequence
            context_sequence = ""
            for i in range(context_start_in_padded, context_end_in_padded):
                if 0 <= i < len(padded_sequence):
                    context_sequence += padded_sequence[i]
                else:
                    context_sequence += 'N'

        guides.append({
            'chr': chrom,
            'start': guide_start,
            'end': guide_end,
            'locus': locus,
            'strand': '+',
            'GuideSequenceWithPAM': guide_with_pam,
            'guideSet': region_name,
            'protospacer': protospacer,
            'context_sequence': context_sequence
        })

    # Scan reverse strand for NGG PAM sites
    # On reverse strand, PAM is CCN in forward coordinates
    # We look for CCN followed by 20bp protospacer
    for match in re.finditer(r'(?=(CC[ACGT].{20}))', sequence):
        guide_region = match.group(1)
        pos_in_seq = match.start()  # Position of first C of CCN

        # In forward coordinates: CC[N] is at pos_in_seq:pos_in_seq+3
        # Protospacer (reverse strand) is at pos_in_seq+3:pos_in_seq+23

        # Get the reverse complement to get the actual guide sequence
        # The guide with PAM on reverse strand: take the 23bp and reverse complement
        guide_with_pam = reverse_complement(guide_region)  # This gives us protospacer + NGG
        protospacer = guide_with_pam[:20]

        # Genomic coordinates for protospacer (in forward strand coordinates)
        # The protospacer on the - strand spans from pos+3 to pos+23 in forward coords
        guide_start = start + pos_in_seq + 3
        guide_end = guide_start + 20

        locus = f"{chrom}:{guide_start}-{guide_end}:-"

        # Extract 30nt context for RS3 scoring (reverse strand)
        # In forward coords: 3bp downstream of protospacer, protospacer (23bp including PAM),
        # 4bp upstream of CCN, then reverse complement the whole thing
        # Forward coords: [pos-3, pos+27) then reverse complement
        context_start_in_padded = offset + pos_in_seq - 3
        context_end_in_padded = offset + pos_in_seq + 27

        if context_start_in_padded >= 0 and context_end_in_padded <= len(padded_sequence):
            context_fwd = padded_sequence[context_start_in_padded:context_end_in_padded]
        else:
            # Pad with N if context extends beyond available sequence
            context_fwd = ""
            for i in range(context_start_in_padded, context_end_in_padded):
                if 0 <= i < len(padded_sequence):
                    context_fwd += padded_sequence[i]
                else:
                    context_fwd += 'N'

        context_sequence = reverse_complement(context_fwd)

        guides.append({
            'chr': chrom,
            'start': guide_start,
            'end': guide_end,
            'locus': locus,
            'strand': '-',
            'GuideSequenceWithPAM': guide_with_pam,
            'guideSet': region_name,
            'protospacer': protospacer,
            'context_sequence': context_sequence
        })

    return guides


def extract_guides_from_bed(bed_path, fasta_path):
    """
    Extract guides from all regions in a BED file.

    Args:
        bed_path: Path to BED file with regions
        fasta_path: Path to genome FASTA file

    Returns:
        A list of guide dicts from all regions
    """
    all_guides = []

    with open(bed_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue

            fields = line.split('\t')
            chrom = fields[0]
            start = int(fields[1])
            end = int(fields[2])

            # Use region name from column 4 if available, otherwise create one
            if len(fields) >= 4 and fields[3]:
                region_name = fields[3]
            else:
                region_name = f"{chrom}:{start}-{end}"

            guides = extract_guides(chrom, start, end, region_name, fasta_path)
            all_guides.extend(guides)

    return all_guides
