import argparse
import graphlib
import itertools
import logging
import os
from enum import StrEnum
from types import UnionType
from typing import Any, get_args

import yaml
from pydantic import BaseModel, Field, PrivateAttr, model_validator

from loculus_preprocessing.datatypes import (
    FileCategory,
    FunctionArgs,
    FunctionInputs,
    FunctionName,
    MoleculeType,
    SegmentClassificationMethod,
    Topology,
)
from loculus_preprocessing.external_services import (
    FileProcessingService,
    TaxonomyService,
)

logger = logging.getLogger(__name__)

# Dataclass types for which we can generate CLI arguments
CLI_TYPES = [str, int, float, bool]

PROCESSED_PREFIX = "processed."
FILES_PREFIX = "files."
NEXTCLADE_PREFIX = "nextclade."
ASSIGNED_REFERENCE_PREFIX = "ASSIGNED_REFERENCE"
INTERNAL_INPUT_PREFIXES = (NEXTCLADE_PREFIX, ASSIGNED_REFERENCE_PREFIX)


class EmblInfoMetadataPropertyNames(BaseModel):
    country_property: str = "geoLocCountry"
    admin_level_properties: list[str] = Field(
        default_factory=lambda: ["geoLocAdmin1", "geoLocAdmin2", "geoLocCity", "geoLocSite"]
    )
    collection_date_property: str = "sampleCollectionDate"
    authors_property: str = "authors"


class ProcessingSpec(BaseModel):
    inputs: FunctionInputs
    function: FunctionName = "identity"
    required: bool = False
    required_when: list[str] = []
    no_input: bool = False
    args: FunctionArgs | None = None


class AlignmentRequirement(StrEnum):
    # Determines whether ALL or ANY segments that a user provides must align.
    # ANY: warn if not all segments can be aligned (all must be assigned) and at least one must fully align
    # ALL: error if any segment fails alignment even if some segments align
    # NONE: do not align any segments, just process them as-is
    # - set if no nextclade dataset is provided
    ANY = "ANY"
    ALL = "ALL"
    NONE = "NONE"


type SegmentName = str
# name of the processed nucleotide sequence, as expected by the backend and LAPIS
type SequenceName = str


class Reference(BaseModel):
    name: str = "singleReference"
    nextclade_dataset_name: str | None = None
    nextclade_dataset_tag: str | None = None
    nextclade_dataset_server: str | None = None
    accepted_dataset_matches: list[str] = Field(default_factory=list)
    genes: list[str] = Field(default_factory=list)
    nextclade_additional_args: list[str] = Field(default_factory=list)


class Segment(BaseModel):
    name: SegmentName = "main"
    references: list[Reference] = Field(default_factory=list)


class NextcladeSequenceAndDataset(BaseModel):
    name: SequenceName = "main"
    reference_name: str = "singleReference"
    segment: SegmentName = "main"
    nextclade_dataset_name: str | None = None
    nextclade_dataset_tag: str | None = None
    nextclade_dataset_server: str | None = None
    # Names of diamond or nextclade sort entries that are acceptable matches for this dataset
    accepted_dataset_matches: list[str] = Field(default_factory=list)
    gene_suffix: str | None = None
    # Names of genes in the Nextclade dataset; when concatenated with gene_suffix
    # this must match the gene names expected by the backend and LAPIS
    genes: list[str] = Field(default_factory=list)
    nextclade_additional_args: list[str] = Field(default_factory=list)


class Config(BaseModel):
    log_level: str = "DEBUG"
    keep_tmp_dir: bool = False
    batch_size: int = 5
    pipeline_version: int = 1
    backend_request_timeout_seconds: int = 30

    backend_host: str = ""  # base API URL and organism - populated in get_config if left empty
    keycloak_host: str = "http://127.0.0.1:8083"
    keycloak_user: str = "preprocessing_pipeline"
    keycloak_password: str = "preprocessing_pipeline"  # noqa: S105
    keycloak_token_path: str = "realms/loculus/protocol/openid-connect/token"  # noqa: S105

    organism: str = "mpox"
    max_sequences_per_entry: int | None = None
    segments: list[Segment] = Field(default_factory=list)
    processing_spec: dict[str, ProcessingSpec] = Field(default_factory=dict)
    processing_order: tuple[str, ...] = ()

    alignment_requirement: AlignmentRequirement = AlignmentRequirement.ALL
    segment_classification_method: SegmentClassificationMethod = SegmentClassificationMethod.ALIGN
    nextclade_dataset_server: str = "https://data.clades.nextstrain.org/v3"

    require_nextclade_sort_match: bool = False
    minimizer_url: str | None = None
    diamond_dmnd_url: str | None = None

    create_embl_file: bool = False
    scientific_name: str = "Orthonairovirus haemorrhagiae"
    molecule_type: MoleculeType = MoleculeType.GENOMIC_RNA
    topology: Topology = Topology.LINEAR
    db_name: str = "Loculus"
    # The 'embl' section of the config contains metadata property names for the EMBL file
    embl: EmblInfoMetadataPropertyNames = Field(default_factory=EmblInfoMetadataPropertyNames)
    insdc_ingest_group_id: int = 1

    # External services
    taxonomy_service_url: str | None = None
    _taxonomy_service: TaxonomyService = PrivateAttr(default=TaxonomyService(None))
    raw_reads_processing_service_url: str | None = None
    raw_reads_processing_service_timeout_seconds: int = 600
    _file_processing_service: FileProcessingService = PrivateAttr(
        default=FileProcessingService(None, 600)
    )

    @model_validator(mode="after")
    def finalize(self):
        if not self.segments:
            self.alignment_requirement = AlignmentRequirement.NONE
        for segment in self.segments:
            if not segment.references or not any(
                ds.nextclade_dataset_name for ds in segment.references
            ):
                self.alignment_requirement = AlignmentRequirement.NONE

        if not self.backend_host:  # Set here so we can use organism
            self.backend_host = f"http://127.0.0.1:8079/{self.organism}"

        self._taxonomy_service = TaxonomyService(self.taxonomy_service_url)
        self._file_processing_service = FileProcessingService(
            self.raw_reads_processing_service_url, self.raw_reads_processing_service_timeout_seconds
        )

        validate_required_when(self)
        self.processing_order = get_processing_order(self)
        self._taxonomy_service = TaxonomyService(self.taxonomy_service_url)

        return self

    @property
    def multi_segment(self) -> bool:
        return len(self.segments) > 1

    @property
    def nextclade_sequence_and_datasets(self) -> list[NextcladeSequenceAndDataset]:
        def build_ds(
            reference: Reference | None, segment_name: str, multi_reference: bool
        ) -> NextcladeSequenceAndDataset:
            base = reference.model_dump() if reference else {}
            ds = NextcladeSequenceAndDataset(
                **base,
                segment=segment_name,
                reference_name=reference.name if reference else "singleReference",
            )
            if ds.nextclade_dataset_server is None:
                ds.nextclade_dataset_server = self.nextclade_dataset_server
            ds.name = set_sequence_name(multi_reference, self.multi_segment, ds)
            ds.gene_suffix = ds.reference_name if multi_reference else None
            return ds

        datasets: list[NextcladeSequenceAndDataset] = []

        for segment in self.segments:
            multi_reference = len(segment.references) > 1
            references: list[Reference] | list[None] = segment.references or [None]
            datasets.extend(build_ds(ref, segment.name, multi_reference) for ref in references)

        return datasets

    @property
    def multi_datasets(self) -> bool:
        return len(self.nextclade_sequence_and_datasets) > 1

    def get_dataset_by_name(self, name: str) -> NextcladeSequenceAndDataset:
        datasets = [ds for ds in self.nextclade_sequence_and_datasets if ds.name == name]
        if len(datasets) == 0:
            msg = f"No dataset found with name: {name}"
            raise ValueError(msg)
        if len(datasets) > 1:
            raise Exception
        return datasets[0]

    def is_user_input(self, field: str) -> bool:
        if (spec := self.processing_spec.get(field)) is not None:
            return not spec.no_input
        # fields without a spec may be `extraInputFields` and therefore still user input
        return not field.startswith(INTERNAL_INPUT_PREFIXES)


def set_sequence_name(
    multi_reference: bool, multi_segment: bool, ds: NextcladeSequenceAndDataset
) -> str:
    match (multi_reference, multi_segment):
        case (False, _):
            return ds.segment
        case (True, True):
            return f"{ds.segment}-{ds.reference_name}"
        case (True, False):
            return ds.reference_name
        case _:
            msg = "Internal Error - unreachable code reached"
            raise AssertionError(msg)


def base_type(field_type: Any) -> type:
    """Pull the non-None type from a Union, e.g. `str | None` -> `str`"""
    if type(field_type) is UnionType:
        return next(t for t in get_args(field_type) if t is not type(None))
    return field_type


def kebab(s: str) -> str:
    """Convert snake_case to kebab-case"""
    return s.replace("_", "-")


def generate_argparse_from_model(config_cls: type[BaseModel]) -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Command-line arguments for Config class")

    parser.add_argument("--config-file")

    for name, field in config_cls.model_fields.items():
        field_type = base_type(field.annotation)
        if field_type not in CLI_TYPES:
            continue

        arg_name = f"--{kebab(name)}"

        if field_type is bool:
            parser.add_argument(arg_name, action=argparse.BooleanOptionalAction)
        else:
            # no default here -> stays None if not provided
            parser.add_argument(arg_name, type=field_type)

    return parser


def validate_required_when(config: Config) -> None:
    """Validate every `required_when` condition in the processing spec."""
    for output_field, spec in config.processing_spec.items():
        if spec.required and spec.required_when:
            msg = (
                f"invalid configuration: field '{output_field}' sets both 'required: true' and "
                "'requiredWhen'. This is not allowed."
            )
            raise ValueError(msg)
        for condition in spec.required_when:
            if condition.startswith(FILES_PREFIX):
                category = condition.removeprefix(FILES_PREFIX)
                if category not in FileCategory:
                    msg = (
                        f"invalid configuration: field '{output_field}' has a requiredWhen "
                        f"condition referencing non-existing file category '{category}'."
                    )
                    raise ValueError(msg)
                continue

            if condition.removeprefix(PROCESSED_PREFIX) == output_field:
                msg = (
                    f"invalid configuration: field '{output_field}' lists itself in `requiredWhen`"
                )
                raise ValueError(msg)


def get_processing_order(config: Config) -> tuple[str, ...]:
    """Return a valid order for processing metadata fields based on their dependencies.

    Dependencies are derived from input fields and requiredWhen fields in `config.processing_spec`.

    A DAG is constructed and topologically sorted to ensure each field is processed after the
    fields it depends on.

    A metadatafield can declare a dependency on a processed metadatafield by prefixing
    the name of the required metadatafield with `processed.` in its inputs.
    E.g.: `processed.collection_date` will look for `collection_date` in the processed metadata,
    whereas `collection_date` will use the unprocessed version
    """
    dag: dict[str, set[str]] = {k: set() for k in config.processing_spec}
    for output_field, spec in config.processing_spec.items():
        for input in itertools.chain(spec.inputs.values(), spec.required_when):
            if not input.startswith(PROCESSED_PREFIX):
                continue
            dependency = input.removeprefix(PROCESSED_PREFIX)

            if dependency not in config.processing_spec:
                msg = (
                    f"invalid configuration: metadata field '{output_field}' requested "
                    f"non-existing field '{dependency}' as input or as a requiredWhen condition."
                )
                raise ValueError(msg)
            dag[output_field].add(dependency)

    ts = graphlib.TopologicalSorter(dag)
    try:
        processing_order = tuple(ts.static_order())
    except graphlib.CycleError as e:
        msg = (
            f"invalid configuration: computation of metadata processing order resulted in a "
            f"Cycle error: {e}"
        )
        raise ValueError(msg) from e

    return processing_order


def get_config(config_file: str | None = None, ignore_args: bool = False) -> Config:
    """
    Config precedence: Direct function args > CLI args > ENV variables > config file > default

    args:
        config_file: Path to YAML config file - only used by tests
    """

    # Set just log level this early from env, so we can debug log during config loading
    env_log_level = os.environ.get("PREPROCESSING_LOG_LEVEL")
    if env_log_level:
        logging.basicConfig(level=env_log_level)

    if not ignore_args:
        parser = generate_argparse_from_model(Config)
        args = parser.parse_args()
    else:
        args = argparse.Namespace()

    # Use first config file present in order of precedence
    config_file_path = (
        config_file or args.config_file or os.environ.get("PREPROCESSING_CONFIG_FILE")
    )

    # Start with lowest precedence config, then overwrite with higher precedence
    if config_file_path:
        with open(config_file_path, encoding="utf-8") as file:
            yaml_config = yaml.safe_load(file)
            logger.debug(f"Loaded config from {config_file_path}: {yaml_config}")
        config = Config(**yaml_config)
    else:
        config = Config()
    # Use environment variables if available
    env_overrides = {}
    for key in Config.model_fields:
        env_var = f"PREPROCESSING_{key.upper()}"
        if env_var in os.environ:
            env_overrides[key] = os.environ[env_var]

    if env_overrides:
        config = Config(**{**config.model_dump(), **env_overrides})

    # Overwrite config with CLI args
    cli_overrides = {
        k: v for k, v in args.__dict__.items() if v is not None and k in Config.model_fields
    }
    if cli_overrides:
        config = Config(**{**config.model_dump(), **cli_overrides})

    return config
