from __future__ import annotations

from typing import Callable
import re

# Lookup lists and dictionaries
# ----------------------------------------------------------------------------------------------------------------------
id_types: list[IdentifierType] = list()
id_type_name_lookup: dict[str, IdentifierType] = dict()
identifiers_org_dict: dict[str, IdentifierType] = dict()


# Types of ID types
# ----------------------------------------------------------------------------------------------------------------------
class IdentifierType:
    """
    Represents a general type of identifier with attributes and methods for standardization,
    recognition, and manipulation.

    This class is designed to encapsulate information about a specific type of identifier,
    including its primary name, alternative names, recognizable patterns, and
    standardization logic.

    The class also registers identifier types in global lookup dictionaries
    to facilitate type recognition and access.
    """
    name: str  #: Primary name of this identifier type.
    shortname: str  #: The name used for the type in the database.
    alt_names: set[str]  #: Alternative names to search for.
    prefixes: set[str]
    suffixes: set[str]
    enum_type: 'IdType'  #: The corresponding Enum value, as present in entries
    stdfunc: Callable[[str], str]  #: Function used to standardize identifiers.
    recogniser: re.Pattern | None  #: Pattern used to recognize identifiers of this type.
    id_org_prefix: str  #: Prefix in identifiers.org
    objectname: str  # name of the assigned object

    def __init__(
            self,
            name: str,
            shortname: str,
            alt_names: set[str] = None,
            prefixes: set[str] = None,
            suffixes: set[str] = None,
            stdfunc: Callable[[str], str] = None,
            recogniser: re.Pattern | None = None,
            objectname: str = None
    ):
        self.name = name
        self.shortname = shortname
        self.alt_names = alt_names
        self.prefixes = prefixes if prefixes else {}
        self.suffixes = suffixes if suffixes else {}
        self.stdfunc = stdfunc
        self.recogniser = recogniser
        self.objectname = objectname

        # Make identifiers org link
        # TODO doesn't always work (mnx?)
        self.identifiers_org_prefix = ''
        for prefix in self.prefixes:
            if prefix.startswith('https://identifiers.org/') or prefix.startswith('http://identifiers.org/'):
                identifiers_org_dict[prefix] = self
                self.identifiers_org_prefix = prefix

        # Register in lookup
        id_types.append(self)

        # Register names
        id_type_name_lookup[self.name] = self
        id_type_name_lookup[self.shortname] = self
        for alt_name in self.alt_names:
            id_type_name_lookup[alt_name] = self

    def __repr__(self):
        return self.objectname

    def __str__(self):
        return self.name

    def __hash__(self):
        return self.shortname.__hash__()

    def std_identifier(self, s: str) -> str:
        """ Standardize a given identifier of this type. """
        if self.stdfunc:
            return self.stdfunc(s)
        else:
            return str(s)

    def trim(self, s: str) -> str:
        """ Remove pre- and suffixes of the string, including identifiers.org urls. """
        s = s.strip()
        for prefix in self.prefixes:
            s = s.removeprefix(prefix)
        for suffix in self.suffixes:
            s = s.removesuffix(suffix)
        return s


# Register subtypes
class IdentifierTypeCompound(IdentifierType):
    enum_type: 'IdTypeCompoundEnum'


class IdentifierTypeStructureRepresentation(IdentifierType):
    enum_type: 'IdTypeStructureRepresentationEnum'


class IdentifierTypeReaction(IdentifierType):
    enum_type: 'IdTypeReactionEnum'


class IdentifierTypeEnzyme(IdentifierType):
    enum_type: 'IdTypeEnzymeEnum'


class IdentifierTypeAAM(IdentifierType):
    enum_type: 'IdTypeAAMEnum'


# Standardisation functions
# ----------------------------------------------------------------------------------------------------------------------
# BiGG
_compartment_suffixes = [
    '_e', '_ex', '_c', '_p', '_m', '_x', '_b',
    '_E', '_EX', '_C', '_P', '_M', '_X', '_B'
]


def _std_bigg(s: str) -> str:
    """ Create a universal BIGG id without the compartment suffix"""
    # TODO double underscore is necessary (indicates stereo - "phe__L" is L-Phenylalanine)?
    s_ = s.strip()
    for suffix in _compartment_suffixes:
        s_ = s_.removesuffix(suffix)
    s_ = s_.replace('_DASH', '')  # Fix weird 'DASH' notation, i.e. M_12ppd_DASH_S -> M_12ppd__S
    s_ = s_.replace('__', '_')  # Hacky fix for double underscore notation
    s_ = s_.removeprefix('m_')
    if not s_.startswith('M_'):
        s_ = f'M_{s_}'

    if s_.isupper():
        # Fix upper case
        s_ = f'M_{s_.removeprefix("M_").lower()}'

    return s_


def _std_bigg_r(s: str) -> str:
    s_ = s.strip()
    for suffix in _compartment_suffixes:
        s_ = s_.removesuffix(suffix)
    s_ = s_.replace('_DASH', '').replace('__', '_').removeprefix('r_')
    if not s_.startswith('R_'):
        s_ = f'R_{s_}'
    if s_.isupper():
        # Fix upper case
        s_ = f'R_{s_.removeprefix("M_").lower()}'

    # Replace LPAREN, RPAREN notation in names
    s_ = s_.replace('LPAREN_', '_').replace('RPAREN_', '_')

    return s_


def _std_chebi(chebi: str) -> str:
    """ Chebi IDs should be prefixed with CHEBI: """
    if not chebi:
        raise ValueError(f'Received empty CHEBI')
    chebi_new = chebi.strip()
    if chebi_new.startswith('http://identifiers.org/chebi/CHEBI'):
        return chebi_new.removeprefix('http://identifiers.org/chebi/')
    if chebi_new.startswith('https://identifiers.org/chebi/CHEBI'):
        return chebi_new.removeprefix('https://identifiers.org/chebi/')

    if chebi_new.startswith('CHEBI:'):
        return chebi_new
    else:
        return f'CHEBI:{chebi_new}'


# C_NAME
def _standardise_cname(cname: str) -> str:
    """ Lowercase, etc """
    return cname.lower().strip()


# Compound Identifier Types
# ---------------------------------------------------------------------------------------------------------------------

# Unknown
C_UNKNOWN = IdentifierTypeCompound(
    name = 'Unknown Compound',
    shortname = 'unknown_c',
    alt_names = set(),
    prefixes = set(),
    objectname = 'C_UNKNOWN'
)

# Name
C_NAME = IdentifierTypeCompound(
    name = 'Compound Name',
    shortname = 'cname',
    alt_names = set(),
    prefixes = set(),
    objectname = 'C_NAME',
)

# MetaNetX
C_MNX = IdentifierTypeCompound(
    name = 'MetaNetX',
    shortname = 'mnx',
    alt_names = {'metanetx', 'mnx'},
    prefixes = {
        'http://identifiers.org/metanetx.chemical/',
        'https://identifiers.org/metanetx.chemical/',
    },
    recogniser = re.compile(r'^(MNXM\d+|MNX\d+|BIOMASS|WATER)$'),
    objectname = 'C_MNX',
)

# BiGG
C_BIGG = IdentifierTypeCompound(
    name = 'BiGG',
    shortname = 'bigg',
    alt_names = {'bigg', 'biggM', 'bigg.metabolite'},
    prefixes = {
        'http://bigg.ucsd.edu/models/universal/metabolites/',
        'http://identifiers.org/bigg.metabolite/',
        'https://identifiers.org/bigg.metabolite/',
        'bigg.metabolite:'
    },
    suffixes = {
        '_e', '_ex', '_c', '_p', '_m', '_x', '_b'
    },
    stdfunc = _std_bigg,
    recogniser = re.compile(r'^[a-z_A-Z0-9]+$'),
    objectname = 'C_BIGG',
)  #: asd

# PubChem CID
C_PUBCHEM = IdentifierTypeCompound(
    name = 'PubChem CiD',
    shortname = 'pubchem_cid',
    alt_names = {'pubchem', 'pc_cid', 'cid'},
    prefixes = {
        'http://identifiers.org/pubchem.compound/',
        'https://identifiers.org/pubchem.compound/'
    },
    recogniser = re.compile(r'^\d+$'),
    objectname = 'C_PUBCHEM',
)

# KEGG Compound
C_KEGG = IdentifierTypeCompound(
    name = 'KEGG',
    shortname = 'kegg',
    alt_names = {'kegg', 'keggC', 'kegg.compound', 'KEGG COMPOUND', 'KEGG', 'KEGG COMPOUND accession'},
    prefixes = {
        'http://identifiers.org/kegg.compound/',
        'https://identifiers.org/kegg.compound/'
    },
    recogniser = re.compile(r'^C\d+$'),
    objectname = 'C_KEGG',
)

# ChEBI
C_CHEBI = IdentifierTypeCompound(
    name = 'ChEBI',
    shortname = 'chebi',
    alt_names = {'chebi', 'CHEBI', 'ChEBI'},
    prefixes = {
        'http://identifiers.org/chebi/',
        'https://identifiers.org/chebi/'
    },
    recogniser = re.compile(r'^CHEBI:\d+$'),
    stdfunc = _std_chebi,
    objectname = 'C_CHEBI',
)

# ECMDB
C_ECMDB = IdentifierTypeCompound(
    name = 'ECMDB',
    shortname = 'ecmdb',
    alt_names = {'ecmdb'},
    recogniser = re.compile(r'^ECMDB\d+$'),
    objectname = 'C_ECMDB',
)

# InChI Key
C_INCHIKEY = IdentifierTypeCompound(
    name = 'InChI key',
    shortname = 'inchi_key',
    alt_names = {'inchi_key', 'inchikey'},
    prefixes = {
        'http://identifiers.org/inchikey/',
        'https://identifiers.org/inchikey/'
    },
    recogniser = re.compile(r'^[A-Z]{14}-[A-Z]{10}(-[A-Z])?$'),
    objectname = 'C_INCHIKEY',
)

# SLM
C_SLM = IdentifierTypeCompound(
    name = 'SwissLipids',
    shortname = 'slm',
    alt_names = {'slm', 'SLM'},
    prefixes = {
        'http://identifiers.org/slm/',
        'https://identifiers.org/slm/'
    },
    recogniser = re.compile(r'^SLM:\d+$'),  # TODO check?
    objectname = 'C_SLM',
)

# LipidMaps
C_LIPIDMAPS = IdentifierTypeCompound(
    name = 'LipidMaps',
    shortname = 'lipidmapsm',
    alt_names = {'LipidMapsM', 'lipidmapsM', 'lipidmaps', 'LIPID MAPS'},
    prefixes = {
        'http://identifiers.org/lipidmaps/',
        'https://identifiers.org/lipidmaps/'
    },
    recogniser = re.compile(r'^LM(FA|GL|GP|SP|ST|PR|SL|PK)[0-9]{4}([0-9a-zA-Z]{4,6})?$'),
    objectname = 'C_LIPIDMAPS',
)

# MetaCyc Compound
C_METACYC = IdentifierTypeCompound(
    name = 'MetaCyc Compound',
    shortname = 'metacyc',
    alt_names = {'metacycM', 'metacyc.compound'},
    prefixes = {
        'http://identifiers.org/metacyc.compound/',
        'https://identifiers.org/metacyc.compound/'
    },
    recogniser = re.compile(r'^[A-Za-z0-9+_.%-:]+$'),
    objectname = 'C_METACYC',
)

# EnviPath (Can be both compound and reaction?)
C_ENVIPATH = IdentifierTypeCompound(
    name = 'enviPath',
    shortname = 'envipath',
    alt_names = {'envipathM'},
    prefixes = {
        'http://identifiers.org/envipath/',
        'https://identifiers.org/envipath/'
    },
    recogniser = re.compile(
        r'^[\w^_]{8}-[\w^_]{4}-[\w^_]{4}-[\w^_]{4}-[\w^_]{12}/[\w-]+/[\w^_]{8}-'
        r'[\w^_]{4}-[\w^_]{4}-[\w^_]{4}-[\w^_]{12}$'
    ),
    objectname = 'C_ENVIPATH',
)

# Seed Compounds
C_SEED = IdentifierTypeCompound(
    name = 'Seed Compound',
    shortname = 'seed',
    alt_names = {'seedM', 'seed.compound'},
    prefixes = {
        'http://identifiers.org/seed.compound/',
        'https://identifiers.org/seed.compound/'
    },
    recogniser = re.compile(r'^cpd\d+$'),
    objectname = 'C_SEED',
)

# Sabio-RK compound
C_SABIORK = IdentifierTypeCompound(
    name = 'Sabio-RK Compound',
    shortname = 'sabiork',
    alt_names = {'sabiork', 'sabiorkM', 'sabiork.compound'},
    prefixes = {
        'http://identifiers.org/sabiork.compound/',
        'https://identifiers.org/sabiork.compound/'
    },
    recogniser = re.compile(r'^\d+$'),
    objectname = 'C_SABIORK',
)

# HMDB
C_HMDB = IdentifierTypeCompound(
    name = 'HMDB',
    shortname = 'hmdb',
    alt_names = {'hmdb', 'HMDB'},
    prefixes = {
        'http://identifiers.org/hmdb/',
        'https://identifiers.org/hmdb/'
    },
    objectname = 'C_HMDB',
)

# Reactome
C_REACTOME = IdentifierTypeCompound(
    name = 'Reactome',
    shortname = 'reactome',
    alt_names = {'reactomeM', 'reactome.compound'},
    prefixes = {
        'http://identifiers.org/reactome/',
        'https://identifiers.org/reactome/'
    },
    recogniser = re.compile(r'(^R-[A-Z]{3}-\d+(-\d+)?(\.\d+)?$)|(^REACT_\d+(\.\d+)?$)'),
    objectname = 'C_REACTOME',
)

# PDBe
C_PDBE = IdentifierTypeCompound(
    name = 'PBDE Compound',
    shortname = 'pdbe',
    alt_names = {'PDBeChem'},
    objectname = 'C_PDBE',
)

C_BIOCYC = IdentifierTypeCompound(
    name = 'BioCyc Compound',
    shortname = 'biocyc',
    alt_names = {'biocyc'},
    prefixes = {
        'https://identifiers.org/biocyc/',
        'http://identifiers.org/biocyc/'
    },
    recogniser = re.compile(r'^[A-Z-0-9]+(:)?[A-Za-z0-9+_.%-:]+$'),
    objectname = 'C_BIOCYC',
)

# MetaMDB
C_METAMDB = IdentifierTypeCompound(
    name = 'MetaMDB Compound',
    shortname = 'metamdb_c',
    alt_names = {'metamdb_c'},
    objectname = 'C_METAMDB',
)

C_BRENDA = IdentifierTypeCompound(
    name = 'Brenda Compound',
    shortname = 'brenda_c',
    alt_names = {'brenda_c'},
    objectname = 'C_BRENDA',
)

# TODO ChemSpider

# Structure representation Identifier Types
# ---------------------------------------------------------------------------------------------------------------------

S_UNKNOWN = IdentifierTypeStructureRepresentation(
    name = 'Unknown structure',
    shortname = 'unknown_s',
    alt_names = set(),
    prefixes = set(),
    objectname = 'S_UNKNOWN',
)

# S_SMILES
S_SMILES = IdentifierTypeStructureRepresentation(
    name = 'S_SMILES',
    shortname = 'smiles',
    alt_names = {'smiles'},
    objectname = 'S_SMILES',
)

# InChI
S_INCHI = IdentifierTypeStructureRepresentation(
    name = 'InChI',
    shortname = 'inchi',
    alt_names = {'inchi'},
    objectname = 'S_INCHI',
)

# MolFile
S_MOLFILE = IdentifierTypeStructureRepresentation(
    name = 'Molfile',
    shortname = 'molfile',
    alt_names = {'molfile'},
    objectname = 'S_MOLFILE',
)

# S_GML
S_GML = IdentifierTypeStructureRepresentation(
    name = 'GML',
    shortname = 'gml',
    alt_names = {'gml'},
    objectname = 'S_GML',
)

# Reaction Identifier Types
# ---------------------------------------------------------------------------------------------------------------------
R_UNKNOWN = IdentifierTypeReaction(
    name = 'Unknown Reaction',
    shortname = 'unknown_r',
    alt_names = set(),
    prefixes = set(),
    objectname = 'R_UNKNOWN',
)

R_NAME = IdentifierTypeReaction(
    name = 'Reaction Name',
    shortname = 'rname',
    alt_names = {'rname'},
    objectname = 'R_NAME',
)

R_MNX = IdentifierTypeReaction(
    name = 'MetaNetX Reaction',
    shortname = 'mnx_r',
    alt_names = {'mnxr'},
    prefixes = {
        'https://identifiers.org/metanetx.reaction/',
        'http://identifiers.org/metanetx.reaction/'
    },
    recogniser = re.compile(r'^(MNXR\d+|EMPTY)$'),
    objectname = 'R_MNX',
)

R_METACYC = IdentifierTypeReaction(
    name = 'MetaCyc Reaction',
    shortname = 'metacyc_r',
    alt_names = {'metacycr', 'metacycR', 'metacyc.reaction'},
    prefixes = {
        'https://identifiers.org/metacyc.reaction/',
        'http://identifiers.org/metacyc.reaction/'
    },
    recogniser = re.compile(r'^[A-Za-z0-9+_.%-:]+$'),
    objectname = 'R_METACYC',

)

R_BIGG = IdentifierTypeReaction(
    name = 'BiGG Reaction',
    shortname = 'bigg_r',
    alt_names = {'biggr', 'bigg_r', 'biggR', 'bigg.reaction'},
    prefixes = {
        'https://identifiers.org/bigg.reaction/',
        'http://identifiers.org/bigg.reaction'
    },
    recogniser = re.compile(r'^[a-z_A-Z0-9]+$'),
    stdfunc = _std_bigg_r,
    objectname = 'R_BIGG',
)

R_SEED = IdentifierTypeReaction(
    name = 'SEED Reaction',
    shortname = 'seed_r',
    alt_names = {'seed_r', 'seedR', 'seed.reaction'},
    recogniser = re.compile(r'^rxn\d+$'),
    objectname = 'R_SEED',
)

R_KEGG = IdentifierTypeReaction(
    name = 'KEGG Reaction',
    shortname = 'kegg_r',
    alt_names = {'kegg_r', 'keggR', 'kegg.reaction'},
    prefixes = {
        'https://identifiers.org/seed.reaction/',
        'http://identifiers.org/seed.reaction'
    },
    recogniser = re.compile(r'^R\d+$'),
    objectname = 'R_KEGG',
)

R_RHEA = IdentifierTypeReaction(
    name = 'RHEA Reaction',
    shortname = 'rhea_r',
    alt_names = {'rheaR', 'rhea'},
    prefixes = {
        'https://identifiers.org/rhea/',
        'http://identifiers.org/rhea/'
    },
    recogniser = re.compile(r'^\d{5}$'),
    objectname = 'R_RHEA',
)

R_SABIORK = IdentifierTypeReaction(
    name = 'Sabio RK Reaction',
    shortname = 'sabiork_r',
    alt_names = {'sabiorkR', 'sabiork.reaction'},
    prefixes = {
        'https://identifiers.org/sabiork.reaction/',
        'http://identifiers.org/sabiork.reaction'
    },
    recogniser = re.compile(r'^\d+$'),
    objectname = 'R_SABIORK',
)

R_METAMDB = IdentifierTypeReaction(
    name = 'MetaMDB reaction',
    shortname = 'metamdb_r',
    alt_names = {'metamdb_r'},
    objectname = 'R_METAMDB',
)

R_MCSA = IdentifierTypeReaction(
    name = 'MCSA reaction',
    shortname = 'mcsa_r',
    alt_names = {'mcsa_r'},
    objectname = 'R_MCSA',
)

R_BRENDA = IdentifierTypeReaction(
    name = 'Brenda Reaction',
    shortname = 'brenda_r',
    alt_names = {'brenda_r'},
    objectname = 'R_BRENDA',
)

# Enzyme Identifier Types
# --------------------------------------------------------------------------------------------------------------
E_UNKNOWN = IdentifierTypeEnzyme(
    name = 'Unknown enzyme',
    shortname = 'unknown_e',
    alt_names = set(),
    prefixes = set(),
    objectname = 'E_UNKNOWN',
)

E_NAME = IdentifierTypeEnzyme(
    name = 'Enzyme Name',
    shortname = 'ename',
    alt_names = {'ename'},
    objectname = 'E_NAME',
)

E_EC = IdentifierTypeEnzyme(
    name = 'EC',
    shortname = 'ec',
    alt_names = {'ec'},
    prefixes = {
        'http://identifiers.org/ec-code/',
        'https://identifiers.org/ec-code/'
    },
    recogniser = re.compile(r'^\d+\.-\.-\.-|\d+\.\d+\.-\.-|\d+\.\d+\.\d+\.-|\d+\.\d+\.\d+\.(n)?\d+$'),
    objectname = 'E_EC',
)

# Atom-to-Atom Map Identifier Types
# --------------------------------------------------------------------------------------------------------------
A_UNKNOWN = IdentifierTypeAAM(
    name = 'Unknown AAM',
    shortname = 'unknown_a',
    alt_names = set(),
    prefixes = set(),
    objectname = 'A_UNKNOWN',
)

A_REACTIONSMILES = IdentifierTypeAAM(
    name = 'Reaction Smiles',
    shortname = 'rsmiles',
    alt_names = {'rsmiles'},
    objectname = 'A_REACTIONSMILES',
)

A_RXN = IdentifierTypeAAM(
    name = 'RXN',
    shortname = 'rxn',
    alt_names = {'rxn'},
    objectname = 'A_RXN',
)

A_GML_RULE = IdentifierTypeAAM(
    name = 'S_GML rule',
    shortname = 'gml_rule',
    alt_names = {'gml_rule'},
    objectname = 'A_GML_RULE',
)

# Lists of ID types
# --------------------------------------------------------------------------------------------------------------
id_types_compound: list[IdentifierTypeCompound] = [
    C_MNX, C_BIGG, C_PUBCHEM, C_KEGG, C_CHEBI, C_ECMDB, C_INCHIKEY, C_SLM, C_LIPIDMAPS, C_METACYC, C_ENVIPATH,
    C_SEED, C_SABIORK, C_HMDB, C_REACTOME, C_BIOCYC, C_METAMDB,
]
id_types_structure_representation: list[IdentifierTypeStructureRepresentation] = [
    S_SMILES, S_INCHI, S_MOLFILE, S_GML,
]
id_types_reaction: list[IdentifierTypeReaction] = [
    R_BIGG, R_MNX, R_METACYC, R_SEED, R_KEGG, R_RHEA, R_SABIORK, R_METAMDB, R_MCSA,
]
id_types_enzyme: list[IdentifierTypeEnzyme] = [
    E_EC
]
id_types_aam: list[IdentifierTypeAAM] = [
    A_UNKNOWN, A_REACTIONSMILES, A_RXN, A_GML_RULE
]
