from typing import Optional

import psycopg as pg
import psycopg.sql as sql

from chemrecon.schema.db_object import Entry, Relation
import chemrecon.core.query_handler

class PopulateQueryHandler(chemrecon.core.query_handler.QueryHandler):
    """ Entries and relations are added to a PopulationQueryHandler, which batches operations and
        sends them to the database.
        When added, the handler will assign a correct ReconID to the object based on the current state of the
        corresponding sequence.
        Should be used as a context.
    """

    # Adding entries to the handler
    # ----------------------------------------------------------------------------------------------------------
    # For adding entries when the assigned reconid is needed. This cannot be batched.
    def add_entry[T: Entry](
        self,
        entry: T,
        returning: bool = True,
        authoritative: bool = False
    ) -> Optional[int]:
        """ Add an entry to the DB, returning its reconid immediately.
            If authoritative, allow overriding fields added previously.
            If not authoritative, do nothing on conflict.
        """
        if returning:
            return self.add_entries([entry], returning = returning, authoritative = authoritative)[0]
        else:
            self.add_entries([entry], returning = returning, authoritative = authoritative)
            return None

    def add_entries[T: Entry](
        self,
        entries: list[T],
        returning: bool = True,
        authoritative: bool = False
    ) -> Optional[list[int]]:
        """ TODO
        """
        entrytype = type(entries[0])
        cols = entrytype.get_columns(include_recon_id = False)
        index_cols = entrytype.get_index_columns()
        table_identifier = sql.Identifier(entrytype.get_table_name())

        # TODO 'returning' clause can be removed if not returning?
        q = sql.SQL("""
            INSERT INTO {entrytypetable} ({columns})
            VALUES ({placeholders})
                ON CONFLICT ({index_columns}) 
                DO UPDATE SET ({columns}) = ROW({coalesced_values})
             RETURNING {reconid_col};
        """).format(
            entrytypetable = table_identifier,
            reconid_col = sql.Identifier('recon_id'),
            columns = sql.SQL(',').join(
                sql.Identifier(c.name) for c in cols
            ),
            index_name = sql.Identifier(f'{entrytype.get_table_name()}_index'),
            index_columns = sql.SQL(',').join(  # TODO better way with named constraint?
                sql.Identifier(c.name) if not c.index_hash
                else sql.SQL('md5({colname})').format(colname = sql.Identifier(c.name))
                for c in index_cols
            ),
            placeholders = sql.SQL(', ').join([sql.SQL('%s') for _ in cols]),
            coalesced_values = sql.SQL(',').join(
                # Coalesce values, preferring existing entry
                sql.SQL('COALESCE(excluded.{col}, {table}.{col})').format(
                    col = sql.Identifier(c.name),
                    table = table_identifier
                )
                for c in cols
            )
        )

        # Execute
        try:
            self.c.executemany(
                q,
                params_seq = [
                    list(e.get_columns_with_values(include_recon_id = False).values())
                    for e in entries
                ],
                returning = returning
            )

            # Return as dict
            if returning:
                return [i[0] for i in self._extract_results_from_cursor()]
            else:
                return None
        except pg.errors.UniqueViolation as e:
            raise KeyError(f'Unique violation: {e}\n\t{e}')
        except pg.errors.SyntaxError as e:
            print(f'ERROR:\n{q.as_string()}\n\t{e}')
            raise RuntimeError

    def add_relation[T1: Entry, T2: Entry](self, relation: Relation[T1, T2]):
        """ When adding a relation, it is expected that the 'entry_1' and 'entry_2' fields are filled such that
            the appropriate key can be looked up.
        """
        self.add_relations([relation])

    def add_relations[T1: Entry, T2: Entry](self, relations: list[Relation[T1, T2]]):
        # TODO dispatch to add by reconid or by entries, depending
        # TODO change to the plural version, specialised for adding only one
        raise NotImplementedError()

    def add_relation_by_recon_ids[T1: Entry, T2: Entry](
        self,
        recon_id_1: int,
        recon_id_2: int,
        relation: Relation[T1, T2]
    ):
        """ Add a relation, using the recon_id fields of the given relation.
        """
        self.add_relations_by_recon_ids([(recon_id_1, recon_id_2, relation)])

    def add_relations_by_recon_ids[T1: Entry, T2: Entry](
        self,
        relations: list[tuple[int, int, Relation[T1, T2]]]
    ):
        """ Add relations by recon_ids. """
        reltype = type(relations[0][2])

        # Add recon_ids to the relation, sort recon_ids of needed
        if reltype.symmetric:
            for r in relations:
                r[2].recon_id_1, r[2].recon_id_2 = sorted((r[0], r[1]))
        else:
            for r in relations:
                r[2].recon_id_1, r[2].recon_id_2 = r[0], r[1]

        cols = reltype.get_columns()
        q = sql.SQL("""
                    INSERT INTO {rel_table_name} ({cols})
                    VALUES ({placeholders})
                        ON CONFLICT DO NOTHING;
                    """).format(
            rel_table_name = sql.Identifier(reltype.get_table_name()),
            cols = sql.SQL(', ').join(sql.Identifier(c.name) for c in cols),
            placeholders = sql.SQL(', ').join([sql.SQL('%s') for _ in cols])
        )
        self.c.executemany(
            q,
            params_seq = [
                list(r[2].get_columns_with_values().values())
                for r in relations
            ]
        )

    def add_relation_by_entries[T1: Entry, T2: Entry](
        self,
        entry_1: T1,
        entry_2: T2,
        relation: Relation[T1, T2]
    ):
        """ Add a relation, using the entry fields of the given relation.
        """
        self.add_relations_by_entries([(entry_1, entry_2, relation)])

    def add_relations_by_entries[T1: Entry, T2: Entry](
        self,
        relations: list[tuple[T1, T2, Relation[T1, T2]]]
    ):
        """ Add relations using the entry fields of each given relation.
        """
        reltype = type(relations[0][2])
        all_rel_cols = reltype.get_columns()
        attrcols = reltype.get_attribute_columns()
        e1_table, e2_table = reltype.get_entry_table_names()
        e1_index_cols = reltype.source_entrytype.get_index_columns()
        e2_index_cols = reltype.target_entrytype.get_index_columns()

        if reltype.symmetric:
            q_ = sql.SQL("""
                 WITH r_id1 AS (
                     SELECT recon_id
                     FROM {entry_1_table}
                     WHERE ({entry_1_index_cols}) = ({entry_1_placeholders})
                 ), r_id2 AS (
                     SELECT recon_id
                     FROM {entry_2_table}
                     WHERE ({entry_2_index_cols}) = ({entry_2_placeholders})
                 )
                 INSERT INTO {rel_table_name} ({all_columns})
                 VALUES (
                     LEAST((SELECT recon_id FROM r_id1), (SELECT recon_id FROM r_id2)),
                     GREATEST((SELECT recon_id FROM r_id1), (SELECT recon_id FROM r_id2))
                     {attr_placeholders}
                     )
                 ON CONFLICT DO NOTHING;
            """)
        else:
            q_ = sql.SQL("""
                 WITH r_id1 AS (
                     SELECT recon_id
                     FROM {entry_1_table}
                     WHERE ({entry_1_index_cols}) = ({entry_1_placeholders})
                 ), r_id2 AS (
                     SELECT recon_id
                     FROM {entry_2_table}
                     WHERE ({entry_2_index_cols}) = ({entry_2_placeholders})
                 )
                 INSERT INTO {rel_table_name} ({all_columns})
                 VALUES ((SELECT recon_id FROM r_id1), (SELECT recon_id FROM r_id2) {attr_placeholders})
                 ON CONFLICT DO NOTHING;
            """)

        # Prepare
        q = q_.format(
            # Entry 1
            entry_1_table = sql.Identifier(e1_table),
            entry_1_index_cols = sql.SQL(', ').join(
                sql.Identifier(c.name) for c in e1_index_cols
            ),
            entry_1_placeholders = sql.SQL(', ').join(sql.SQL('%s') for _ in e1_index_cols),
            # Entry2
            entry_2_table = sql.Identifier(e2_table),
            entry_2_index_cols = sql.SQL(', ').join(
                sql.Identifier(c.name) for c in e2_index_cols
            ),
            entry_2_placeholders = sql.SQL(', ').join(sql.SQL('%s') for _ in e2_index_cols),
            # Attributes
            rel_table_name = sql.Identifier(reltype.get_table_name()),
            all_columns = sql.SQL(', ').join(
                sql.Identifier(c.name) for c in all_rel_cols
            ),
            attr_placeholders = (
                    sql.SQL(', ') + sql.SQL(', ').join(sql.SQL('%s') for _ in attrcols)
            ) if len(attrcols) > 0 else sql.SQL('')
        )

        # Execute
        self.c.executemany(
            q,
            params_seq = [
                [
                    *e1.get_index_columns_with_values().values(),
                    *e2.get_index_columns_with_values().values(),
                    *r.get_columns_with_values(include_recon_ids = False).values()
                ]
                for e1, e2, r in relations
            ],
            returning = False
        )

    # Connected adders
    def add_entry_with_relations[T1: Entry, T2: Entry](
        self, entry: T1, relations: list[tuple[Relation[T1, T2], T2]]
    ) -> int:
        """ An entry and a number of relations to possibly new entries with the first entry as the source.
            Relations do not need to have their recon_id_1, recon_id_2, or entries filled.
            Returns id of source entry
        """
        # Add the T1 entry
        e_recon_id = self.add_entry(entry, returning = True)

        if len(relations) == 0:
            # Short circuit if no relations given
            return e_recon_id

        rels, t2_entries = zip(*relations)
        rel_type = type(rels[0])
        t2_type = type(t2_entries[0])

        # Add relations and secondary entries
        self.add_relations_to_entry_with_reconid(
            recon_id = e_recon_id,
            entry_table = type(entry),
            relations = relations
        )

        # Return id of added/found source entry
        return e_recon_id


    def add_relations_to_entry_with_reconid[T1: Entry, T2: Entry](
        self,
        recon_id: int,
        entry_table: type[Entry],
        relations: list[tuple[Relation[T1, T2], T2]] | list[tuple[Relation[T2, T1], T2]]
    ) -> list[int]:
        """ Add a number of relations from the given reconid.
            Returns a list of assigned Recon IDs in the same order as the input.
        """
        if len(relations) == 0:
            return []

        # Gather types
        rels, t2_entries = zip(*relations)
        rel_type = type(rels[0])
        t2_type = type(t2_entries[0])

        # Add all T2 entries
        t2_ids: list[int] = self.add_entries(t2_entries, returning = True, authoritative = False)
        if len(t2_ids) != len(relations):
            raise ValueError

        # Dispatch, setting recon_ids of added relations depending on type
        if rel_type.symmetric:
            # T2 = T1, Relation[T1, T1]
            for t2_id, rel in zip(t2_ids, rels):
                rel: Relation[T1, T1]
                rel.recon_id_1 = recon_id
                rel.recon_id_2 = t2_id
                rel.recon_id_1, rel.recon_id_2 = sorted((rel.recon_id_1, rel.recon_id_2))  # Keep invariant

        elif rel_type.source_entrytype is entry_table:
            # Relation[T1, T2]
            for t2_id, rel in zip(t2_ids, rels):
                rel: Relation[T1, T2]
                rel.recon_id_1, rel.recon_id_2 = (recon_id, t2_id)

        elif rel_type.target_entrytype is entry_table:
            # Relation[T2, T1]
            for t2_id, rel in zip(t2_ids, rels):
                rel: Relation[T2, T1]
                rel.recon_id_1, rel.recon_id_2 = (t2_id, recon_id)

        else:
            raise ValueError()

        # Add relations by recon ids
        self.add_relations_by_recon_ids(
            [(rel.recon_id_1, rel.recon_id_2, rel) for rel in rels]
        )

        # Return recon_ids assigned by db
        return t2_ids

    # Utility functions
    # ----------------------------------------------------------------------------------------------------------
