from __future__ import annotations

from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, ClassVar, Optional, OrderedDict

from chemrecon.core.id_types import IdentifierType
from chemrecon.schema.enums import Quality, SourceDatabase


class DatabaseObject(ABC):
    """ Most generic database object, covers both entries and relations.
    """
    entrytype_name: ClassVar[str]
    _table_name: ClassVar[str]

    @classmethod
    def get_table_name(cls) -> str:
        return cls._table_name


class Column:
    """ Wrapper for database columns.
    """
    name: str
    col_type: type
    serial: bool  # If numeric, whether serial
    index_hash: bool  # Whether this should be hashed if used in an index.

    def __init__(self, name: str, col_type: type, serial: bool = False, index_hash: bool = False):
        self.name = name
        self.col_type = col_type
        self.serial = serial
        self.index_hash = index_hash

    def __repr__(self):
        return f'col: {self.name}'


# Entry abstract base class
# --------------------------------------------------------------------------------------------------------------
class Entry(DatabaseObject, ABC):
    """ Generic base class for entries.
    """
    # Attributes

    #: Internal identifier in the *ChemRecon* database.
    #: Normally, entries have a nonnegative `recon_id`, unique to the table.
    #: A negative `recon_id` indicates that the object is 'virtual', that is, it was created by a procedural relation,
    #: and does not exist in the database.
    #: A `recon_id` of `None` indicates that the entry is not stored in the database.
    recon_id: Optional[int]

    # Database
    _columns: ClassVar[list[Column]]  # List of columns. The zeroth (recon_id) will be considered the PK.
    _index: ClassVar[list[int]]  # List of the columns on which to create a (possibly hashed) index.

    # For visualization
    _draw_colour: ClassVar[str]

    def get_columns_with_values(self, include_recon_id = True) -> dict[Column, Any]:
        """ Get the columns of this entry with values. """
        return OrderedDict(
            (c, self.__getattribute__(c.name)) for c in self._columns if
            (c is not col_recon_id) or include_recon_id
        )

    def get_index_columns_with_values(self) -> dict[Column, Any]:
        """ Get the index (primary key) columns of this entry with values. """
        return OrderedDict(
            (c, self.__getattribute__(c.name)) for c in self.get_index_columns()
        )

    # Methods for database interaction
    @classmethod
    def get_columns(cls, include_recon_id = True) -> list[Column]:
        if include_recon_id:
            return cls._columns
        else:
            return cls._columns[1:]  # Assuming recon_id is columns[0]

    @classmethod
    def get_index_indices(cls) -> list[int]:
        return cls._index

    @classmethod
    def get_index_columns(cls) -> list[Column]:
        return [cls._columns[i] for i in cls._index]

    # Misc
    def __init__(self, recon_id: Optional[int] = None):
        super().__init__()
        self.recon_id = recon_id

    def __repr__(self):
        attr_cols: str = ', '.join(
            f'{c.name}: {v}'
            for c, v in self.get_columns_with_values(include_recon_id = False).items()
            if (v is not None) and (v != [])
        )
        return f'<{self._table_name} {self.recon_id or '-'}: {attr_cols}>'

    # Comparison and identity
    def __eq__(self, other: Entry):
        if type(self) is not type(other):
            return False
        elif self.recon_id is not None and other.recon_id is not None:
            # Compare recon_id if applicable
            return self.recon_id == other.recon_id
        else:
            # Else, compare by index columns
            return all(
                self.__getattribute__(col.name) == other.__getattribute__(col.name)
                for col in self.get_index_columns()
            )

    def __hash__(self):
        return tuple(self.__getattribute__(col.name) for col in self.get_index_columns()).__hash__()

    def __lt__(self, other: Entry):
        # Compare by reconid, otherwise by index
        if (self.recon_id is not None) and (other.recon_id is not None):
            return self.recon_id.__lt__(other.recon_id)
        else:
            # Compare by index columns
            return tuple(self.get_columns_with_values().values()) < tuple(
                other.get_columns_with_values().values())

    @classmethod
    def get_supertype_of_id_types(cls) -> Optional[type[IdentifierType]]:
        for c in cls.get_columns():
            if c.name == 'id_type':
                if not issubclass(c.col_type, IdentifierType):
                    raise ValueError('Error in identifier type assignment of column.')
                return c.col_type
        return None

    # Visualisation
    @abstractmethod
    def _vis_str(self) -> str:
        """ Get the string of the primary information for visualization.
        """
        pass

    def _vis_attrs(self) -> dict[str, str]:
        """ Define extra attributes used in visualization.
        """
        return {
            'fillcolor': f'"#{self._draw_colour}"'
        }

    # Encoding (for memoization/caching)
    def encode(self):
        return str(self.recon_id).encode()

    def serialise(self, index_only: bool = False) -> dict:
        d = dict()
        if index_only:
            for col, val in self.get_index_columns_with_values().items():
                d[col.name] = serialise_col(col, val)
        else:
            for col, val in self.get_columns_with_values().items():
                d[col.name] = serialise_col(col, val)
        return d


# Entry as a database stand-in
class SourceEntry(Entry, ABC):
    """ An entry which stands for a database entry in one of the source databases.
    """
    source_id: str
    id_type: Enum

    def _vis_str(self) -> str:
        return f'{self.id_type.name}: {self.source_id}'


# Relation
# --------------------------------------------------------------------------------------------------------------
class Relation[T1, T2](DatabaseObject, ABC):
    """ Generic base class for relations between entries.
    """
    # Attributes
    recon_id_1: Optional[int]  #: Recon ID of source
    recon_id_2: Optional[int]  #: Recon ID of target
    entry_1: Optional[T1]  # Entries on either end. Used on prototypes, and can also be fetched via the db.
    entry_2: Optional[T2]

    # Database
    ignore_generation_limit: ClassVar[bool] = False  # If true, will continue exploring after generation limit
    symmetric: ClassVar[bool]  # If symmetric, will always have recon_id_1 <= recon_id_2

    # Set if this is the main table of another inverse relation
    has_inverse: ClassVar[Optional[type[Relation]]]  #: :meta private:

    _attribute_columns: ClassVar[list[Column]]
    _index: ClassVar[list[int]]  # List of col index in attr list on which to create an index.

    source_entrytype: ClassVar[type[Entry]]  #: :meta private:
    target_entrytype: ClassVar[type[Entry]]  #: :meta private:

    # Methods for database interaction
    @classmethod
    def get_attribute_columns(cls) -> list[Column]:
        # Note: does not include source/target columns
        return cls._attribute_columns

    @classmethod
    def get_columns(cls, include_recon_ids: bool = True) -> list[Column]:
        if include_recon_ids:
            return [col_recon_id_1, col_recon_id_2, *cls._attribute_columns]
        else:
            return cls._attribute_columns

    @classmethod
    def get_index_indices(cls) -> list[int]:
        return cls._index

    @classmethod
    def get_index_columns(cls) -> list[Column]:
        return [col_recon_id_1, col_recon_id_2, *[
            cls._attribute_columns[i]
            for i in cls._index
        ]]

    def get_columns_with_values(self, include_recon_ids: bool = True) -> dict[Column, Any]:
        """ Get the columns of this entry with values. """
        return OrderedDict(
            (c, self.__getattribute__(c.name)) for c in self.get_columns(include_recon_ids)
        )

    def get_index_columns_with_values(self) -> dict[Column, Any]:
        """ Get the index (primary key) columns of this entry with values.
            This always includes `recon_id_1` and `recon_id_2`, and may include attribute columns.
        """
        return OrderedDict(
            (c, self.__getattribute__(c.name)) for c in self.get_index_columns()
        )

    @classmethod
    def get_entry_table_names(cls) -> tuple[str, str]:
        return cls.source_entrytype.get_table_name(), cls.target_entrytype.get_table_name()

    # Misc
    def __init__(
            self,
            recon_id_1: Optional[int] = None,
            recon_id_2: Optional[int] = None,
    ):
        super().__init__()
        self.recon_id_1, self.recon_id_2 = recon_id_1, recon_id_2

    def __repr__(self):
        arrow = '<->' if self.symmetric else '->'
        attr_cols: str = ', '.join(
            f'{c.name}: {v}'
            for c, v in self.get_columns_with_values(include_recon_ids = False).items()
        )
        return f'<{self._table_name} {self.recon_id_1} {arrow} {self.recon_id_2}] {attr_cols}>'

    # Visualisation
    def _vis_attrs(self) -> dict[str, str]:
        """ Define extra attributes used in visualization.
        """
        return dict()

    def _vis_str(self) -> str:
        """ Define extra attributes
        """
        return ''

# Inverse relations
class InverseRelation[T1: Entry, T2: Entry](Relation[T1, T2], ABC):
    """ Represents the inverse of another (main) relation.
        The main relation should have the has_inverse tag set.
        Exists in the database only as views of the main relation.
    """
    inverse_main_relation: ClassVar[type[Relation]]  #: :meta private:

    # Check that attribute columns are the same as for the main relation


# Procedural Relations
class ProceduralRelation[T1: Entry, T2: Entry](Relation[T1, T2], ABC):
    """ Procedural relations are not stored in the database but computed at runtime.
        Other than this, they have the same interface as normal relations.
        Some can be computed only "one-way" (e.g. a compound can be standardized, but not un-standardized)
    """

    @classmethod
    @abstractmethod
    def generate(
            cls,
            take_entry: T1,
    ) -> list[tuple[ProceduralRelation[T1, T2], T2]]:
        """ Given a T1, generate relations from that.
        """
        raise NotImplementedError()


class ProceduralGeneratorError(Exception):
    """ Should be thrown by generator methods.
    """
    pass


# Composed Relations
class ComposedRelation[T1, T2, T_intermediate](Relation[T1, T2], ABC):
    # Composed of two other relations
    rel_type_1: ClassVar[type[Relation]]  #: :meta private:
    rel_type_2: ClassVar[type[Relation]]  #: :meta private:
    intermediate_entrytype: ClassVar[type[Entry]]  #: :meta private:

    # Init based on the given relations and intermediate
    @abstractmethod
    def __init__(
            self,
            rel_1: Relation[T1, T_intermediate],
            rel_2: Relation[T_intermediate, T2],
            intermediate: T_intermediate,
            recon_id_1: Optional[int] = None,
            recon_id_2: Optional[int] = None,
    ):
        """ Init for composed relations should be overriden to set attributes of the composed relation based on the
            attributes of relations and entries used to create it.
        """
        super().__init__(recon_id_1, recon_id_2)

    # Filter functions on the composed edges, which determine whether to create the procedural relations
    @classmethod
    def filter_rel_1(cls, r: Relation[T1, T_intermediate]) -> bool:
        return True

    @classmethod
    def filter_intermediate(cls, e: T_intermediate) -> bool:
        return True

    @classmethod
    def filter_rel_2(cls, r: Relation[T_intermediate, T2]) -> bool:
        return True


# Predefined columns
# ----------------------------------------------------------------------------------------------------------------------
col_recon_id = Column('recon_id', int, serial = True)
col_source_id = Column('source_id', str)
col_source_id_hashed = Column('source_id', str, index_hash = True)
col_name = Column('name', str)
col_src = Column('src', SourceDatabase)
col_quality = Column('quality', Quality)
col_properties = Column('properties', list[str])
col_recon_id_2 = Column('recon_id_2', int)
col_recon_id_1 = Column('recon_id_1', int)
col_score = Column('score', float)


# Util
# ----------------------------------------------------------------------------------------------------------------------
def serialise_col(col: Column, value: Any) -> Any:
    if isinstance(value, Enum):
        return value.name
    else:
        return value
