""" Functions for finding entries given an index.
"""
from typing import Optional

import chemrecon.connection as connection

from chemrecon import IdentifierType, Entry, IdentifierTypeCompound, IdentifierTypeReaction, IdentifierTypeEnzyme, \
    IdentifierTypeStructureRepresentation, IdentifierTypeAAM
from chemrecon.schema.entry_types.aam import AAM
from chemrecon.schema.entry_types.aam_repr import AAMRepr
from chemrecon.schema.entry_types.compound import Compound
from chemrecon.schema.entry_types.enzyme import Enzyme
from chemrecon.schema.entry_types.reaction import Reaction
from chemrecon.schema.entry_types.molstructure_repr import MolStructureRepr
from chemrecon.schema.entry_types.molstructure import MolStructure


def find_entry(id_type: IdentifierType, source_id: str) -> Optional[Entry]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    match id_type:
        case IdentifierTypeCompound():
            return find_compound_entry(id_type, source_id)
        case IdentifierTypeReaction():
            return find_reaction_entry(id_type, source_id)
        case IdentifierTypeEnzyme():
            return find_enzyme_entry(id_type, source_id)
        case IdentifierTypeStructureRepresentation():
            return find_structure_representation_entry(id_type, source_id)
        case IdentifierTypeAAM():
            return find_aam_representation_entry(id_type, source_id)
        case _:
            # Not implemented
            raise NotImplementedError()

def find_compound_entry(id_type: IdentifierTypeCompound, source_id: str) -> Optional[Compound]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = Compound(id_type = id_type.enum_type, source_id = id_type.std_identifier(source_id))
    return connection.handler.get_entry_by_index(e)

def find_reaction_entry(id_type: IdentifierTypeReaction, source_id: str) -> Optional[Reaction]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = Reaction(id_type = id_type.enum_type, source_id = id_type.std_identifier(source_id))
    return connection.handler.get_entry_by_index(e)

def find_enzyme_entry(id_type: IdentifierTypeEnzyme, source_id: str) -> Optional[Enzyme]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = Enzyme(id_type = id_type.enum_type, source_id = id_type.std_identifier(source_id))
    return connection.handler.get_entry_by_index(e)

def find_structure_representation_entry(id_type: IdentifierTypeStructureRepresentation, source_id: str) -> Optional[MolStructureRepr]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = MolStructureRepr(id_type = id_type.enum_type, source_id = id_type.std_identifier(source_id), implicit = False)
    return connection.handler.get_entry_by_index(e)

def find_aam_representation_entry(id_type: IdentifierTypeAAM, source_id: str) -> Optional[AAMRepr]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = AAMRepr(id_type = id_type.enum_type, source_id = id_type.std_identifier(source_id))
    return connection.handler.get_entry_by_index(e)

def find_structure_entry(smiles: str) -> Optional[MolStructure]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = MolStructure(smiles = smiles)
    return connection.handler.get_entry_by_index(e)

def find_aam_entry(reaction_smiles: str) -> Optional[AAM]:
    """ Look for an entry with the specified type and id in the connected Database.
        If not found, returns None.
    """
    e = AAM(reaction_smiles = reaction_smiles)
    return connection.handler.get_entry_by_index(e)
