# -*- coding: utf-8 -*-
"""GDICT implementation for handling gene nomenclature data.

This module provides classes and functions to read, write, and manipulate gene
nomenclature data in GDICT format. It includes the `GeneDict` class for storing
gene information, as well as utility functions for reading GDICT files, sorting
gene entries, and creating stripped GDICT version.

"""

import re
from collections import UserDict, defaultdict
from dataclasses import dataclass, field, replace
from datetime import datetime
from pathlib import Path
from typing import Callable, Iterable, Mapping, Optional, TypeVar, Union

__gdict_version__ = "0.0.2"


@dataclass(slots=True)
class _Gene:
    label: str
    c: Optional[str] = field(default=None, kw_only=True)


@dataclass(slots=True)
class DbxrefGeneID(_Gene):
    """A gene with a Dbxref GeneID."""

    an_source: str
    GeneID: int


@dataclass(slots=True)
class GeneGeneric(_Gene):
    """A generic gene ID with multiple AN sources."""

    an_sources: list[str] = field(default_factory=list)


@dataclass(slots=True)
class GeneDescription(_Gene):
    """A description of a gene."""

    source: str


@dataclass(slots=True)
class GeneDictInfo:
    """Information about the gene dictionary.

    Attributes:
        labels (int): Number of unique labels in the dictionary.
        total_entries (int): Total number of entries in the dictionary.
        gene_descriptions (int): Count of `GeneDescription` entries.
        gene_generics (int): Count of `GeneGeneric` entries.
        dbxref_GeneIDs (int): Count of `DbxrefGeneID` entries.

    """

    labels: int = 0
    total_entries: int = 0
    gene_descriptions: int = 0
    gene_generics: int = 0
    dbxref_GeneIDs: int = 0


T = TypeVar("T")
G = TypeVar("G", bound=_Gene)

GeneUnion = Union[GeneDescription, GeneGeneric, DbxrefGeneID]

GeneList = list[tuple[str, G]]

SortedGeneGroups = tuple[
    GeneList[GeneDescription], GeneList[GeneGeneric], GeneList[DbxrefGeneID]
]


class GeneDict(UserDict[str, GeneUnion]):
    """A dictionary to store gene nomeclature information with additional metadata.

    Attributes:
        version (str): Version of the GDICT format.
        header (list[str]): Header lines for the GDICT file.
        info (GeneDictInfo): Information about the gene dictionary.

    """

    def __init__(
        self,
        initial: Optional[Mapping[str, GeneUnion]] = None,
        *,
        version: str = __gdict_version__,
        header: Optional[list[str]] = None,
        info: Optional[GeneDictInfo] = None,
        lazy_info: bool = True,
    ):
        """Initialize a GeneDict with optional initial data, version, header, and info.

        Args:
            initial (dict): Initial data for the GeneDict.
            version (str): Version of the GDICT format. Default is __gdict_version__.
            header (list[str]): Header lines for the GDICT file.
            info (GeneDictInfo): Information about the gene dictionary. This normally
                                 is generated by the `update_info` method,
                                 but can be set manually.
            lazy_info (bool): If False, `update_info` will be called at initialization.
                              Default is True, meaning the info will not be updated
                              until `update_info` is called.

        """
        super().__init__(initial or {})

        self.version: str = version
        self.header: list[str] = header or []
        self.info: GeneDictInfo = info or GeneDictInfo()

        if not lazy_info:
            self.update_info()

    @classmethod
    def _from_data(
        cls,
        inital: dict[str, GeneUnion],
        *,
        version: str = __gdict_version__,
        header: Optional[list[str]] = None,
        info: Optional[GeneDictInfo] = None,
        lazy_info: bool = True,
    ) -> "GeneDict":
        """Fast constructor for internal use with pre-validated data.

        This method is intended for internal use when the data has already been
        validated (mostly by coming from another GeneDict) and is ready to be used
        to create a GeneDict instance. It skips the `__init__` method, which would've
        call `self.update()` (`self.data.update(data)`) on all key/value pairs from
        `initial`.
        """
        instance = cls.__new__(cls)
        instance.data = inital

        instance.version = version
        instance.header = header or []
        instance.info = info or GeneDictInfo()

        if not lazy_info:
            instance.update_info()

        return instance

    def to_gdict(
        self,
        gdict_file: Union[str, Path],
        overwrite: bool = False,
    ) -> None:
        """Write a gene dictionary to a GDICT file, sorted by label.

        Args:
            gdict_file (str): Path to the GDICT file.
            overwrite (bool): Whether to overwrite the GDICT file if it already exists.
                Default is False.

        """
        gdict_file = Path(gdict_file).resolve()

        if gdict_file.exists() and not overwrite:
            raise FileExistsError(
                f"GDICT file already exists: {gdict_file}. Use overwrite=True to "
                "overwrite."
            )

        if self.version != __gdict_version__:
            raise ValueError(
                f"GDICT not on version {__gdict_version__}. GDICT version: "
                f"{self.version}"
            )

        all_labels: list[str] = natural_sort(
            {gene.label for gene in self.data.values()}
        )

        label_as_key = defaultdict(list)
        for key, value in self.data.items():
            label_as_key[value.label].append((key, value))

        all_sorted: dict[str, SortedGeneGroups] = {}
        for label, values in label_as_key.items():
            gd: GeneList[GeneDescription] = []
            gn: GeneList[GeneGeneric] = []
            dx: GeneList[DbxrefGeneID] = []
            for key, value in values:
                match value:
                    case GeneDescription():
                        gd.append((key, value))
                    case GeneGeneric():
                        gn.append((key, value))
                    case DbxrefGeneID():
                        dx.append((key, value))

            # Sort each group once
            all_sorted[label] = (
                natural_sort(gd, key=lambda x: x[0]),
                natural_sort(gn, key=lambda x: x[0]),
                natural_sort(dx, key=lambda x: x[0]),
            )

        with open(gdict_file, "w") as f:
            for line in self.header:
                f.write(f"#! {line}\n")

            for label in all_labels:
                f.write(f"\n[{label}]\n")
                gd, gn, dx = all_sorted[label]

                for key, value in gd:
                    f.write(
                        f"{key} #gd {value.source}"
                        f"{' #c ' + value.c if value.c else ''}\n"
                    )

                for key, value in gn:
                    if value.an_sources:
                        f.write(
                            f"{key} #gn {' '.join(value.an_sources)}"
                            f"{' #c ' + value.c if value.c else ''}\n"
                        )
                    else:
                        f.write(f"{key} #gn{' #c ' + value.c if value.c else ''}\n")

                for key, value in dx:
                    f.write(
                        f"{key} #dx {value.an_source}:{value.GeneID}"
                        f"{' #c ' + value.c if value.c else ''}\n"
                    )

    def update_info(self) -> None:
        """Update the information about the gene dictionary.

        This method populates the `info` attribute with the number of unique labels,
        total entries, and counts of each type of gene entry.

        `info` attribute is an instance of `GeneDictInfo` with the following attributes:
            - labels: Number of unique labels in the dictionary.
            - total_entries: Total number of entries in the dictionary.
            - gene_descriptions: Count of `GeneDescription` entries.
            - gene_generics: Count of `GeneGeneric` entries.
            - dbxref_GeneIDs: Count of `DbxrefGeneID` entries.

        """
        self.info = GeneDictInfo()  # Reset info
        labels = set()

        for value in self.data.values():
            labels.add(value.label)
            match value:
                case DbxrefGeneID():
                    self.info.dbxref_GeneIDs += 1

                case GeneGeneric():
                    self.info.gene_generics += 1

                case GeneDescription():
                    self.info.gene_descriptions += 1

                case _:
                    print(f"[INFO] Unknown type for {value}")

        self.info.labels = len(labels)
        self.info.total_entries = len(self.data)

    def create_stripped(
        self,
        keep_gn: bool = False,
        lazy_info: bool = True,
    ) -> "GeneDict":
        """Create a stripped version of the GeneDict.

        This method creates a new GeneDict instance that contains only the
        `GeneDescription` and optionally `GeneGeneric` entries, removing any
        `DbxrefGeneID` entries. The new GeneDict will have an updated header
        indicating that it is a stripped version.

        Args:
            keep_gn (bool): If True, `GeneGeneric` entries will be kept in the
                            stripped version. Default is False, meaning only
                            `GeneDescription` entries will be kept.
            lazy_info (bool): If False, `update_info` will be called on the stripped
                              GeneDict. Default is True, meaning the info will not
                              be updated until `update_info` is called.

        Returns:
            GeneDict: A new GeneDict instance with stripped data and an updated header.

        """
        header = self.header.copy()
        header.append(f"{time_now()} - Stripped GDICT version from original GDICT file")

        check = (GeneDescription, GeneGeneric) if keep_gn else GeneDescription

        stripped_data = {
            key: replace(value)
            for key, value in self.data.items()
            if isinstance(value, check)
        }

        # Since stripped came from an already existing GeneDict (self),
        # we can optimize the creation of the new GeneDict
        return self._from_data(
            stripped_data,
            version=self.version,
            header=header,
            lazy_info=lazy_info,
        )

    def rename_labels(
        self,
        names: Mapping[str, str],
        lazy_info: bool = True,
    ) -> "GeneDict":
        """Rename labels in the GeneDict.

        This method replaces the labels of all entries in the GeneDict according to
        the provided mapping. If a label is not found in the mapping, it remains
        unchanged.

        Args:
            names (Mapping[str, str]): A mapping of old label names to new label names.
            lazy_info (bool): If False, `update_info` will be called on the renamed
                              GeneDict. Default is True, meaning the info will not
                              be updated until `update_info` is called.

        Returns:
            GeneDict: A new GeneDict instance with renamed labels and an updated header.

        """
        header = self.header.copy()
        header.append(f"{time_now()} - Labels renamed")

        renamed_data = {
            key: (
                replace(value, label=names[value.label])
                if value.label in names
                else replace(value)
            )
            for key, value in self.data.items()
        }

        # Since renamed_data came from an already existing GeneDict (self),
        # we can optimize the creation of the new GeneDict
        return self._from_data(
            renamed_data,
            version=self.version,
            header=header,
            lazy_info=lazy_info,
        )

    def remove_labels(
        self,
        labels: Iterable[str],
        lazy_info: bool = True,
    ) -> "GeneDict":
        """Remove entries with specified labels from the GeneDict.

        This method removes all entries that have a label in the provided iterable
        of labels.

        Args:
            labels (Iterable[str]): An iterable of labels to remove from the GeneDict.
            lazy_info (bool): If False, `update_info` will be called on the modified
                              GeneDict. Default is True, meaning the info will not
                              be updated until `update_info` is called.

        Returns:
            GeneDict: A new GeneDict instance with specified labels removed and an
                      updated header.

        """
        if isinstance(labels, str):
            labels = [labels]

        if not all(isinstance(label, str) for label in labels):
            raise TypeError("All labels must be strings")

        header = self.header.copy()
        header.append(f"{time_now()} - Labels removed")

        filtered_data = {
            key: value for key, value in self.data.items() if value.label not in labels
        }

        # Since filtered_data came from an already existing GeneDict (self),
        # we can optimize the creation of the new GeneDict
        return self._from_data(
            filtered_data,
            version=self.version,
            header=header,
            lazy_info=lazy_info,
        )

    def __str__(self) -> str:
        """Return a string representation of the GeneDict."""
        if self.info.labels == 0:
            self.update_info()

        return (
            f"GeneDict(version={self.version}, header={self.header}, info={self.info})"
        )


def time_now() -> str:
    """Return the current time formatted as a string."""
    return datetime.now().strftime("%Y-%m-%d %H:%M")


def read_gdict(
    gdict_file: Union[str, Path],
    max_an_sources: int = 0,
    lazy_info: bool = True,
) -> GeneDict:
    """Read a gdict file into GeneDict.

    Args:
        gdict_file (str): Path to the GDICT file.
        max_an_sources (int): Maximum number of AN sources to include in GeneGeneric.
                            If set to 0, all sources will be included. Default is 0.
        lazy_info (bool): If False, `update_info` will be called on the created
                          GeneDict. Default is True, meaning the info will not be
                          updated until `update_info` is called.

    Returns:
        GeneDict: A dictionary containing the gdict file information.

    """
    gdict_file = Path(gdict_file).resolve()

    if not gdict_file.exists():
        raise FileNotFoundError(f"GDICT file not found: {gdict_file}")

    with open(gdict_file, "r") as f:
        lines = [line.strip() for line in f.read().split("\n") if line.strip()]

    if lines[0] != "#! version 0.0.2":
        raise ValueError(
            f"Invalid GDICT file version: {lines[0]}. Expected '#! version 0.0.2'"
        )

    result = GeneDict()
    for line in lines:
        if line.startswith("#!"):
            result.header.append(line[2:].strip())
            continue
        else:
            break

    current_section = None
    for line in lines:
        if line.startswith("[") and line.endswith("]"):
            current_section = line[1:-1]
            continue

        # Skip if no section is defined
        if not line or not current_section:
            continue

        # Parse entry line
        tag = line.split("#", 1)[0].strip()
        if not tag:
            print(f"Skipping empty tag in line: {line}")
            continue

        if "#c" in line:
            line, comment = line.split("#c", 1)
            comment = comment.strip()
            line = line.strip()
        else:
            comment = None

        if "#dx" in line:  # DbxrefGeneID dx
            stuff = line.split("#dx", 1)[1].strip()
            an_source = stuff.split(":")[0].strip()
            dbxref = int(stuff.split(":")[1].strip())

            result[tag] = DbxrefGeneID(
                label=current_section,
                an_source=an_source,
                GeneID=dbxref,
                c=comment,
            )

        elif "#gn" in line:  # GeneGeneric gn
            an_sources = [s.strip() for s in line.split("#gn", 1)[1].strip().split()]
            an_sources = an_sources if an_sources else []

            if len(an_sources) >= max_an_sources and max_an_sources > 0:
                an_sources = an_sources[:max_an_sources]
                comment = comment if comment else ""
                comment += f" |More than {max_an_sources} sources,"
                f"adding only the first {max_an_sources}|"

            result[tag] = GeneGeneric(
                label=current_section,
                an_sources=an_sources,
                c=comment,
            )

        elif "#gd" in line:  # GeneDescription gd
            source = line.split("#gd", 1)[1].strip()

            result[tag] = GeneDescription(
                label=current_section,
                source=source,
                c=comment,
            )

    if not lazy_info:
        result.update_info()

    return result


def natural_sort_key(s: str) -> list[Union[int, str]]:
    """Generate a natural sort key for a string.

    This function splits the string into parts of digits and non-digits,
    converting digit parts to integers and leaving non-digit parts as strings.

    Args:
        s (str): The string to generate a natural sort key for.

    Returns:
        list[Union[int, str]]: A list where digit parts are converted to integers
                               and non-digit parts are converted to lowercase strings.

    """
    return [int(c) if c.isdigit() else c.lower() for c in re.split(r"(\d+)", s)]


def natural_sort(
    iterable: Iterable[T],
    key: Callable[[T], str] | None = None,
    reverse: bool = False,
) -> list[T]:
    """Sort a list in natural order.

    This function sorts an iterable in natural order, which means that it sorts
    strings in a way that is intuitive to humans, taking into account numeric
    values within the strings. If a key function is provided, it will be used
    to extract the comparison key from each element before sorting.
    If the key is None, it will sort the strings directly using the natural sort key.

    Args:
        iterable: An iterable to sort.
        key: Function to extract comparison key from each element.
        reverse: Whether to sort in reverse order.

    Returns:
        A sorted list in natural order.

    """
    if key is None:  # Original behavior for simple strings
        # Type narrowing: if key is None, T must be str
        return sorted(iterable, key=natural_sort_key, reverse=reverse)  # type: ignore[arg-type]

    return sorted(iterable, key=lambda x: natural_sort_key(key(x)), reverse=reverse)


def create_empty_gdict(
    gdict_file: Union[Path, str],
    default_text: str = "Empty gdict file",
    overwrite: bool = False,
) -> None:
    """Create an empty GDICT file.

    This function creates an empty GDICT file with a default header and a timestamp.

    Args:
        gdict_file (Union[Path, str]): Path to the GDICT file to create.
        default_text (str): Default text to include in the header.
                            Default is 'Empty gdict file'.
        overwrite (bool): Whether to overwrite the file if it already exists.
                          Default is False.

    """
    gdict_file = Path(gdict_file).resolve()

    if not overwrite and gdict_file.exists():
        raise FileExistsError(
            f"GDICT file already exists: {gdict_file}. Use overwrite=True to overwrite."
        )

    with open(gdict_file, "w") as f:
        f.write(f"#! version {__gdict_version__}\n")
        f.write(f"#! {time_now()} - {default_text}\n")
