from __future__ import annotations

from abc import ABC, abstractmethod
from enum import IntEnum
from typing import TYPE_CHECKING, ClassVar, Optional

import chemrecon
from chemrecon.chem.mol import Mol, MolTemplate
from chemrecon.chem.sumformula import SumFormula

if TYPE_CHECKING:
    from aam_annotator.mapped_network import MappedMetabolite, MappedReaction
    from aam_annotator.matching import Matching


class Severity(IntEnum):
    UNKNOWN = 10
    INFO = 20
    WARNING = 30
    ERROR = 40


UNKNOWN = Severity.UNKNOWN
INFO = Severity.INFO
WARNING = Severity.WARNING
ERROR = Severity.ERROR


class Problem(Exception, ABC):
    """Generic base class for 'problems' identified by the algorithm, which are passed to the user so they
    can provide fixes in the form of curations.
    """

    severity: ClassVar[Severity] = UNKNOWN

    @abstractmethod
    def message(self) -> str:
        """What is printed to the JSON file."""


# Curation problems
# --------------------------------------------------------------------------------------------------------------
class CurationProblem(Problem, ABC):
    def __init__(self):
        pass


class CuratedStructureNotFound(CurationProblem):
    severity = WARNING
    id: str
    mol: MolTemplate

    def __init__(self, id: str, mol: MolTemplate):
        self.id = id
        self.mol = mol

    def message(self) -> str:
        return f"Curated structure not found: {self.mol.to_smiles()}"


class CuratedAAMNotFound(CurationProblem):
    severity = WARNING
    id: str
    aam: chemrecon.ChemReaction

    def __init__(self, id: str, aam: chemrecon.ChemReaction):
        self.id = id
        self.aam = aam

    def message(self) -> str:
        return f"Curated AAM not found: {self.id}"


# Metabolite problems
# --------------------------------------------------------------------------------------------------------------
class MetaboliteProblem(Problem, ABC):
    metabolite: MappedMetabolite

    def __init__(self, m: MappedMetabolite):
        self.metabolite = m


class NoStructuresFound(MetaboliteProblem):
    severity = WARNING
    eg: chemrecon.EntryGraph

    def __init__(self, m: MappedMetabolite, eg: chemrecon.EntryGraph):
        super().__init__(m)
        self.eg = eg

    def message(self) -> str:
        return (
            f"No enrygraph structures for compound {self.metabolite}."
            f"Entry graph contains {sum(1 for _ in self.eg.vertices())} entries."
        )


class IncompatibleStructuresSelectedWarning(MetaboliteProblem):
    severity = WARNING
    representative_structure: Mol
    incompatible_structures: list[Mol]

    def __init__(
        self,
        m: MappedMetabolite,
        representative_structure: Mol,
        structures: list[Mol],
    ):
        super().__init__(m)
        self.representative_structure = representative_structure
        self.structures = structures

    def message(self) -> str:
        return (
            f"Incompatible structures selected for {self.metabolite}. "
            f"Representative structure: {self.representative_structure}. "
            f"Conflicting structures: {self.structures}"
        )


class CouldNotGenerateMapToCanonicalStructure(MetaboliteProblem):
    severity = WARNING
    rdt_query: str

    def __init__(self, m: MappedMetabolite, rdt_query: str):
        super().__init__(m)
        self.rdt_query = rdt_query

    def message(self) -> str:
        return f"Could not generate map to canonical structure for {self.metabolite} using RDT query: {self.rdt_query}"


# Reaction problems
# --------------------------------------------------------------------------------------------------------------
class ReactionProblem(Problem, ABC):
    reaction: MappedReaction

    def __init__(self, r: MappedReaction):
        self.reaction = r


class NoAAMsFound(ReactionProblem):
    severity = WARNING
    eg: chemrecon.EntryGraph

    def __init__(self, r: MappedReaction, eg: chemrecon.EntryGraph):
        super().__init__(r)
        self.eg = eg

    def message(self) -> str:
        return f"No AAMs found for reaction {self.reaction}."


# AAM generation errors


class RDTError(ReactionProblem):
    severity = ERROR
    rdt_query: str

    def __init__(self, r: MappedReaction, rdt_query: str):
        super().__init__(r)
        self.rdt_query = rdt_query

    def message(self) -> str:
        return f"RDT error on query: {self.rdt_query}"


class RDTTimeout(ReactionProblem):
    severity = ERROR
    rdt_query: str

    def __init__(self, r: MappedReaction, rdt_query: str):
        super().__init__(r)
        self.rdt_query = rdt_query

    def message(self) -> str:
        return f"RDT timeout on query: {self.rdt_query}"


class CannotGenerateAAMNoStructures(ReactionProblem):
    severity = ERROR
    metabolites_missing_structures: list[MappedMetabolite]

    def __init__(
        self,
        reaction: MappedReaction,
        metabolites_missing_structures: list[MappedMetabolite],
    ):
        super().__init__(reaction)
        self.metabolites_missing_structures = metabolites_missing_structures

    def message(self) -> str:
        return f"Cannot generate AAM for reaction {self.reaction.id} due to missing structures for metabolites: {[m.id for m in self.metabolites_missing_structures]}"


class MissingStructuresWarning(ReactionProblem):
    severity = WARNING
    metabolites_missing_structures: list[MappedMetabolite]

    def __init__(
        self,
        reaction: MappedReaction,
        metabolites_missing_structures: list[MappedMetabolite],
    ):
        super().__init__(reaction)
        self.metabolites_missing_structures = metabolites_missing_structures

    def message(self) -> str:
        return f"Missing structures for metabolites in reaction {self.reaction.id}: {[m.id for m in self.metabolites_missing_structures]}"


# Database AAM reactions
# ----------------------------------------------------------------------------------------------------------------------
class DatabaseAAMProblem(ReactionProblem, ABC):
    aam_entry: chemrecon.AAM
    aam: Optional[chemrecon.ChemReaction]

    def __init__(
        self,
        reaction: MappedReaction,
        aam_entry: chemrecon.AAM,
        aam: chemrecon.ChemReaction,
    ):
        super().__init__(reaction)
        self.aam_entry = aam_entry
        self.aam = aam


class AAMUnbalancedProblem(DatabaseAAMProblem):
    severity = WARNING
    difference: SumFormula

    def __init__(self, reaction, aam_entry: Optional[chemrecon.AAM], aam):
        assert aam_entry is not None
        super().__init__(reaction, aam_entry, aam)
        self.difference = aam.get_balance_difference()

    def message(self) -> str:
        return (
            f"Database AAM for {self.reaction.id} is unbalanced."
            f"AAM: {self.aam.to_reaction_smiles() if self.aam else 'Curated'}."
            f"Balance difference: {self.difference}."
        )


# Matching
# -----------------------------------------------
class MatchingProblem(Problem, ABC):
    def __init__(self):
        pass


class UnmatchedCuratedStructure(MatchingProblem):
    severity = WARNING

    metabolite: MappedMetabolite
    curated_mol: MolTemplate

    def __init__(self, metabolite: MappedMetabolite, curated_mol: MolTemplate):
        super().__init__()
        self.metabolite = metabolite
        self.curated_mol = curated_mol

    def message(self) -> str:
        return (
            f"Curated structure for {self.metabolite} does not match any structure in the AAM."
            f"Curated structure: {self.curated_mol.to_smiles()}"
        )


class CuratedStructureNotInAAM(MatchingProblem):
    severity = WARNING
    curated_metabolite: MappedMetabolite
    curated_mol: MolTemplate

    def __init__(
        self, curationed_metabolite: MappedMetabolite, curated_mol: MolTemplate
    ):
        super().__init__()
        self.curated_metabolite = curationed_metabolite
        self.curated_mol = curated_mol

    def message(self) -> str:
        return (
            f"Curated structure for {self.curated_metabolite} does not match any structure in the AAM."
            f"Curated structure: {self.curated_mol.to_smiles()}"
        )


class MultiplicityMismatch(MatchingProblem):
    severity = WARNING
    metabolite: MappedMetabolite
    aam_mol: MolTemplate

    def __init__(self, metabolite: MappedMetabolite, aam_mol: MolTemplate):
        self.metabolite = metabolite
        self.aam_mol = aam_mol

    def message(self) -> str:
        return f"Multiplicity mismatch for metabolite {self.metabolite} and structure {self.aam_mol}."


class IncompleteMatchingProblem(MatchingProblem):
    severity = WARNING

    def message(self) -> str:
        # TODO print unmatched items
        return "Incomplete matching: "


class LowConfidenceMatchingWarning(MatchingProblem):
    severity = WARNING
    conf: float

    def __init__(self, conf: float):
        self.conf = conf

    def message(self) -> str:
        return f"Low confidence matching: {self.conf:.2f}"


class AnotherMatchingHasHigherConfidence(MatchingProblem):
    severity = WARNING
    max_conf_matching: Matching

    def __init__(self, max_conf_matching: Matching):
        self.max_conf_matching = max_conf_matching

    def message(self) -> str:
        return f"Another matching has higher confidence: {self.max_conf_matching.get_conf():.2f}"


# Exceptions
# ---------------------------------------------------------
class CannotMatch(MatchingProblem):
    severity = ERROR

    def message(self) -> str:
        return "Cannot match."


class CannotMatchNoMetabolites(CannotMatch):
    severity = ERROR
    pass


class CannotMatchNoStructures(CannotMatch):
    severity = ERROR
    pass
