from pydantic import BaseModel, PositiveInt
from pydantic.functional_validators import AfterValidator
from typing import Annotated
import re
from enum import Enum

from primal_page.bedfiles import BedfileVersion

INFO_SCHEMA = "v1.3.0"

SCHEMENAME_PATTERN = r"^[a-z0-9][a-z0-9-]*[a-z0-9]$"
VERSION_PATTERN = r"^v\d+\.\d+\.\d+$"


class PrimerClass(Enum):
    PRIMERSCHEMES = "primerschemes"


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


def validate_schemeversion(version: str) -> str:
    if not re.match(VERSION_PATTERN, version):
        raise ValueError(
            f"Invalid version: {version}. Must match be in form of v(int).(int).(int)"
        )
    return version


def validate_schemename(schemename: str) -> str:
    if not re.match(SCHEMENAME_PATTERN, schemename):
        raise ValueError(
            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) -> list | set:
    if len(x) == 0:
        raise ValueError("Cannot be empty")
    return x


class Collection(Enum):
    # Authors
    ARTIC = "ARTIC"
    MODJADJI = "MODJADJI"
    QUICK_LAB = "QUICK-LAB"
    COMMUNITY = "COMMUNITY"
    # Uses
    WASTE_WATER = "WASTE-WATER"
    CLINAL_ISOLATES = "CLINAL-ISOLATES"
    # Coverage types
    WHOLE_GENOME = "WHOLE-GENOME"
    PANEL = "PANEL"
    MULTI_TARGET = "MULTI-TARGET"


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[set[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
    # Add the optional fields
    description: str | None = None
    derivedfrom: str | None = None
    collections: set[Collection] = set()


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=set("artic"),
        algorithmversion="test",
        species=set("sars-cov-2"),
        articbedversion=BedfileVersion.V3,
        collections=set(),
    )

    # indexv = IndexVersion(
    #    ampliconsize=400,
    #    schemeversion="v0.0.0",
    #    schemename="test",
    #    primer_bed_md5="hello",
    #    reference_fasta_md5="world",
    #    status=SchemeStatus.DRAFT,
    #    citations=[],
    #    authors=[],
    #    algorithmversion="test",
    #    species=[1, "hello", 1],
    #    IndexVersion="test",
    # )

    print(info.model_dump_json(indent=4))

    info.authors.add("hello")
    print(info.model_dump_json(indent=4))

    # print(indexv.model_dump_json(indent=4))
