"""
CyVCF2-based VCF/BCF reader implementation for WASP2.

This module provides CyVCF2Source, a high-performance VariantSource implementation
using cyvcf2 library (6.9x faster than pysam). Offers the same interface as VCFSource
but with significantly improved performance for VCF parsing operations.

Performance:
    - 6.9x faster than pysam for VCF parsing
    - Zero-copy numpy array access to genotype data
    - Direct memory access to htslib structures

Requirements:
    pip install wasp2[cyvcf2]
"""

import subprocess
from collections.abc import Iterator
from pathlib import Path

try:
    import cyvcf2

    CYVCF2_AVAILABLE = True
except ImportError:
    CYVCF2_AVAILABLE = False

from .variant_source import (
    Genotype,
    Variant,
    VariantGenotype,
    VariantSource,
)

# Only register if cyvcf2 is available
if CYVCF2_AVAILABLE:

    @VariantSource.register(
        "cyvcf2.vcf", "cyvcf2.vcf.gz", "cyvcf2.vcf.bgz", "cyvcf2.bcf", "cyvcf2.bcf.gz"
    )
    class CyVCF2Source(VariantSource):
        """High-performance VariantSource implementation using cyvcf2.

        Reads variant data from VCF/BCF files using cyvcf2 (cython + htslib),
        providing 6.9x faster performance compared to pysam. Uses zero-copy
        numpy arrays for efficient genotype access.

        The class handles:
        - Standard VCF/BCF parsing (faster than pysam)
        - Genotype extraction via numpy arrays
        - Sample-specific filtering
        - Heterozygous-only filtering
        - Region queries (if indexed)
        - BED format export using bcftools for efficiency

        Attributes:
            path: Path to the VCF/BCF file
            vcf: cyvcf2.VCF handle
            _samples: Cached list of sample IDs
            _variant_count: Cached variant count (lazy computed)

        Example:
            >>> with CyVCF2Source("variants.vcf.gz") as vcf:
            ...     for vg in vcf.iter_variants(het_only=True):
            ...         print(f"{vg.variant.chrom}:{vg.variant.pos}")
        """

        def __init__(self, path: str, **kwargs) -> None:
            """Initialize CyVCF2 source.

            Args:
                path: Path to VCF/BCF file (str or Path-like)
                **kwargs: Additional arguments (reserved for future use)

            Raises:
                ImportError: If cyvcf2 is not installed
                FileNotFoundError: If file doesn't exist
                ValueError: If file cannot be opened or parsed
            """
            if not CYVCF2_AVAILABLE:
                raise ImportError(
                    "cyvcf2 is not installed. Install with: pip install wasp2[cyvcf2]"
                )

            self.path = Path(path)

            # Open VCF file with cyvcf2
            try:
                self.vcf = cyvcf2.VCF(str(self.path))
            except Exception as e:
                raise ValueError(f"Failed to open VCF file {self.path}: {e}") from e

            # Cache samples from header
            self._samples = self.vcf.samples

            # Lazy-computed variant count
            self._variant_count: int | None = None

            # Track if iterator has been used (cyvcf2 doesn't support seek)
            self._iterator_used = False

        @property
        def samples(self) -> list[str]:
            """Get list of sample IDs from VCF header.

            Returns:
                List of sample ID strings in file order
            """
            return list(self._samples)

        @property
        def variant_count(self) -> int:
            """Get total number of variants in the file.

            Counts variants by iterating through the file. Result is cached
            for subsequent calls.

            Returns:
                Total number of variants
            """
            if self._variant_count is None:
                # Count variants by iterating through file
                count = 0
                for _ in self.vcf:
                    count += 1
                self._variant_count = count

                # Mark iterator as used and reopen for future use
                self._iterator_used = True
                self.vcf.close()
                self.vcf = cyvcf2.VCF(str(self.path))
                self._iterator_used = False

            return self._variant_count

        @property
        def sample_count(self) -> int:
            """Get total number of samples.

            Returns:
                Total number of samples
            """
            return len(self._samples)

        def iter_variants(
            self, samples: list[str] | None = None, het_only: bool = False
        ) -> Iterator[VariantGenotype]:
            """Iterate over variants with optional filtering.

            Yields one VariantGenotype per variant for the first sample in the list
            (or first sample in file if samples=None).

            Args:
                samples: Optional list of sample IDs. If None, uses first sample.
                        Currently only supports single sample iteration.
                het_only: If True, only yield heterozygous variants

            Yields:
                VariantGenotype objects for each variant

            Example:
                >>> for vg in source.iter_variants(samples=["sample1"], het_only=True):
                ...     print(vg.variant.pos, vg.genotype)
            """
            # Determine which sample to iterate
            if samples is None:
                target_samples = [self._samples[0]] if self._samples else []
            else:
                # Validate samples exist
                for s in samples:
                    if s not in self._samples:
                        raise ValueError(f"Sample '{s}' not found in VCF")
                target_samples = samples

            if not target_samples:
                return

            # Currently support single sample iteration
            sample_id = target_samples[0]
            sample_idx = self._samples.index(sample_id)

            # cyvcf2 doesn't support rewind/seek, so reopen if iterator was used
            if self._iterator_used:
                self.vcf.close()
                self.vcf = cyvcf2.VCF(str(self.path))
                self._iterator_used = False

            # Mark iterator as used
            self._iterator_used = True

            # Iterate through VCF records
            for variant in self.vcf:
                # Get genotype using numpy array (zero-copy access)
                # gt_types: 0=HOM_REF, 1=HET, 2=HOM_UNKNOWN, 3=HOM_ALT
                gt_type = variant.gt_types[sample_idx]

                # Convert cyvcf2 gt_type to our Genotype enum
                if gt_type == 0:
                    genotype = Genotype.HOM_REF
                elif gt_type == 1:
                    genotype = Genotype.HET
                elif gt_type == 3:
                    genotype = Genotype.HOM_ALT
                else:  # gt_type == 2 (HOM_UNKNOWN) or other
                    genotype = Genotype.MISSING

                # Filter by het_only if requested
                if het_only and genotype != Genotype.HET:
                    continue

                # Create Variant object (use first ALT if multi-allelic)
                alt = variant.ALT[0] if variant.ALT else variant.REF
                var = Variant(
                    chrom=variant.CHROM,
                    pos=variant.POS,
                    ref=variant.REF,
                    alt=alt,
                    id=variant.ID if variant.ID else None,
                )

                # Get allele sequences from genotype array
                # gt_bases gives actual allele sequences for each sample
                gt_bases = variant.gt_bases[sample_idx]
                if gt_bases and "/" in gt_bases:
                    alleles = gt_bases.split("/")
                    allele1 = alleles[0] if alleles[0] != "." else None
                    allele2 = alleles[1] if len(alleles) > 1 and alleles[1] != "." else None
                elif gt_bases and "|" in gt_bases:
                    alleles = gt_bases.split("|")
                    allele1 = alleles[0] if alleles[0] != "." else None
                    allele2 = alleles[1] if len(alleles) > 1 and alleles[1] != "." else None
                else:
                    allele1, allele2 = None, None

                yield VariantGenotype(
                    variant=var, genotype=genotype, allele1=allele1, allele2=allele2
                )

        def get_genotype(self, sample: str, chrom: str, pos: int) -> Genotype:
            """Get genotype for a specific sample at a genomic position.

            Args:
                sample: Sample ID
                chrom: Chromosome name
                pos: 1-based genomic position

            Returns:
                Genotype enum value

            Raises:
                ValueError: If sample not found or position has no variant
            """
            # Validate sample exists
            if sample not in self._samples:
                raise ValueError(f"Sample '{sample}' not found in VCF")

            sample_idx = self._samples.index(sample)

            # Query the position using cyvcf2 (requires indexed file)
            try:
                # cyvcf2 uses 1-based coordinates for queries
                region = f"{chrom}:{pos}-{pos}"
                records = list(self.vcf(region))
            except Exception as e:
                raise ValueError(f"Failed to query position {chrom}:{pos}: {e}") from e

            if not records:
                raise ValueError(f"No variant found at {chrom}:{pos}")

            # Get genotype from first matching record
            variant = records[0]
            gt_type = variant.gt_types[sample_idx]

            # Convert to Genotype enum
            if gt_type == 0:
                return Genotype.HOM_REF
            elif gt_type == 1:
                return Genotype.HET
            elif gt_type == 3:
                return Genotype.HOM_ALT
            else:
                return Genotype.MISSING

        def query_region(
            self, chrom: str, start: int, end: int, samples: list[str] | None = None
        ) -> Iterator[VariantGenotype]:
            """Query variants in a genomic region.

            Requires the VCF to be indexed (.tbi or .csi). Uses 1-based inclusive
            coordinates (VCF standard).

            Args:
                chrom: Chromosome name
                start: 1-based start position (inclusive)
                end: 1-based end position (inclusive)
                samples: Optional list of sample IDs. If None, uses first sample.

            Yields:
                VariantGenotype objects in the region

            Raises:
                ValueError: If the file is not indexed or region is invalid
            """
            # Determine target sample
            if samples is None:
                target_samples = [self._samples[0]] if self._samples else []
            else:
                for s in samples:
                    if s not in self._samples:
                        raise ValueError(f"Sample '{s}' not found in VCF")
                target_samples = samples

            if not target_samples:
                return

            sample_id = target_samples[0]
            sample_idx = self._samples.index(sample_id)

            # Query region (cyvcf2 uses 1-based coordinates)
            try:
                region = f"{chrom}:{start}-{end}"
                records = self.vcf(region)
            except Exception as e:
                raise ValueError(
                    f"Failed to query region {chrom}:{start}-{end}. File may not be indexed: {e}"
                ) from e

            # Yield VariantGenotype for each record
            for variant in records:
                gt_type = variant.gt_types[sample_idx]

                # Convert to Genotype enum
                if gt_type == 0:
                    genotype = Genotype.HOM_REF
                elif gt_type == 1:
                    genotype = Genotype.HET
                elif gt_type == 3:
                    genotype = Genotype.HOM_ALT
                else:
                    genotype = Genotype.MISSING

                # Create Variant (use first ALT)
                alt = variant.ALT[0] if variant.ALT else variant.REF
                var = Variant(
                    chrom=variant.CHROM,
                    pos=variant.POS,
                    ref=variant.REF,
                    alt=alt,
                    id=variant.ID if variant.ID else None,
                )

                # Get allele sequences
                gt_bases = variant.gt_bases[sample_idx]
                if gt_bases and "/" in gt_bases:
                    alleles = gt_bases.split("/")
                    allele1 = alleles[0] if alleles[0] != "." else None
                    allele2 = alleles[1] if len(alleles) > 1 and alleles[1] != "." else None
                elif gt_bases and "|" in gt_bases:
                    alleles = gt_bases.split("|")
                    allele1 = alleles[0] if alleles[0] != "." else None
                    allele2 = alleles[1] if len(alleles) > 1 and alleles[1] != "." else None
                else:
                    allele1, allele2 = None, None

                yield VariantGenotype(
                    variant=var, genotype=genotype, allele1=allele1, allele2=allele2
                )

        def to_bed(
            self,
            output: Path,
            samples: list[str] | None = None,
            het_only: bool = True,
            include_genotypes: bool = True,
        ) -> Path:
            r"""Export variants to BED format file.

            Uses bcftools for efficient filtering and export. BED format uses
            0-based start, 1-based end coordinates.

            Format:
            - Without genotypes: chrom\\tstart\\tend\\tref\\talt
            - With genotypes: chrom\\tstart\\tend\\tref\\talt\\tgenotype

            Args:
                output: Output BED file path
                samples: Optional list of sample IDs to include
                het_only: If True, only export heterozygous variants
                include_genotypes: If True, include genotype column(s)

            Returns:
                Path to the created BED file

            Raises:
                IOError: If bcftools fails or file cannot be written
                ValueError: If samples not found
            """
            # Validate samples if provided
            if samples is not None:
                for s in samples:
                    if s not in self._samples:
                        raise ValueError(f"Sample '{s}' not found in VCF")

            # Build bcftools commands based on parameters
            # This follows the pattern from VCFSource for consistency

            # Base view command: filter to biallelic SNPs
            view_cmd = [
                "bcftools",
                "view",
                str(self.path),
                "-m2",
                "-M2",  # min/max alleles
                "-v",
                "snps",  # SNPs only
                "-Ou",  # uncompressed BCF output
            ]

            # Build query command
            query_cmd = ["bcftools", "query", "-o", str(output), "-f"]

            # Configure based on samples and het_only
            if samples is None:
                # No samples: drop genotypes
                view_cmd.append("--drop-genotypes")
                query_cmd.append("%CHROM\t%POS0\t%END\t%REF\t%ALT\n")

                view_process = subprocess.run(view_cmd, stdout=subprocess.PIPE, check=True)
            else:
                samples_arg = ",".join(samples)
                num_samples = len(samples)

                if num_samples > 1:
                    # Multi-sample: filter to variants with at least one non-ref allele
                    view_cmd.extend(
                        ["-s", samples_arg, "--min-ac", "1", "--max-ac", str((num_samples * 2) - 1)]
                    )
                    view_process = subprocess.run(view_cmd, stdout=subprocess.PIPE, check=True)
                else:
                    # Single sample
                    view_cmd.extend(["-s", samples_arg])
                    subset_process = subprocess.run(view_cmd, stdout=subprocess.PIPE, check=True)

                    if het_only:
                        # Filter to het genotypes
                        het_view_cmd = ["bcftools", "view", "--genotype", "het", "-Ou"]
                        view_process = subprocess.run(
                            het_view_cmd,
                            input=subset_process.stdout,
                            stdout=subprocess.PIPE,
                            check=True,
                        )
                    else:
                        view_process = subset_process

                # Add genotype column if requested
                if include_genotypes:
                    query_cmd.append("%CHROM\t%POS0\t%END\t%REF\t%ALT[\t%TGT]\n")
                else:
                    query_cmd.append("%CHROM\t%POS0\t%END\t%REF\t%ALT\n")

            # Run query command
            try:
                subprocess.run(query_cmd, input=view_process.stdout, check=True)
            except subprocess.CalledProcessError as e:
                raise OSError(f"bcftools failed: {e}") from e

            return output

        def close(self) -> None:
            """Close the cyvcf2.VCF handle.

            Releases file resources. Should be called when done with the source,
            or use context manager protocol.
            """
            if hasattr(self, "vcf") and self.vcf is not None:
                self.vcf.close()
else:
    # Create dummy class if cyvcf2 not available (for documentation/type checking)
    class CyVCF2Source:  # type: ignore[no-redef]
        """Placeholder class when cyvcf2 is not installed."""

        def __init__(self, *args, **kwargs):
            raise ImportError("cyvcf2 is not installed. Install with: pip install wasp2[cyvcf2]")
