"""
Gene BED Parser for MSK-ACCESS Panel Data.

Parses MSK-ACCESS target BED files and groups regions by gene.
Supports V1 (xs1) and V2 (xs2) naming conventions.

V1 format: exon_GENE_exon_probe (e.g., exon_MTOR_48a.1_2)
V2 format: GENE_target_NN (e.g., MTOR_target_02)

Non-gene regions (FP_, msi_, snp_FP_, BAT-*) are filtered out.
"""

import gzip
import re
import logging
from pathlib import Path
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Generator, Dict, List, Optional, IO

logger = logging.getLogger(__name__)


@contextmanager
def _open_bed(path: Path) -> Generator[IO[str], None, None]:
    """Open a BED file, handling gzip compression transparently."""
    path = Path(path)
    if str(path).endswith(".gz"):
        f = gzip.open(path, "rt")
    else:
        f = open(path, "r")
    try:
        yield f
    finally:
        f.close()


@dataclass
class Region:
    """A genomic region with gene annotation."""

    chrom: str
    start: int
    end: int
    name: str
    gene: str


def detect_assay(target_bed: Path) -> Optional[str]:
    """
    Detect assay type from target BED filename.

    Looks for 'xs1' or 'xs2' in filename, or legacy 'v1.0'/'v2.0' patterns.

    Args:
        target_bed: Path to target regions BED file

    Returns:
        Assay code ('xs1' or 'xs2') or None if not detected

    Examples:
        >>> detect_assay(Path("xs2.targets.bed"))
        'xs2'
        >>> detect_assay(Path("MSK-ACCESS-v1_0-probe-A.sorted.bed"))
        'xs1'
        >>> detect_assay(Path("custom_panel.bed"))
        None
    """
    filename = target_bed.name.lower()

    # Direct match for standardized names
    if "xs1" in filename:
        logger.debug(f"Detected assay 'xs1' from filename: {filename}")
        return "xs1"
    if "xs2" in filename:
        logger.debug(f"Detected assay 'xs2' from filename: {filename}")
        return "xs2"

    # Legacy patterns
    if (
        "v1.0" in filename
        or "v1_0" in filename
        or "access-v1" in filename.replace("_", "-")
    ):
        logger.debug(f"Detected assay 'xs1' from legacy pattern: {filename}")
        return "xs1"
    if (
        "v2.0" in filename
        or "v2_0" in filename
        or "access-v2" in filename.replace("_", "-")
    ):
        logger.debug(f"Detected assay 'xs2' from legacy pattern: {filename}")
        return "xs2"

    logger.debug(f"Could not detect assay from filename: {filename}")
    return None


def detect_version_from_content(bed_path: Path) -> str:
    """
    Detect BED version (V1, V2, or GENE_BED) by examining content.

    V1 uses: exon_GENE_exon_probe pattern
    V2 uses: GENE_target_NN pattern
    GENE_BED: Pre-parsed gene BED with #chrom\\tstart\\tend\\tgene\\tname header

    Args:
        bed_path: Path to target BED file

    Returns:
        'V1', 'V2', or 'GENE_BED' based on content analysis
    """
    v1_pattern = re.compile(r"^exon_[A-Z0-9]+_\d+")
    v2_pattern = re.compile(r"^[A-Z0-9]+_target_\d+$")

    v1_count = 0
    v2_count = 0
    is_gene_bed = False

    with _open_bed(bed_path) as f:
        for i, line in enumerate(f):
            if i >= 100:  # Sample first 100 lines
                break

            # Check for gene BED header
            if line.startswith("#chrom"):
                if "gene" in line.lower():
                    is_gene_bed = True
                    break
                continue

            if line.startswith("#") or not line.strip():
                continue

            fields = line.strip().split("\t")
            if len(fields) < 4:
                continue

            # V1 format has name in column 5 (strand in column 4)
            # V2 format has name in column 4
            # Check both columns for pattern matching
            for idx in [4, 3]:
                if len(fields) > idx:
                    name = fields[idx]
                    if v1_pattern.match(name):
                        v1_count += 1
                        break
                    elif v2_pattern.match(name):
                        v2_count += 1
                        break

    if is_gene_bed:
        logger.debug("Detected GENE_BED format (pre-parsed)")
        return "GENE_BED"

    version = "V1" if v1_count > v2_count else "V2"
    logger.debug(f"Detected version {version} (V1: {v1_count}, V2: {v2_count})")
    return version


def _is_gene_region(name: str) -> bool:
    """
    Check if region name represents a gene target (not fingerprint/MSI).

    Excludes:
      - FP_* (fingerprint probes)
      - snp_FP_* (SNP fingerprint probes)
      - msi_* (microsatellite instability markers)
      - BAT-* (BAT markers)

    Args:
        name: Region name from BED file

    Returns:
        True if this is a gene region, False otherwise
    """
    name_lower = name.lower()

    # Exclude patterns
    exclude_prefixes = ("fp_", "snp_fp_", "msi_", "bat-")

    for prefix in exclude_prefixes:
        if name_lower.startswith(prefix):
            return False

    return True


def _extract_gene_v1(name: str) -> Optional[str]:
    """
    Extract gene name from V1 format: exon_GENE_exon_probe

    Examples:
        exon_MTOR_48a.1_2 -> MTOR
        exon_ARID1A_1a.1_1 -> ARID1A
        exon_H3F3A_2_1 -> H3F3A
    """
    # Pattern: exon_GENE_...
    match = re.match(r"^exon_([A-Z0-9]+)_", name, re.IGNORECASE)
    if match:
        return match.group(1).upper()
    return None


def _extract_gene_v2(name: str) -> Optional[str]:
    """
    Extract gene name from V2 format: GENE_target_NN

    Examples:
        MTOR_target_02 -> MTOR
        ARID1A_target_01 -> ARID1A
    """
    # Pattern: GENE_target_NN
    match = re.match(r"^([A-Z0-9]+)_target_\d+$", name, re.IGNORECASE)
    if match:
        return match.group(1).upper()
    return None


def parse_gene_bed(bed_path: Path, version: str = "auto") -> Dict[str, List[Region]]:
    """
    Parse MSK-ACCESS target BED file and group regions by gene.

    Filters out non-gene regions (FP, MSI, BAT markers) and groups
    the remaining regions by their associated gene name.

    Args:
        bed_path: Path to target BED file
        version: 'V1', 'V2', or 'auto' for auto-detection

    Returns:
        Dictionary mapping gene names to lists of Region objects

    Example:
        >>> genes = parse_gene_bed(Path("xs2.targets.bed"))
        >>> genes['MTOR']
        [Region(chrom='1', start=11168235, end=11168345, ...)]
    """
    bed_path = Path(bed_path)

    if not bed_path.exists():
        raise FileNotFoundError(f"BED file not found: {bed_path}")

    # Auto-detect version if needed
    if version == "auto":
        version = detect_version_from_content(bed_path)
        logger.info(f"Auto-detected BED format: {version}")

    genes: Dict[str, List[Region]] = {}
    total_regions = 0
    filtered_regions = 0

    with _open_bed(bed_path) as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue

            fields = line.strip().split("\t")
            if len(fields) < 3:
                continue

            total_regions += 1

            chrom = fields[0]
            try:
                start = int(fields[1])
                end = int(fields[2])
            except ValueError:
                continue

            # Handle different formats
            if version == "GENE_BED":
                # Pre-parsed format: chrom, start, end, gene, name
                gene = fields[3] if len(fields) > 3 else None
                name = fields[4] if len(fields) > 4 else gene or ""
            else:
                # Get region name - V1 has name in column 5, V2 in column 4
                name = ""
                if len(fields) >= 5:
                    # Check if column 5 looks like a gene name (V1 format)
                    if fields[4] and (
                        fields[4].startswith("exon_") or "target" in fields[3].lower()
                    ):
                        name = fields[4] if fields[4].startswith("exon_") else fields[3]
                    else:
                        name = fields[3]
                elif len(fields) >= 4:
                    name = fields[3]

                # Filter non-gene regions
                if not _is_gene_region(name):
                    filtered_regions += 1
                    continue

                # Extract gene name based on version
                extract_gene = _extract_gene_v1 if version == "V1" else _extract_gene_v2
                gene = extract_gene(name)
                if not gene:
                    # Fallback: try other version's pattern
                    alt_extract = (
                        _extract_gene_v2 if version == "V1" else _extract_gene_v1
                    )
                    gene = alt_extract(name)

            if not gene:
                logger.debug(f"Could not extract gene from: {name}")
                continue

            region = Region(chrom=chrom, start=start, end=end, name=name, gene=gene)

            if gene not in genes:
                genes[gene] = []
            genes[gene].append(region)

    logger.info(
        f"Parsed {len(genes)} genes from {total_regions} regions "
        f"(filtered {filtered_regions} non-gene regions)"
    )

    return genes


def get_gene_count(genes: Dict[str, List[Region]]) -> Dict[str, int]:
    """Get count of regions per gene."""
    return {gene: len(regions) for gene, regions in genes.items()}


def write_gene_bed(
    genes: Dict[str, List[Region]],
    output_path: Path,
    sort: bool = True,
    compress: bool = True,
) -> int:
    """
    Write gene-grouped regions to a BED file (optionally BGZF compressed).

    Args:
        genes: Dictionary from parse_gene_bed()
        output_path: Output path for gene BED file
        sort: If True, sort by chromosome and position
        compress: If True, write as BGZF (.bed.gz)

    Returns:
        Number of regions written
    """
    import gzip

    output_path = Path(output_path)

    # Collect all regions
    all_regions = []
    for gene, regions in genes.items():
        for region in regions:
            all_regions.append(
                (region.chrom, region.start, region.end, gene, region.name)
            )

    # Sort by chromosome and position
    if sort:
        # Custom chromosome sort (1, 2, ..., 22, X, Y)
        def chrom_key(r):
            chrom = r[0].replace("chr", "")
            if chrom.isdigit():
                return (0, int(chrom), r[1])
            return (1, chrom, r[1])

        all_regions.sort(key=chrom_key)

    # Build content
    lines = ["#chrom\tstart\tend\tgene\tname\n"]
    for chrom, start, end, gene, name in all_regions:
        lines.append(f"{chrom}\t{start}\t{end}\t{gene}\t{name}\n")
    content = "".join(lines)

    # Write output (compressed or plain)
    if compress or str(output_path).endswith(".gz"):
        with gzip.open(output_path, "wt") as f:
            f.write(content)
    else:
        with open(output_path, "w") as f:
            f.write(content)

    logger.info(f"Wrote {len(all_regions)} gene regions to {output_path}")
    return len(all_regions)


def get_bundled_gene_bed(assay: str, genome: str = "GRCh37") -> Optional[Path]:
    """
    Get path to bundled gene BED file for a known assay.

    Args:
        assay: Assay code ('xs1' or 'xs2')
        genome: Genome build ('GRCh37' or 'GRCh38')

    Returns:
        Path to bundled gene BED file, or None if not found

    Example:
        >>> path = get_bundled_gene_bed('xs2', 'GRCh37')
        >>> path
        PosixPath('.../data/genes/GRCh37/xs2.genes.bed.gz')
    """
    import importlib.resources

    assay = assay.lower()
    if assay not in ("xs1", "xs2"):
        logger.debug(f"Unknown assay '{assay}', no bundled gene BED")
        return None

    # Use importlib.resources for package data
    try:
        data_pkg = f"krewlyzer.data.genes.{genome}"
        filename = f"{assay}.genes.bed.gz"

        # Try to get the file path
        ref = importlib.resources.files(data_pkg) / filename
        if ref.is_file():
            # Return actual path for downstream compatibility
            return Path(str(ref))
    except (ModuleNotFoundError, FileNotFoundError, TypeError) as e:
        logger.debug(f"importlib.resources lookup failed for {assay}/{genome}: {e}")

    # Fallback: check relative to this file
    data_dir = Path(__file__).parent.parent / "data" / "genes" / genome
    bundled_path = data_dir / f"{assay}.genes.bed.gz"

    if bundled_path.exists():
        logger.debug(f"Found bundled gene BED: {bundled_path}")
        return bundled_path

    logger.debug(f"Bundled gene BED not found for assay={assay}, genome={genome}")
    return None


def load_gene_bed(
    target_bed: Optional[Path] = None,
    gene_bed: Optional[Path] = None,
    assay: Optional[str] = None,
    genome: str = "GRCh37",
) -> Dict[str, List[Region]]:
    """
    Load gene regions from appropriate source.

    Priority:
    1. Explicit gene_bed file if provided
    2. Bundled gene BED if assay is xs1/xs2
    3. Parse target_bed on-the-fly

    Args:
        target_bed: Path to target regions BED (for on-the-fly parsing)
        gene_bed: Explicit path to pre-parsed gene BED
        assay: Assay code ('xs1', 'xs2', or None for auto-detect)
        genome: Genome build for bundled asset lookup

    Returns:
        Dictionary mapping gene names to Region lists

    Raises:
        ValueError: If no valid source is available
    """
    # Option 1: Explicit gene BED file
    if gene_bed and Path(gene_bed).exists():
        logger.info(f"Loading gene BED from explicit path: {gene_bed}")
        return parse_gene_bed(gene_bed)

    # Option 2: Auto-detect or use provided assay for bundled file
    if assay is None and target_bed:
        assay = detect_assay(Path(target_bed))

    if assay:
        bundled = get_bundled_gene_bed(assay, genome)
        if bundled:
            logger.info(f"Loading bundled gene BED for assay {assay}: {bundled}")
            return parse_gene_bed(bundled)

    # Option 3: Parse target BED on-the-fly
    if target_bed and Path(target_bed).exists():
        logger.info(f"Parsing target BED on-the-fly: {target_bed}")
        return parse_gene_bed(target_bed)

    raise ValueError(
        "No valid gene source: provide gene_bed, target_bed, or known assay"
    )
