from __future__ import annotations

from typing import Optional

from rdkit import Chem as rdk

from chemrecon.chem.gml.gml_to_rdk import gml_to_rdkit_mol
from chemrecon.chem.mol import MolTemplate, MolInstance, feat_enum_map
from chemrecon.core.id_types import S_SMILES, S_INCHI, S_GML, S_MOLFILE
from chemrecon.chem.mol import Mol
from chemrecon.schema.entry_types.molstructure import MolStructure
from chemrecon.schema.entry_types.molstructure_repr import MolStructureRepr


def mol_from_smiles(smiles: str, provenance: Optional[str] = None, template: bool = False, safe = False) -> Mol:
    if template:
        mol = MolTemplate(
            rdk.MolFromSmiles(smiles),
            provenance = provenance
        )
    else:
        mol = MolInstance(
            rdk.MolFromSmiles(smiles),
            provenance = provenance
        )
    return mol


def mol_from_inchi(inchi: str, provenance: Optional[str] = None, template: bool = False, safe = False) -> Mol:
    if template:
        mol = MolTemplate(rdk.MolFromInchi(inchi), provenance = provenance)
    else:
        mol = MolInstance(rdk.MolFromInchi(inchi), provenance = provenance)
    return mol


def mol_from_molblock(molblock: str, provenance: Optional[str] = None, template: bool = False, safe = False) -> Mol:
    if template:
        mol = MolTemplate(
            rdk.MolFromMolBlock(molblock, sanitize = not safe, removeHs = not safe),
            provenance = provenance
        )
    else:
        mol = MolInstance(
            rdk.MolFromMolBlock(molblock, sanitize = not safe, removeHs = not safe),
            provenance = provenance
        )
    return mol


def mol_from_gml(gml: str, provenance: Optional[str] = None, template: bool = False) -> Mol:
    rdk_mol: rdk.Mol = gml_to_rdkit_mol(gml)
    mol = mol_from_rdk_mol(rdk_mol, provenance, template)
    return mol


def mol_from_rdk_mol(rdk_mol: rdk.Mol, provenance: Optional[str] = None, template: bool = False) -> Mol:
    if template:
        mol = MolTemplate(rdk_mol, provenance = provenance)
    else:
        mol = MolInstance(rdk_mol, provenance = provenance)
    return mol


def mol_from_structurerepresentation(
    entry: MolStructureRepr,
    provenance: Optional[str] = None,
    template: bool = False,
    safe: bool = False
):
    match entry.id_type:
        case S_SMILES.enum_type:
            return mol_from_smiles(entry.source_id, provenance, template = template, safe = safe)
        case S_INCHI.enum_type:
            return mol_from_inchi(entry.source_id, provenance, template = template, safe = safe)
        case S_MOLFILE.enum_type:
            return mol_from_molblock(entry.source_id, provenance, template = template, safe = safe)
        case S_GML.enum_type:
            return mol_from_gml(entry.source_id, provenance, template = template)
        case _:
            raise ValueError


def mol_from_struct_entry(entry: MolStructure):
    mol = Mol(
        rdk.MolFromSmiles(entry.smiles),
        set_features ={
            feat_enum_map[f] for f in entry.std_feats
        }
    )
    return mol
