import logging
import os
import os.path
from os.path import dirname
from os.path import join as pjoin

import pysam
from simppl.simple_pipeline import SimplePipeline
from ugbio_core.exec_utils import print_and_execute


class VcfUtils:
    """Utilities of vcf pipeline, mostly wrappers around shell scripts

    Attributes
    ----------
    sp : SimplePipeline
        Simple pipeline object
    logger : logging.Logger
        Logger instance
    """

    def __init__(self, simple_pipeline: SimplePipeline | None = None, logger: logging.Logger | None = None):
        """Combines VCF in parts from GATK and indices the result

        Parameters
        ----------
        simple_pipeline : SimplePipeline, optional
            Optional SimplePipeline object for executing shell commands
        logger : logging.Logger, optional
            Optional logger instance
        """
        self.sp = simple_pipeline
        if logger is None:
            self.logger = logging.getLogger(__name__)
        else:
            self.logger = logger

    def __execute(self, command: str, output_file: str | None = None):
        """Summary

        Parameters
        ----------
        command : str
            Description
        output_file : str, optional
            Description
        """
        print_and_execute(command, output_file=output_file, simple_pipeline=self.sp, module_name=__name__)

    def index_vcf(self, vcf: str):
        """Tabix index on VCF

        Parameters
        ----------
        vcf : str
            Input vcf.gz file
        """
        self.__execute(f"bcftools index -tf {vcf}")

    def sort_vcf(self, input_file: str, output_file: str):
        """Sort VCF file

        Parameters
        ----------
        input_file : str
            Input file name
        output_file : str
            Output file name

        Returns
        -------
        None
            Generates `output_file`.
        """
        tempdir = os.path.dirname(os.path.abspath(output_file))
        self.__execute(f"bcftools sort -o {output_file} -O z {input_file} -T {tempdir}/")

    def concat_vcf(self, input_files: list[str], output_file: str):
        """Concatenate VCF files using bcftools concat and index

        Parameters
        ----------
        input_files : list[str]
            List of input VCF files
        output_file : str
            Output VCF file

        Returns
        -------
        None
            Generates `output_file`.
        """
        input_files_str = " ".join(input_files)
        tmpdir = os.path.dirname(output_file)
        out_tmp_file = f"{os.path.join(tmpdir, os.path.basename(output_file))}.tmp.vcf.gz"
        self.__execute(f"bcftools concat -a -o {out_tmp_file} -O z {input_files_str}")
        self.sort_vcf(out_tmp_file, output_file)
        self.index_vcf(output_file)
        os.unlink(out_tmp_file)

    def filter_vcf(
        self,
        input_vcf: str,
        output_vcf: str,
        filter_name: str,
        n_threads: int = 1,
        include_expression: str | None = None,
        exclude_expression: str | None = None,
    ) -> None:
        """Filter VCF file using bcftools filter and add a filter with a given name

        Parameters
        ----------
        input_vcf : str
            Input VCF file path
        output_vcf : str
            Output VCF file path
        filter_name : str
            Name of the filter to add to variants that don't pass the criteria
        n_threads : int, optional
            Number of threads to use (default is 1)
        include_expression : str, optional
            bcftools include expression (variants NOT matching this will be filtered)
        exclude_expression : str, optional
            bcftools exclude expression (variants matching this will be filtered)

        Raises
        ------
        ValueError
            If neither include_expression nor exclude_expression is provided, or if both are provided
        """
        if include_expression is None and exclude_expression is None:
            msg = "At least one of include_expression or exclude_expression must be provided"
            raise ValueError(msg)

        if include_expression is not None and exclude_expression is not None:
            msg = "Only one of include_expression or exclude_expression can be provided at a time"
            raise ValueError(msg)

        # Build the bcftools filter command
        cmd_parts = ["bcftools", "filter"]

        # Add include or exclude expression
        if include_expression is not None:
            cmd_parts.extend(["-i", f"'{include_expression}'"])
        else:  # exclude_expression is not None
            cmd_parts.extend(["-e", f"'{exclude_expression}'"])

        cmd_parts.extend(["--threads", str(n_threads)])
        # Add filter name and output format
        cmd_parts.extend(["-s", filter_name, "-m", "+", "-O", "z", "-o", output_vcf, input_vcf])

        # Execute the filtering command
        self.__execute(" ".join(cmd_parts))

    def intersect_with_intervals(self, input_fn: str, intervals_fn: str, output_fn: str) -> None:
        """Intersects VCF with intervalList. Writes output_fn file

        Parameters
        ----------
        input_fn : str
            Input file
        intervals_fn : str
            Interval_list file
        output_fn : str
            Output file

        Writes output_fn file
        """
        self.__execute(f"gatk SelectVariants -V {input_fn} -L {intervals_fn} -O {output_fn}")

    def view_vcf(
        self,
        input_vcf: str,
        output_vcf: str,
        n_threads: int = 1,
        extra_args: str = "",
    ) -> None:
        """View/subset VCF file using bcftools view

        Parameters
        ----------
        input_vcf : str
            Input VCF file path
        output_vcf : str
            Output VCF file path
        n_threads : int, optional
            Number of threads to use (default is 1)
        extra_args : str, optional
            Additional arguments to pass to bcftools view (default is empty string)
        """
        # Build the bcftools view command
        cmd_parts = ["bcftools", "view"]

        # Add threads
        cmd_parts.extend(["--threads", str(n_threads)])

        # Add extra arguments if provided
        if extra_args:
            cmd_parts.extend(extra_args.split())

        # Add output format and files
        cmd_parts.extend(["-O", "z", "-o", output_vcf, input_vcf])

        # Execute the view command
        self.__execute(" ".join(cmd_parts))

    def remove_filter_annotations(self, input_vcf: str, output_vcf: str, n_threads: int = 1) -> None:
        """
        Remove all filter annotations from a VCF file using bcftools.

        This function removes:
        1. All ##FILTER header lines from the VCF header
        2. Sets all FILTER column values to '.' (missing) in variant records

        Parameters
        ----------
        input_vcf : str
            Path to the input VCF file (can be .vcf or .vcf.gz)
        output_vcf : str
            Path to the output VCF file where filtered result will be written
        n_threads : int, optional
            Number of threads to use for bcftools operations (default: 1)

        Returns
        -------
        None
            Writes the filtered VCF to output_vcf and creates an index
        """
        # Build the bcftools annotate command
        cmd_parts = ["bcftools", "annotate"]

        # Remove FILTER column data and FILTER header lines
        cmd_parts.extend(["-x", "FILTER"])

        # Add threading
        cmd_parts.extend(["--threads", str(n_threads)])

        # Add output format and files
        cmd_parts.extend(["-o", output_vcf, "-O", "z", input_vcf])

        # Execute the command
        self.__execute(" ".join(cmd_parts))

        # Index the output VCF
        self.index_vcf(output_vcf)

    def collapse_vcf(  # noqa: PLR0913
        self,
        vcf: str,
        output_vcf: str,
        bed: str | None = None,
        refdist: int = 500,
        pctseq: float = 0.0,
        pctsize: float = 0.0,
        maxsize: int = 50000,
        *,
        ignore_filter: bool = False,
        ignore_sv_type: bool = True,
        pick_best: bool = False,
        erase_removed: bool = True,
    ) -> str | None:
        """
        Collapse SV/CNV VCF using truvari collapse

        Parameters
        ----------
        vcf : str
            Input VCF file
        output_vcf : str
            Output VCF file
        bed : str, optional
            Restrict collapsing to this bed file, by default None
        refdist : int, optional
            Reference distance for collapsing variants, by default 500
        pctseq : float, optional
            Percentage of sequence identity, by default 0.0
        pctsize : float, optional
            Percentage of size identity, by default 0.0
        maxsize : int, optional
            Maximum size for SV comparison, by default 50000. For CNV
            comparison, consider increasing this value or use -1 for
            unlimited.
        ignore_filter : bool, optional
            If True, ignore FILTER field (remove truvari's --passonly flag), by default False
        ignore_sv_type : bool, optional
            If True, ignore SVTYPE when collapsing variants, by default True
        pick_best : bool, optional
            If True, pick the best variant among those being merged, by default False, pick first
        erase_removed: bool, optional
            If True, delete the temporary file with removed variants, by default True,
            if not return the location of the file

        Returns
        -------
        None | str
        """

        removed_vcf_path = pjoin(dirname(output_vcf), "tmp.vcf")
        truvari_cmd = ["truvari", "collapse", "-i", vcf, "-c", removed_vcf_path, "--chain"]

        if not ignore_filter:
            truvari_cmd.append("--passonly")
        if ignore_sv_type:
            truvari_cmd.append("-t")
        if pick_best:
            truvari_cmd.extend(["--keep", "maxqual"])
        else:
            truvari_cmd.extend(["--keep", "first"])
        if bed:
            truvari_cmd.extend(["--bed", bed])
        truvari_cmd.extend(["--pctseq", str(pctseq)])
        truvari_cmd.extend(["--pctsize", str(pctsize)])
        truvari_cmd.extend(["--refdist", str(int(refdist))])
        truvari_cmd.extend(["--sizemax", str(maxsize)])
        truvari_cmd.extend(["-o", output_vcf])
        self.logger.info(f"truvari command: {' '.join(truvari_cmd)}")
        self.__execute(" ".join(truvari_cmd))

        # Parameterize the file path
        if erase_removed:
            os.unlink(removed_vcf_path)
            self.logger.info(f"Deleted temporary file: {removed_vcf_path}")
        else:
            return removed_vcf_path

    def annotate_vcf(
        self,
        input_vcf: str,
        output_vcf: str,
        annotation_file: str,
        header_file: str | None = None,
        columns: str | None = None,
        n_threads: int = 1,
        extra_args: str = "",
    ) -> None:
        """
        Annotate VCF file using bcftools annotate

        Parameters
        ----------
        input_vcf : str
            Input VCF file path
        output_vcf : str
            Output VCF file path
        annotation_file : str
            Path to the annotation file (e.g., tabix-indexed TSV file)
        header_file : str, optional
            Path to header (txt) file containing (only) new INFO/FORMAT field definitions
        columns : str, optional
            Comma-separated list of columns to annotate (e.g., CHROM,POS,INFO/DP)
        n_threads : int, optional
            Number of threads to use (default is 1)
        extra_args : str, optional
            Additional arguments to pass to bcftools annotate (default is empty string)

        Usage example:
        >>> vcf_utils = VcfUtils()
        >>> vcf_utils.annotate_vcf(
        ...     input_vcf="input.vcf.gz",
        ...     output_vcf="output.vcf.gz",
        ...     annotation_file="annotations.tsv.gz",
        ...     header_file="new_info_header.txt",
        ...     columns="CHROM,POS,INFO/NEW_FIELD",
        ...     n_threads=4,
        ...     extra_args="--some-extra-arg"
        ... )
        """
        # Build the bcftools annotate command
        cmd_parts = ["bcftools", "annotate"]

        # Add output format
        cmd_parts.extend(["-O", "z", "-o", output_vcf])

        # Add annotation file
        cmd_parts.extend(["-a", annotation_file])

        # Add header file if provided
        if header_file:
            cmd_parts.extend(["-h", header_file])

        # Add columns if provided
        if columns:
            cmd_parts.extend(["-c", columns])

        # Add threading
        cmd_parts.extend(["--threads", str(n_threads)])

        # Add extra arguments if provided
        if extra_args:
            cmd_parts.extend(extra_args.split())

        # Add input file
        cmd_parts.append(input_vcf)

        # Execute the command
        self.__execute(" ".join(cmd_parts))

        # Index the output VCF
        self.index_vcf(output_vcf)

    def annotate_tandem_repeats(self, input_file: str, reference_fasta: str) -> str:
        """Runs VariantAnnotator on the input file to add tandem repeat annotations (maybe others)

        Parameters
        ----------
        input_file : str
            vcf.gz file
        reference_fasta : str
            Reference file (should have .dict file nearby)

        Creates a copy of the input_file with .annotated.vcf.gz and the index
        Returns
        -------
        path to output file: str
        """

        output_file = input_file.replace("vcf.gz", "annotated.vcf.gz")
        self.__execute(f"gatk VariantAnnotator -V {input_file} -O {output_file} -R {reference_fasta} -A TandemRepeat")
        return output_file

    def update_vcf_contigs_from_fai(
        self,
        input_vcf: str,
        output_vcf: str,
        fasta_fai: str,
    ) -> None:
        """
        Update VCF header contig information using FASTA FAI file and index the output.

        This function updates the VCF header contig lines using information from a FASTA
        index (.fai) file. The contig lines are replaced with entries from the FAI file.

        Parameters
        ----------
        input_vcf : str
            Input VCF file path (can be .vcf or .vcf.gz)
        output_vcf : str
            Output VCF file path (.vcf.gz recommended)
        fasta_fai : str
            Path to the FASTA index (.fai) file

        Returns
        -------
        None
            Writes the updated VCF to output_vcf and creates an index

        """
        # Use bcftools reheader to update contig lines from FAI file
        self.__execute(f"bcftools reheader -f {fasta_fai} -o {output_vcf} {input_vcf}")

        # Index the output VCF
        self.index_vcf(output_vcf)

    def remove_filters(self, input_vcf: str, output_vcf: str, filters_to_remove: list[str] | None = None) -> None:
        """Remove specific filters or all filters from VCF records using bcftools annotate.

        Parameters
        ----------
        input_vcf : str
            Path to input VCF file
        output_vcf : str
            Path to output VCF file
        filters_to_remove : list[str], optional
            List of specific filter names to remove (e.g., ['LowQual', 'LowDP']), if None - all are removed
        """

        # Build annotation removal string
        if filters_to_remove is None:
            annotation = "FILTER"
        else:
            annotation = f"FILTER/{','.join(filters_to_remove)}"
        self.__execute(f"bcftools annotate -x {annotation} -o {output_vcf} {input_vcf}")

    @staticmethod
    def copy_vcf_record(rec: pysam.VariantRecord, new_header: pysam.VariantHeader) -> pysam.VariantRecord:
        """
        Create a new VCF record with the same data as the input record, but using a new header.

        Parameters
        ----------
        rec : pysam.VariantRecord
            The original VCF record to copy.
        new_header : pysam.VariantHeader
            The new VCF header to use for the copied record.

        Returns
        -------
        pysam.VariantRecord
            A new VCF record with the same data as `rec`, but using `new_header`.
        """
        new_record = new_header.new_record(
            contig=rec.chrom,
            start=rec.start,
            stop=rec.stop,
            id=rec.id,
            qual=rec.qual,
            alleles=rec.alleles,
            filter=rec.filter.keys(),
        )

        # copy INFO fields
        for k, v in rec.info.items():
            if k in new_header.info:
                new_record.info[k] = v

        # copy FORMAT fields
        for sample in rec.samples:
            src = rec.samples[sample]
            tgt = new_record.samples[sample]
            for k, v in src.items():
                if v in (None, (None,)):
                    continue  # no need to assign missing values
                tgt[k] = v

        return new_record
