import json
from curses import meta
from typing import Any, Optional

import chemrecon
from chemrecon.chem.mol import Mol, MolInstance, MolTemplate
from mntools import Side

from aam_annotator.curations import Curations
from aam_annotator.mapped_network import (
    MappedMetabolite,
    MappedNetwork,
    MappedReaction,
    MappedSpecies,
)
from aam_annotator.matching import Matching
from aam_annotator.problem import ERROR
from aam_annotator.provenance import ProvenanceStruct


def serialize_network(network: MappedNetwork, curations: Curations) -> str:

    out: dict[str, Any] = dict()

    # Overall statistics
    # ----------------------------------------------------------------------------------------------------------
    out["stats"] = network.stats()

    # Curations
    # ----------------------------------------------------------------------------------------------------------
    curations_metabolites_dict: dict[str, str] = dict()
    unmatched_curations_metabolites: set[str] = set()
    for metabolite_id, mol in curations.structures.items():
        curations_metabolites_dict[metabolite_id] = mol.to_smiles()
        if metabolite_id not in curations.structure_matchings:
            unmatched_curations_metabolites.add(metabolite_id)

    curations_reactions_dict: dict[str, str] = dict()
    unmatched_curations_reactions: set[str] = set()
    for reaction_id, aam in curations.aams.items():
        curations_reactions_dict[reaction_id] = aam.to_reaction_smiles()
        if reaction_id not in curations.reaction_matchings:
            unmatched_curations_reactions.add(reaction_id)

    curated_hydrogen: list[str] = list()
    for metabolite_id in curations.hydrogen_metabolites:
        curated_hydrogen.append(metabolite_id)

    # Collect curations
    curations_dict: dict[str, Any] = {
        "hydrogen metabolites": curated_hydrogen,
        "curated_metabolites": curations_metabolites_dict,
        "curated_reactions": curations_reactions_dict,
    }

    # Un-matched curations, if any.
    if len(unmatched_curations_metabolites) > 0:
        curations_dict["unmatched_metabolites"] = [
            m_id for m_id in unmatched_curations_metabolites
        ]
    if len(unmatched_curations_reactions) > 0:
        curations_dict["unmatched_reactions"] = [
            r_id for r_id in unmatched_curations_reactions
        ]

    # Done.
    out["curations"] = curations_dict

    # Metabolites
    # ----------------------------------------------------------------------------------------------------------
    metabolites_dict: dict[str, dict[str, Any]] = dict()
    for m in network.all_metabolites():
        metabolites_dict[m.id] = serialize_metabolite(m)
    out["metabolites"] = metabolites_dict

    # Reactions
    # ----------------------------------------------------------------------------------------------------------
    reactions_dict: dict[str, dict[str, Any]] = dict()
    for r in network.all_reactions():
        reactions_dict[r.id] = serialize_reaction(r)
    out["reactions"] = reactions_dict

    # Problems
    # ----------------------------------------------------------------------------------------------------------
    # TODO

    # Missing information
    # ----------------------------------------------------------------------------------------------------------
    # TODO metabolites missing structures
    # TODO reactions missing aams

    # Finalize
    # ----------------------------------------------------------------------------------------------------------
    json_str = json.dumps(out, indent=2)
    return json_str


def serialize_metabolite(metabolite: MappedMetabolite) -> dict[str, Any]:
    out: dict[str, Any] = dict()

    out["id"] = metabolite.id
    out["name"] = metabolite.name

    # Species associated with this metabolite (by compartment)
    out["sbml_species"] = {sp.compartment.id: sp.id for sp in metabolite.species()}

    # Biomass
    if metabolite.is_biomass:
        out["biomass"] = True
        return out

    if metabolite.is_hydrogen:
        out["is_hydrogen_only"] = True

    # Selected
    primary_structure: Optional[MolTemplate] = None
    structure_source: Optional[str] = None

    if metabolite.selected_mol is not None:
        assert metabolite.selection_provenance is not None
        primary_structure = metabolite.selected_mol
        structure_source = metabolite.selection_provenance.serialize()

    # Add selected structure
    if primary_structure is not None:
        out["primary_structure"] = primary_structure.to_smiles()
        out["primary_structure_source"] = structure_source
        out["primary_structure_std_features"] = primary_structure.feature_string()
    else:
        out["primary_structure"] = None

    # Other structures
    if metabolite.candidate_structures_aam_selected is not None:
        if len(metabolite.candidate_structures_aam_selected) > 1:
            alt_structs_dict = dict()
            for alt_struct in metabolite.candidate_structures_aam_selected:
                if alt_struct == metabolite.selected_mol:
                    continue
                # Add to dict
                struct_dict: dict[str, Any] = dict()
                struct_dict["smiles"] = alt_struct.to_smiles()
                struct_dict["std_features"] = alt_struct.feature_string()
                try:
                    struct_dict["map_to_canonical"] = metabolite.canonical_maps[
                        alt_struct
                    ].to_reaction_smiles()
                except KeyError:
                    struct_dict["map_to_canonical"] = "CANONICAL MAP ERROR."
                struct_dict["in reactions"] = [
                    reaction.id
                    for reaction, matching in metabolite.candidate_structures_aam_selected[
                        alt_struct
                    ]
                ]
                # TODO add any prolbmes

                alt_structs_dict[alt_struct.smiles] = struct_dict

            # Add to output
            out["alternative_structures"] = alt_structs_dict

    # Problems
    problems_list: list[str] = list()
    for problem in metabolite.problems:
        problems_list.append(problem.message())

    if len(problems_list) > 0:
        out["problems"] = problems_list

    # Done.
    return out


def serialize_reaction(reaction: MappedReaction) -> dict[str, Any]:
    out: dict[str, Any] = dict()

    # Basic information
    out["id"] = reaction.id
    out["name"] = reaction.name

    # If biomass, skip
    if reaction.is_biomass_reaction:
        out["biomass_reaction"] = True
        return out

    # Network
    out["is_transport"] = reaction.is_transport_reaction
    out["is_transport_1to1"] = reaction.is_transport_reaction_1to1
    compartment = reaction.get_compartment()
    if compartment is not None:
        out["compartment"] = compartment.id

    # Stoichiometry
    out["lhs"] = {sp.id: stoich for sp, stoich in reaction.get_lhs().items()}
    out["rhs"] = {sp.id: stoich for sp, stoich in reaction.get_rhs().items()}

    # Matching
    if reaction.selected_aam_matching is not None:
        out["selected_matching"] = serialize_matching(reaction.selected_aam_matching)
        assert reaction.selection_provenance is not None
        out["matching_source"] = reaction.selection_provenance.serialize()

    # Problems
    errors_list: list[str] = list()
    warnings_list: list[str] = list()
    for problem in reaction.problems:
        if problem.severity == ERROR:
            errors_list.append(problem.message())
        else:
            warnings_list.append(problem.message())

    if len(errors_list) > 0:
        out["errors"] = errors_list
    if len(warnings_list) > 0:
        out["warnings"] = warnings_list

    return out


def serialize_matching(matching: Matching) -> dict[str, Any]:
    out: dict[str, Any] = dict()

    # AAM
    out["aam"] = matching.aam.to_reaction_smiles()

    # Matching itself
    match: dict[Side, list] = dict()
    for side in Side:
        matchlist = list()
        for metabolite, mol in matching.match[side].items():
            matchlist.append(
                {
                    "id": metabolite.id,
                    "smiles": mol.to_smiles(),
                }
            )
        match[side] = matchlist

    out["lhs"] = match[Side.L]
    out["rhs"] = match[Side.R]

    # Other
    out["conf"] = matching.get_conf()

    # Unmatched
    # TODO warn

    # Mapping to non-primary structure
    # TODO warn

    return out
