"""Defines curations used in the algorithm.

There are multiple types:
 - Structure curations (input by the user, possibly ahead of time - correct smiles for given primary ids)
 - Match curation (created by the program with input from the user to fix weird matchings)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from chemrecon.chem.chemreaction import chem_reaction_from_reactionsmiles

from aam_annotator.utils.hydrogen import is_hydrogen

if TYPE_CHECKING:
    from aam_annotator.mapped_network import MappedMetabolite, MappedReaction

import json
from typing import Optional

import chemrecon
from chemrecon import (
    ChemReaction,
    Mol,
)
from chemrecon.chem.mol import MolTemplate

from aam_annotator.utils.std_id import (
    standardize_metabolite_id,
    standardize_reaction_id,
)


# Curations object
# --------------------------------------------------------------------------------------------------------------
class Curations:
    """Contains curations."""

    # TODO also allow curating by name?

    structures: dict[str, MolTemplate]  # By key - should be standardized
    aams: dict[
        str, ChemReaction
    ]  # Reaction key - reaction and optional matching for MolTemplates

    # Misc
    hydrogen_metabolites: set[
        str
    ]  # List of metabolites assigned with a hydrogen formula

    # Matchings
    structure_matchings: dict[str, Optional[MappedMetabolite]]
    reaction_matchings: dict[str, Optional[MappedReaction]]

    def __init__(self):
        """Initialized as empty, should be populated by loading JSONs."""
        self.structures = dict()
        self.aams = dict()
        self.hydrogen_metabolites = set()

        self.structure_matchings = dict()
        self.reaction_matchings = dict()


# Loader
# ----------------------------------------------------------------------------------------------------------------------
def load_curations_from_directory(curations_path: str) -> Curations:
    print("Loading curation data...")
    curations = Curations()

    # Curated structures
    # ------------------------------------------------------------------------------------------------------------------
    try:
        with open(curations_path, "r") as f:
            curations_json = json.load(f)

            # Structures
            structures_json: Optional[dict[str, str]]
            try:
                structures_json = curations_json["structures"]
            except KeyError:
                print(" - No structures found in curations file!")
                structures_json = None

            if structures_json is not None:
                for st_id, st_smiles in structures_json.items():
                    std_id = standardize_metabolite_id(st_id)
                    try:
                        mol = chemrecon.mol_from_smiles(st_smiles, template=True)
                        assert isinstance(mol, MolTemplate)

                        # Set curation
                        curations.structures[std_id] = mol

                        # Check whether the mol is hydrogen
                        if is_hydrogen(mol):
                            curations.hydrogen_metabolites.add(std_id)

                    except ValueError:
                        # TODO add curation  error
                        print(f" - Could not parse curated SMILES for {st_id}!")
                        continue

            # AAMs
            aams_json: Optional[dict[str, str]]
            try:
                aams_json = curations_json["aams"]
            except KeyError:
                print(" - No AAMs found in curations file!")
                aams_json = None

            if aams_json is not None:
                for st_id, st_aam in aams_json.items():
                    st_id = standardize_reaction_id(st_id)
                    try:
                        # Set curation
                        aam = chem_reaction_from_reactionsmiles(st_aam)
                        curations.aams[st_id] = aam

                    except ValueError:
                        # TODO add curation error
                        print(f" - Could not parse curated AAM for {st_id}!")
                        continue

    except FileNotFoundError:
        print(" - Could not open curations file.")

    # Finalize
    return curations
