from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Optional, ClassVar

import psycopg as pg
from psycopg import sql as sql

from chemrecon.entrygraph.filter import EntryFilter, RelationFilter
from chemrecon.schema.db_object import Entry, Relation


class ExploreProcedure[T_takes: Entry, T_gives: Entry](ABC):
    takes_entrytype: type[T_takes]
    gives_entrytype: type[T_gives]
    relation_type: type[Relation]  # Filter on the produced relations
    entry_filter: Optional[EntryFilter[T_gives]]   # Filter on T_gives
    relation_filter: Optional[RelationFilter]

    # Pre-computed SQL snippets
    _opt_relation_filter_clause: sql.Composable
    _opt_entry_filter_clause: sql.Composable

    n_params: ClassVar[int]

    # SQL
    q: sql.Composed
    cursor: pg.Cursor

    def __init__(
        self,
        relationtype: type[Relation],
        relation_filter: Optional[RelationFilter],
        entry_filter: Optional[EntryFilter],
    ):
        # Init
        self.relation_type = relationtype
        self.relation_filter = relation_filter
        self.entry_filter = entry_filter

        if relation_filter is not None:
            #assert relation_filter.relation_type == relationtype
            pass
        if entry_filter is not None:
            #assert entry_filter.entry_type == self.gives_entrytype
            pass

        # Generate the various clauses ahead of time
        match self.relation_filter:
            # TODO filtering
            #case RelationFilterSource():
            #    # Filter
            #    self._opt_relation_filter_clause = sql.SQL("""
            #        AND (rel_src in {src_set})
            #    """).format(
            #        src_set = list(relation_filter.allowed_sources)
            #    )
            case _:
                # No filter
                self._opt_relation_filter_clause = sql.SQL('')

        # Entry clause
        # TODO change to filter entries on the DB level (hard for symmetric relations?)
        #   for symmetric - BOTH recon ids need to be in the set of allowed id types?
        self._opt_entry_filter_clause = self.make_entry_filter_clause()

        # Generate the query except for the formats
        # Clause is prepared with as much formating as possible now, %(reconids)s param will be used when
        # calling only
        self.q = sql.SQL("""
          SELECT *
          FROM {rel_view}
          WHERE {where_clause} {opt_relation_clause} {opt_entry_clause}
          ;
        """).format(
            rel_view = sql.Identifier(f'{self.relation_type.get_table_name()}_v'),
            where_clause = self.make_where_clause(),
            opt_relation_clause = self._opt_relation_filter_clause,
            opt_entry_clause = self._opt_entry_filter_clause
        )

    @abstractmethod
    def make_where_clause(self) -> sql.SQL:
        """ Returns a parametrised SQL string to be used when calling. """
        pass

    @abstractmethod
    def make_entry_filter_clause(self) -> sql.SQL:
        """ Returns an SQL string filtering correctly by recon_id_1, recon_id_2, or both. """
        pass


class ExploreProcedureT1[T_takes: Entry, T_gives: Entry](ExploreProcedure[T_takes, T_gives]):
    # Explore from T1
    takes_entrytype: type[T_takes]
    gives_entrytype: type[T_gives]
    relation_type: type[Relation[T_takes, T_gives]]
    entry_filter: Optional[EntryFilter[T_gives]]
    relation_filter: Optional[RelationFilter[T_takes, T_gives]]

    n_params: ClassVar[int] = 1

    def __init__(
        self,
        relationtype: type[Relation[T_takes, T_gives]],
        relation_filter: Optional[RelationFilter[T_takes, T_gives]],
        entry_filter: Optional[EntryFilter[T_gives]]
    ):
        assert not relationtype.symmetric
        self.takes_entrytype = relationtype.source_entrytype
        self.produces_entrytype = relationtype.target_entrytype
        self.relation = relationtype

        super().__init__(relationtype, relation_filter, entry_filter)

    def make_where_clause(self) -> sql.SQL:
        return sql.SQL("recon_id_1 = %s")

    def make_entry_filter_clause(self) -> sql.SQL:
        # TODO
        return sql.SQL('')


class ExploreProcedureT2[T_takes: Entry, T_gives: Entry](ExploreProcedure[T_takes, T_gives]):
    # Explore from T2
    takes_entrytype: type[T_takes]
    gives_entrytype: type[T_gives]
    relation_type: type[Relation[T_gives, T_takes]]
    entry_filter: Optional[EntryFilter[T_gives]]
    relation_filter: Optional[RelationFilter[T_gives, T_takes]]

    n_params: ClassVar[int] = 1

    def __init__(
        self,
        relationtype: type[Relation[T_gives, T_takes]],
        relation_filter: Optional[RelationFilter[T_gives, T_takes]],
        entry_filter: Optional[EntryFilter[T_gives]]
    ):
        assert not relationtype.symmetric
        self.takes_entrytype = relationtype.target_entrytype
        self.produces_entrytype = relationtype.source_entrytype
        self.relation = relationtype

        super().__init__(relationtype, relation_filter, entry_filter)

    def make_where_clause(self) -> sql.SQL:
        return sql.SQL("recon_id_2 = %s")

    def make_entry_filter_clause(self) -> sql.SQL:
        # TODO
        return sql.SQL('')


class ExploreProcedureSym[T: Entry](ExploreProcedure[T, T]):
    # Explore from T
    takes_entrytype: type[T]
    gives_entrytype: type[T]
    relation_type: type[Relation[T, T]]
    entry_filter: Optional[EntryFilter[T]]
    relation_filter: Optional[RelationFilter[T, T]]

    n_params: ClassVar[int] = 2

    def __init__(
        self,
        relationtype: type[Relation[T, T]],
        relation_filter: Optional[RelationFilter[T, T]],
        entry_filter: Optional[EntryFilter[T]]
    ):
        assert relationtype.symmetric
        self.takes_entrytype = relationtype.source_entrytype
        self.produces_entrytype = relationtype.target_entrytype
        self.relation = relationtype

        super().__init__(relationtype, relation_filter, entry_filter)

    def make_where_clause(self) -> sql.SQL:
        return sql.SQL("(recon_id_1 = %s) OR (recon_id_2 = %s)")

    def make_entry_filter_clause(self) -> sql.SQL:
        # TODO
        return sql.SQL('')
