"""Generate allele-swapped reads for remapping.

Provides functions for creating FASTQ files with haplotype-swapped reads
that need to be remapped to check for mapping bias.
"""

from __future__ import annotations

import shutil
import tempfile
from pathlib import Path

import pysam

from wasp2.cli import (
    create_progress,
    detail,
    error,
    print_file_path,
    rust_status,
    status,
    success,
)

# Rust acceleration (required; no fallback)
from wasp2_rust import remap_all_chromosomes, remap_chromosome, remap_chromosome_multi


def _write_remap_bam_rust_optimized(
    bam_file: str,
    intersect_file: str,
    r1_out: str,
    r2_out: str,
    max_seqs: int = 64,
    parallel: bool = True,
) -> None:
    """
    Optimized Rust remapping - parses intersect file ONCE, processes chromosomes in parallel.

    This is the fastest implementation:
    - Parses intersect file once (22x fewer parse operations for RNA-seq)
    - Uses rayon for parallel chromosome processing (4-8x speedup with 8 cores)
    - Total expected speedup: ~100x for large RNA-seq datasets
    """
    import inspect

    mode = "parallel" if parallel else "sequential"
    rust_status(f"Using optimized Rust remapper (parse-once, {mode})")

    # Check if the Rust function accepts 'parallel' parameter (backward compatibility)
    sig = inspect.signature(remap_all_chromosomes)
    has_parallel_param = "parallel" in sig.parameters

    if has_parallel_param:
        # New version with parallel parameter
        pairs, haps = remap_all_chromosomes(
            bam_file, intersect_file, r1_out, r2_out, max_seqs=max_seqs, parallel=parallel
        )
    else:
        # Old version without parallel parameter (always runs in parallel)
        detail("Using Rust version without 'parallel' parameter (parallel by default)")
        pairs, haps = remap_all_chromosomes(
            bam_file, intersect_file, r1_out, r2_out, max_seqs=max_seqs
        )

    success(f"Rust remapper (optimized): {pairs:,} pairs → {haps:,} haplotypes")
    print_file_path("R1 output", r1_out)
    print_file_path("R2 output", r2_out)


def _write_remap_bam_rust(
    bam_file: str, intersect_file: str, r1_out: str, r2_out: str, max_seqs: int = 64
) -> None:
    """Rust-accelerated remapping implementation (5-7x faster than Python) - LEGACY per-chromosome version"""
    # Get chromosomes that have variants in the intersect file
    # This avoids processing ~170 empty chromosomes (major speedup!)
    intersect_chroms = set()
    with open(intersect_file) as f:
        for line in f:
            chrom = line.split("\t")[0]
            intersect_chroms.add(chrom)

    # Filter BAM chromosomes to only those with variants
    with pysam.AlignmentFile(bam_file, "rb") as bam:
        chromosomes = [c for c in bam.header.references if c in intersect_chroms]

    status(f"Processing {len(chromosomes)} chromosomes with variants")

    # Create temp directory for per-chromosome outputs
    with tempfile.TemporaryDirectory() as tmpdir:
        total_pairs = 0
        total_haps = 0

        with create_progress() as progress:
            task = progress.add_task("Remapping chromosomes", total=len(chromosomes))

            # Process each chromosome with Rust
            for chrom in chromosomes:
                chrom_r1 = f"{tmpdir}/{chrom}_r1.fq"
                chrom_r2 = f"{tmpdir}/{chrom}_r2.fq"

                try:
                    pairs, haps = remap_chromosome(
                        bam_file, intersect_file, chrom, chrom_r1, chrom_r2, max_seqs=max_seqs
                    )
                    total_pairs += pairs
                    total_haps += haps
                    if pairs > 0:
                        detail(f"{chrom}: {pairs:,} pairs → {haps:,} haplotypes")
                except (RuntimeError, OSError) as e:
                    error(f"{chrom}: Error - {e}")

                progress.update(task, advance=1)

        # Concatenate all R1 files
        r1_files = sorted(Path(tmpdir).glob("*_r1.fq"))
        with open(r1_out, "wb") as outfile:
            for fq_path in r1_files:
                with open(fq_path, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

        # Concatenate all R2 files
        r2_files = sorted(Path(tmpdir).glob("*_r2.fq"))
        with open(r2_out, "wb") as outfile:
            for fq_path in r2_files:
                with open(fq_path, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

        success(f"Rust remapper: {total_pairs:,} pairs → {total_haps:,} haplotypes")
        print_file_path("R1 output", r1_out)
        print_file_path("R2 output", r2_out)


def _write_remap_bam_rust_multi(
    bam_file: str,
    intersect_file: str,
    r1_out: str,
    r2_out: str,
    num_samples: int,
    max_seqs: int = 64,
) -> None:
    """Rust-accelerated multi-sample remapping implementation"""
    # Get chromosomes that have variants in the intersect file
    intersect_chroms = set()
    with open(intersect_file) as f:
        for line in f:
            chrom = line.split("\t")[0]
            intersect_chroms.add(chrom)

    # Filter BAM chromosomes to only those with variants
    with pysam.AlignmentFile(bam_file, "rb") as bam:
        chromosomes = [c for c in bam.header.references if c in intersect_chroms]

    status(f"Processing {len(chromosomes)} chromosomes with variants ({num_samples} samples)")

    # Create temp directory for per-chromosome outputs
    with tempfile.TemporaryDirectory() as tmpdir:
        total_pairs = 0
        total_haps = 0

        with create_progress() as progress:
            task = progress.add_task("Multi-sample remapping", total=len(chromosomes))

            # Process each chromosome with Rust multi-sample
            for chrom in chromosomes:
                chrom_r1 = f"{tmpdir}/{chrom}_r1.fq"
                chrom_r2 = f"{tmpdir}/{chrom}_r2.fq"

                try:
                    pairs, haps = remap_chromosome_multi(
                        bam_file,
                        intersect_file,
                        chrom,
                        chrom_r1,
                        chrom_r2,
                        num_samples=num_samples,
                        max_seqs=max_seqs,
                    )
                    total_pairs += pairs
                    total_haps += haps
                    if pairs > 0:
                        detail(f"{chrom}: {pairs:,} pairs → {haps:,} haplotypes")
                except (RuntimeError, OSError) as e:
                    error(f"{chrom}: Error - {e}")

                progress.update(task, advance=1)

        # Concatenate all R1 files
        r1_files = sorted(Path(tmpdir).glob("*_r1.fq"))
        with open(r1_out, "wb") as outfile:
            for fq_path in r1_files:
                with open(fq_path, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

        # Concatenate all R2 files
        r2_files = sorted(Path(tmpdir).glob("*_r2.fq"))
        with open(r2_out, "wb") as outfile:
            for fq_path in r2_files:
                with open(fq_path, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

        success(f"Rust multi-sample remapper: {total_pairs:,} pairs → {total_haps:,} haplotypes")
        print_file_path("R1 output", r1_out)
        print_file_path("R2 output", r2_out)


def write_remap_bam(
    bam_file: str,
    intersect_file: str,
    r1_out: str,
    r2_out: str,
    samples: list[str],
    max_seqs: int = 64,
    include_indels: bool = False,
    insert_qual: int = 30,
) -> None:
    """Rust-accelerated remapping - parses intersect file once, processes chromosomes in parallel.

    Uses Rust acceleration (required; no fallback).

    Args:
        bam_file: Input BAM file
        intersect_file: Intersect BED file
        r1_out: Output FASTQ for read 1
        r2_out: Output FASTQ for read 2
        samples: List of sample IDs
        max_seqs: Maximum haplotype sequences per read pair
        include_indels: Include indels in remapping (not yet supported in Rust)
        insert_qual: Quality score for inserted bases (not yet supported in Rust)
    """
    num_samples = len(samples)

    if num_samples == 1:
        # Single sample: use optimized all-chromosome Rust
        _write_remap_bam_rust_optimized(
            bam_file, intersect_file, r1_out, r2_out, max_seqs, parallel=True
        )
    else:
        # Multi-sample: use per-chromosome Rust
        _write_remap_bam_rust_multi(bam_file, intersect_file, r1_out, r2_out, num_samples, max_seqs)
