import logging
import os
import os.path
from dataclasses import dataclass, field
from enum import Enum
from os.path import dirname
from os.path import join as pjoin
from pathlib import Path

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


class VcfMetaType(str, Enum):
    """VCF header meta types for structured field definitions."""

    INFO = "INFO"
    FORMAT = "FORMAT"


@dataclass(frozen=True)
class VcfField:
    """VCF field definition for INFO or FORMAT header lines.

    Parameters
    ----------
    field_id : str
        Field identifier (e.g. ``"DP"``, ``"GT"``).
    number : str
        VCF Number attribute (e.g. ``"1"``, ``"A"``, ``"."``, ``"G"``).
    field_type : str
        VCF Type attribute (e.g. ``"Integer"``, ``"Float"``, ``"String"``).
    description : str
        Human-readable description.
    meta_type : VcfMetaType
        Whether this is an INFO or FORMAT field. Defaults to ``VcfMetaType.INFO``.
    """

    field_id: str
    number: str
    field_type: str
    description: str
    meta_type: VcfMetaType = field(default=VcfMetaType.INFO)

    def to_header_line(self) -> str:
        """Render as a VCF header line string.

        Returns
        -------
        str
            A complete VCF header line, e.g.
            ``##INFO=<ID=DP,Number=1,Type=Integer,Description="Read depth">``.
        """
        return (
            f"##{self.meta_type.value}=<ID={self.field_id},Number={self.number},"
            f'Type={self.field_type},Description="{self.description}">'
        )


def write_vcf_header_file(
    fields: list[VcfField], header_file: Path, additional_header_lines: list[str] | None = None
) -> None:
    """Write a VCF header file with field definitions and optional additional header lines.

    Creates a header file suitable for use with bcftools annotate -h option.
    Each field's ``meta_type`` determines whether it is rendered as ``##INFO=`` or ``##FORMAT=``.

    Parameters
    ----------
    fields : list[VcfField]
        List of VcfField objects defining the fields to add.
    header_file : Path
        Path to the output header file.
    additional_header_lines : list[str], optional
        List of additional header lines to append (e.g., ["##tumor_sample=<sample_name>"]).
        These lines will be written as-is after the field definitions.
    """
    with open(header_file, "w") as f:
        for fld in fields:
            f.write(fld.to_header_line() + "\n")

        if additional_header_lines:
            for line in additional_header_lines:
                formatted_line = line if line.startswith("##") else f"##{line}"
                f.write(formatted_line + "\n")


def get_vcf_sample_names(vcf_path: str | Path) -> list[str]:
    """Get the list of sample names from a VCF file header.

    Parameters
    ----------
    vcf_path : str | Path
        Path to the VCF file.

    Returns
    -------
    list[str]
        List of sample names found in the VCF header.
    """
    with pysam.VariantFile(str(vcf_path)) as vcf:
        return list(vcf.header.samples)
