import hashlib
import pathlib
import re
from enum import Enum
from typing import Annotated

from Bio import SeqIO
from primalbedtools.bedfiles import BedLineParser
from pydantic import BaseModel, PositiveInt
from pydantic.functional_validators import AfterValidator

from primal_page.bedfiles import BedfileVersion
from primal_page.errors import InvalidSchemeID, InvalidSchemeName, InvalidSchemeVersion


def hash_file(fname: pathlib.Path) -> str:
    hash_md5 = hashlib.md5()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


INFO_SCHEMA = "v2.1.0"

SCHEMENAME_PATTERN = r"^[a-z0-9][a-z0-9-]*[a-z0-9]$"
VERSION_PATTERN = r"^v\d+\.\d+\.\d+(-[a-z0-9]+)?$"
DOI_PATTERN = r"^(https:\/\/|dx.)*(doi.org\/|doi:)10.\d+\/.*$"


class PrimerClass(Enum):
    PRIMERSCHEMES = "primerschemes"


class SchemeStatus(Enum):
    WITHDRAWN = "withdrawn"
    DEPRECATED = "deprecated"
    AUTOGENERATED = "autogenerated"
    DRAFT = "draft"
    TESTING = "tested"
    VALIDATED = "validated"


def validate_scheme_id(schemeid) -> tuple[str, str, str]:
    """
    Parse the schemeid into its components
    :raises InvalidSchemeID: if the schemeid is invalid
    """
    try:
        (
            schemename,
            ampliconsize,
            schemeversion,
        ) = schemeid.split("/")
    except ValueError as e:
        raise InvalidSchemeID(
            f"{schemeid} needs to be in the form (schemename)/(ampliconsize)/(schemeversion)"
        ) from e

    # Validate each part separately
    try:
        schemename = validate_schemename(schemename)
        schemeversion = validate_schemeversion(schemeversion)
        ampliconsize = int(ampliconsize)
    except InvalidSchemeName as e:
        raise InvalidSchemeID(
            f"{schemename} is an invalid schemename in {schemeid}"
        ) from e
    except InvalidSchemeVersion as e:
        raise InvalidSchemeID(
            f"{schemeversion} is an invalid schemeversion in {schemeid}"
        ) from e
    except ValueError as e:
        raise InvalidSchemeID(
            f"{ampliconsize} is an invalid ampliconsize in {schemeid}"
        ) from e

    return schemename, str(ampliconsize), schemeversion


def validate_schemeversion(version: str) -> str:
    """
    Validate the schemeversion. Must be in the form of v(int).(int).(int) with optional suffix -[a-z0-9].
    :raises InvalidSchemeVersion: if the schemeversion is invalid
    """
    if not re.match(VERSION_PATTERN, version):
        raise InvalidSchemeVersion(
            f"Invalid version: ({version}). Must match be in form of v(int).(int).(int) with optional suffix -[a-z0-9]"
        )
    return version


def validate_schemename(schemename: str) -> str:
    """
    Validate the schemename
    :raises InvalidSchemeName: if the schemename is invalid
    """
    if not re.match(SCHEMENAME_PATTERN, schemename):
        raise InvalidSchemeName(
            f"Invalid schemename: ({schemename}). Must only contain a-z, 0-9, and -. Cannot start or end with -"
        )
    return schemename


def not_empty(x: list | set | str) -> list | set | str:
    if len(x) == 0:
        raise ValueError("Cannot be empty")
    return x


IUPACAmbiguousDNA = {
    "A",
    "G",
    "K",
    "Y",
    "B",
    "S",
    "N",
    "H",
    "C",
    "W",
    "D",
    "R",
    "M",
    "T",
    "V",
}


class Collection(Enum):
    # Authors
    ARTIC = "ARTIC"
    MODJADJI = "MODJADJI"
    QUICK_LAB = "QUICK-LAB"
    COMMUNITY = "COMMUNITY"
    # Uses
    WASTE_WATER = "WASTE-WATER"
    CLINICAL_ISOLATES = "CLINICAL-ISOLATES"
    # Coverage types
    WHOLE_GENOME = "WHOLE-GENOME"
    PANEL = "PANEL"
    MULTI_TARGET = "MULTI-TARGET"
    # Other Scheme Generators
    VARVAMP = "VARVAMP"

    def __ord__(self):
        return self.value


class Links(BaseModel):
    protocols: list[str] = []
    validation: list[str] = []
    homepage: list[str] = []
    vendors: list[str] = []
    misc: list[str] = []

    def getattr(self, link_type: str):
        """
        Get the list of links
        Raises:
        -   AttributeError if the link_type is not in the model_fields
        """
        return getattr(self, link_type)

    def append_link(
        self,
        link_type: str,
        link: str,
    ):
        """
        Add a link to the correct list
        Raises:
        -   AttributeError if the link_type is not in the model_fields
        """
        # Append the link to the correct list
        getattr(self, link_type).append(link)

    def link_remove(
        self,
        link_type: str,
        link: str,
    ):
        """
        Remove a the first occurrence of the link from the correct list
        Raises:
        -   ValueError if the link is not in the list
        -   AttributeError if the link_type is not in the model_fields
        """
        getattr(self, link_type).remove(link)


class Info(BaseModel):
    ampliconsize: PositiveInt
    schemeversion: Annotated[str, AfterValidator(validate_schemeversion)]
    schemename: Annotated[str, AfterValidator(validate_schemename)]
    primer_bed_md5: str
    reference_fasta_md5: str
    status: SchemeStatus
    citations: set[str]
    authors: Annotated[list[str], AfterValidator(not_empty)]
    algorithmversion: str
    species: Annotated[set[int | str], AfterValidator(not_empty)]
    license: str = "CC BY-SA 4.0"
    primerclass: PrimerClass = PrimerClass.PRIMERSCHEMES
    infoschema: str = INFO_SCHEMA
    articbedversion: BedfileVersion
    collections: set[Collection] = set()
    links: Links = Links()
    refselect: dict[str, dict[str, str]] | None = None

    # Add the optional fields
    description: str | None = None
    derivedfrom: str | None = None
    contactinfo: str | None = None

    def add_refselect(
        self, ref_select_file: pathlib.Path, chrom: str, info_path: pathlib.Path
    ):
        """
        Add a reference selection file to the ref select dictionary
        """
        if "-" in self.schemeversion:
            raise ValueError(
                "Cannot add reference selection to a version with a mapping suffix"
            )

        # Validate the reference selection file
        validate_ref_select_file(self, chrom, ref_select_file, info_path)

        # Read the ref selection into memory
        new_ref_select_path = info_path.parent / f"{chrom}_refselect.fasta"

        ref_select = list(SeqIO.parse(ref_select_file, "fasta"))
        with open(new_ref_select_path, "w") as f:
            SeqIO.write(ref_select, f, "fasta")

        # Add the reference selection file
        if self.refselect is None:
            self.refselect = {}
        self.refselect[chrom] = {
            "md5": hash_file(new_ref_select_path),
            "filename": new_ref_select_path.name,
        }  # type: ignore

    def collection_add(self, collection: Collection):
        """Add a collection to the collections set"""
        self.collections.add(collection)

    def collection_remove(self, collection: Collection):
        """
        Remove a collection from the collections set.
        Raises: KeyError if the collection is not in the set
        """
        self.collections.remove(collection)

    def citation_add(self, citation: str):
        """Add a citation to the citations set"""
        self.citations.add(citation)

    def citation_remove(self, citation: str):
        """
        Remove a citation from the citations set.
        Raises: KeyError if the citation is not in the set
        """
        self.citations.remove(citation)

    def author_add(self, author: str, author_index: int | None):
        """Add an author to the authors list"""
        if author not in self.authors:
            if author_index is None:
                self.authors.append(author)
            else:
                # Insert is safe, will append if index is out of range
                self.authors.insert(author_index, author)

    def author_remove(self, author: str):
        """
        Remove an author from the authors list.
        Raises: ValueError if the author is not in the list
        """
        self.authors.remove(author)

    def author_reorders(self, new_order: list[int]):
        """
        Reorder the authors list
        Raises:
            ValueError if duplicate indexes are found
            IndexError if the index is out of range
        """
        if len(new_order) != len(set(new_order)):
            raise ValueError("Duplicate indexes found")
            # Append authors in the new order

        new_authors = []
        for new_index in new_order:
            if new_index >= len(self.authors) or new_index < 0:
                raise IndexError(f"{new_index} is out of range")
            new_authors.append(self.authors[new_index])

        # Append any authors not in the new order
        for index, author in enumerate(self.authors):
            if index not in new_order:
                new_authors.append(author)

        self.authors = new_authors

    def change_primerclass(self, new_class: PrimerClass):
        """Change the primerclass"""
        self.primerclass = new_class

    def change_status(self, new_status: SchemeStatus):
        """Change the status"""
        self.status = new_status

    def change_license(self, new_license: str):
        """Change the license"""
        self.license = new_license

    def change_derivedfrom(self, new_derivedfrom: str):
        """Change the derivedfrom"""
        self.derivedfrom = new_derivedfrom

    def change_description(self, new_description: str):
        """Change the description"""
        self.description = new_description

    def change_contactinfo(self, new_contactinfo: str | None):
        """Change the contactinfo"""
        self.contactinfo = new_contactinfo

    def get_schemepath(self) -> str:
        return f"{self.schemename}/{self.ampliconsize}/{self.schemeversion}"


def validate_ref_select_file(
    info: Info, chrom: str, ref_select: pathlib.Path, infopath: pathlib.Path
):
    """
    Validate the reference selection file
    """
    # Read in the bedfile
    _header, bedlines = BedLineParser.from_file(infopath.parent / "primer.bed")
    bedline_chrom_names = {bedline.chrom for bedline in bedlines}

    # Check the chrom is in the bedfile
    if chrom not in bedline_chrom_names:
        raise ValueError(f"Chromosome ({chrom}) not in primer.bed")

    # Check the file exists
    if not ref_select.is_file():
        raise FileNotFoundError(f"{ref_select} is not a file")

    # Check the file it self
    ref_select_index = SeqIO.index(str(ref_select), "fasta")
    # Check genomes len
    ref_select_lengths = {len(seq) for seq in ref_select_index.values()}
    if len(ref_select_lengths) != 1:
        raise ValueError(
            f"Chrom ({chrom}) in {infopath}:{ref_select.name} has different lengths: {ref_select_lengths}"
        )
    # Check chrom inside fasta
    if chrom not in ref_select_index:
        raise ValueError(f"Chrom ({chrom}) not in {infopath}:{ref_select.name}")

    # Check the fasta headers
    for record in ref_select_index.values():
        if len(record.description.split()) != 2:
            raise ValueError(
                f"Chrom {chrom} in {infopath}:{ref_select.name} has fasta headers without clade tags"
            )

    return


if __name__ == "__main__":
    info = Info(
        ampliconsize=400,
        schemeversion="v0.0.0",
        schemename="test",
        primer_bed_md5="hello",
        reference_fasta_md5="world",
        status=SchemeStatus.DRAFT,
        citations=set(),
        authors=["artic"],
        algorithmversion="test",
        species=set(["sars-cov-2"]),
        articbedversion=BedfileVersion.V3,
        collections=set(),
        links=Links(),
    )

    info_json = info.model_dump_json()
