"""Gene annotation parsing and data management."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import NamedTuple

import polars as pl

logger = logging.getLogger(__name__)


class ParsedGeneData(NamedTuple):
    """Parsed gene data from GTF/GFF3 file."""

    gene_df: pl.DataFrame
    feature: str
    attribute: str
    parent_attribute: str


class WaspGeneData:
    """Container for gene annotation file paths and configuration.

    Attributes
    ----------
    gene_file : str | Path
        Path to the gene annotation file.
    feature : str | None
        Feature type to extract.
    attribute : str | None
        Attribute for region names.
    parent_attribute : str | None
        Parent attribute for hierarchical features.
    """

    def __init__(
        self,
        gene_file: str | Path,
        feature: str | None = None,
        attribute: str | None = None,
        parent_attribute: str | None = None,
    ) -> None:
        self.gene_file = gene_file
        self.feature = feature
        self.attribute = attribute
        self.parent_attribute = parent_attribute

    def update_data(self, data: ParsedGeneData) -> None:
        """Update attributes with parsed data.

        Parameters
        ----------
        data : ParsedGeneData
            Parsed gene data to update from.
        """
        for key in data._fields:
            if hasattr(self, key):
                setattr(self, key, getattr(data, key))


def parse_gene_file(
    gene_file: str | Path,
    feature: str | None = None,
    attribute: str | None = None,
    parent_attribute: str | None = None,
) -> ParsedGeneData:
    """Parse GTF/GFF3 gene annotation file.

    Parameters
    ----------
    gene_file : str | Path
        Path to GTF/GFF3 file.
    feature : str | None, optional
        Feature type to extract (auto-detected if None).
    attribute : str | None, optional
        Attribute for region names (auto-detected if None).
    parent_attribute : str | None, optional
        Parent attribute for hierarchical features (auto-detected if None).

    Returns
    -------
    ParsedGeneData
        Named tuple with (gene_df, feature, attribute, parent_attribute).
    """
    # Use gtf col names
    gtf_cols = [
        "seqname",
        "source",
        "feature",
        "start",
        "end",
        "score",
        "strand",
        "frame",
        "attribute",
    ]

    # Cant use lazyframe in case of compressed
    df = pl.read_csv(
        gene_file, separator="\t", comment_prefix="#", has_header=False, new_columns=gtf_cols
    )

    # Check if we need to preparse feature
    if feature is None:
        feature_list = df.select(pl.col("feature").unique()).to_series()

        if "exon" in feature_list:
            feature = "exon"
        elif "transcript" in feature_list:
            feature = "transcript"
        elif "gene" in feature_list:
            feature = "gene"
        else:
            logger.warning("exon, gene or transcript not found in feature list: %s", feature_list)

    assert feature is not None, "Could not determine feature type from gene file"

    # feature filter
    df = df.filter(pl.col("feature") == feature)

    # Parse attributes
    if attribute is None:
        if df.get_column("attribute").str.contains(f"{feature}_id").any():
            attribute = f"{feature}_id"
        elif df.get_column("attribute").str.contains("ID").any():
            attribute = "ID"
        elif df.get_column("attribute").str.contains("Name").any():
            attribute = "Name"
        else:
            # TODO return an error
            # TODO maybe just use region or coords as a feature
            logger.warning("No 'ID', '%s_id' or 'Name' attribute found. Please include ID", feature)

    # TODO: Figure out best way to handle parent attribute

    # Parse parent attributes
    # Figure out best way to parse and handle this
    if parent_attribute is None:
        # Defaults to gene(possibly transcript???)

        if df.get_column("attribute").str.contains("Parent").any():
            parent_attribute = "Parent"
        elif df.get_column("attribute").str.contains("transcript_id").any():
            parent_attribute = "transcript_id"
        elif df.get_column("attribute").str.contains("gene_id").any():
            parent_attribute = "gene_id"
        else:
            parent_attribute = attribute

    # TODO: Allow for count output without parent column
    assert attribute is not None, "Could not determine attribute from gene file"
    assert parent_attribute is not None, "Could not determine parent_attribute from gene file"
    if parent_attribute == attribute:
        parent_col = f"groupby_{attribute}"
    else:
        parent_col = parent_attribute

    # Extract relevant attributes
    attr_regex = rf"{attribute}[=\s]\"?\'?(.*?)\"?\'?;"
    parent_regex = rf"{parent_attribute}[=\s]\"?\'?(.*?)\"?\'?;"

    df = df.with_columns(
        pl.col("start").sub(1),
        pl.col("attribute").str.extract(attr_regex).alias(attribute),
        pl.col("attribute").str.extract(parent_regex).alias(parent_col),
    ).select(["seqname", "start", "end", attribute, parent_col])

    return ParsedGeneData(df, feature, attribute, parent_col)


def make_gene_data(
    gene_file: str | Path,
    out_bed: str | Path,
    feature: str | None = None,
    attribute: str | None = None,
    parent_attribute: str | None = None,
) -> WaspGeneData:
    """Parse gene file and create BED for intersection.

    Parameters
    ----------
    gene_file : str | Path
        Path to GTF/GFF3 file.
    out_bed : str | Path
        Output BED file path.
    feature : str | None, optional
        Feature type to extract.
    attribute : str | None, optional
        Attribute for region names.
    parent_attribute : str | None, optional
        Parent attribute for hierarchical features.

    Returns
    -------
    WaspGeneData
        Container with parsed gene data and configuration.
    """
    gene_data = WaspGeneData(
        gene_file=gene_file, feature=feature, attribute=attribute, parent_attribute=parent_attribute
    )

    parsed_file_data = parse_gene_file(
        gene_data.gene_file,
        feature=gene_data.feature,
        attribute=gene_data.attribute,
        parent_attribute=gene_data.parent_attribute,
    )

    # Update gene_data obj
    gene_data.update_data(parsed_file_data)

    # Write out_bed
    parsed_file_data.gene_df.write_csv(out_bed, separator="\t", include_header=False)

    return gene_data


def parse_intersect_genes(
    intersect_file: str | Path,
    attribute: str | None = None,
    parent_attribute: str | None = None,
) -> pl.DataFrame:
    """Parse gene intersection file (legacy version).

    Parameters
    ----------
    intersect_file : str | Path
        Path to bedtools intersection output.
    attribute : str | None, optional
        Attribute column name, by default 'ID'.
    parent_attribute : str | None, optional
        Parent attribute column name, by default 'Parent'.

    Returns
    -------
    pl.DataFrame
        Parsed intersection data.
    """
    if attribute is None:
        attribute = "ID"

    if parent_attribute is None:
        parent_attribute = "Parent"

    # AFTER performing gtf_to_bed and intersecting!
    df = pl.scan_csv(intersect_file, separator="\t", has_header=False, infer_schema_length=0)

    # Should i poossibly consider diff number of cols?

    # Might want to do checks for wrong number of columns
    subset_cols = [df.columns[i] for i in [0, 2, 3, 4, -2, -1]]
    new_cols = ["chrom", "pos", "ref", "alt", attribute, parent_attribute]

    # Parse dataframe columns
    rename_cols = dict(zip(subset_cols, new_cols))
    df = df.select(subset_cols).rename(rename_cols).with_columns(pl.col("pos").cast(pl.UInt32))

    return df.unique(maintain_order=True).collect()


def parse_intersect_genes_new(
    intersect_file: str | Path,
    attribute: str | None = None,
    parent_attribute: str | None = None,
) -> pl.DataFrame:
    """Parse gene intersection file with typed columns.

    Parameters
    ----------
    intersect_file : str | Path
        Path to bedtools intersection output.
    attribute : str | None, optional
        Attribute column name, by default 'ID'.
    parent_attribute : str | None, optional
        Parent attribute column name, by default 'Parent'.

    Returns
    -------
    pl.DataFrame
        Parsed intersection data with typed columns.
    """
    if attribute is None:
        attribute = "ID"

    if parent_attribute is None:
        parent_attribute = "Parent"

    # AFTER performing gtf_to_bed and intersecting!
    df = pl.scan_csv(intersect_file, separator="\t", has_header=False, infer_schema_length=0)

    vcf_schema = [
        pl.col("chrom").cast(pl.Categorical),
        pl.col("pos").cast(pl.UInt32),
        pl.col("ref").cast(pl.Categorical),
        pl.col("alt").cast(pl.Categorical),
    ]

    # Expect at min 10 cols, 11 if GT included
    if len(df.columns) > 10:
        subset_cols = [df.columns[i] for i in [0, 2, 3, 4, 5, -2, -1]]
        new_cols = ["chrom", "pos", "ref", "alt", "GT", attribute, parent_attribute]
        vcf_schema.append(pl.col("GT").cast(pl.Categorical))
    else:
        subset_cols = [df.columns[i] for i in [0, 2, 3, 4, -2, -1]]
        new_cols = ["chrom", "pos", "ref", "alt", attribute, parent_attribute]

    # Parse dataframe columns
    rename_cols = dict(zip(subset_cols, new_cols))
    df = df.select(subset_cols).rename(rename_cols).with_columns(vcf_schema)

    return df.unique(maintain_order=True).collect()
