import logging
from .base import DataSource
from ..exceptions import NdffException
import psycopg2
from psycopg2.extras import DictCursor

log = logging.getLogger(__name__)


class PostgresDataSource(DataSource):

    def __init__(self, settings=None):
        super().__init__(settings)
        log.debug(f'Init PostgresDataSource, settings: {self.settings}')
        self.host = None
        self.port = None
        self.user = None
        self.password = None
        self.dbname = None
        self.timeout = None
        self.table = None
        # save and test connection
        self.conn = None
        self._set_settings()

    def _set_settings(self):
        if 'postgres_host' in self.settings:
            self.host = self.settings['postgres_host']
        if 'postgres_port' in self.settings:
            self.port = self.settings['postgres_port']
        if 'postgres_user' in self.settings:
            self.user = self.settings['postgres_user']
        if 'postgres_password' in self.settings:
            self.password = self.settings['postgres_password']
        if 'postgres_dbname' in self.settings:
            self.dbname = self.settings['postgres_dbname']
        if 'postgres_timeout' in self.settings:
            self.timeout = self.settings['postgres_timeout']
        if 'postgres_table' in self.settings:
            self.table = self.settings['postgres_table']

        # only if seemingly valid host and user etc.:
        if self.host is not None and self.user is not None and self.dbname is not None:
            try:
                self.conn = psycopg2.connect(
                                host=self.host,
                                port=self.port,
                                user=self.user,
                                password=self.password,
                                dbname=self.dbname,
                                connect_timeout=self.timeout)
            except NdffException:
                log.error('Unable to connect to {} database using "{}@{}:{}" within {} seconds'.format(self.host, self.user, self.dbname, self.port, self.timeout))

    def get_records(self):
        # we need ordering, else it's impossible to write tests, as the order is not defined...
        """
        https://stackoverflow.com/questions/12379221/sql-query-to-find-primary-key-of-a-table/12379241
        
        SELECT k.COLUMN_NAME
        FROM information_schema.table_constraints t
        LEFT JOIN information_schema.key_column_usage k
        USING(constraint_name,table_schema,table_name)
        WHERE t.constraint_type='PRIMARY KEY'
            AND t.table_schema='schema'
            AND t.table_name='table';
        """
        schema_table = self.table.split('.')
        if len(schema_table) == 2:
            schema = schema_table[0]
            table = schema_table[1]
        elif len(schema_table) == 1:
            schema = 'public'
            table = self.table
        else:
            raise NdffException('In get_records(): no schema/table defined to get records from')

        primary_key_query = f"""SELECT k.COLUMN_NAME
        FROM information_schema.table_constraints t
        LEFT JOIN information_schema.key_column_usage k
        USING(constraint_name,table_schema,table_name)
        WHERE t.constraint_type='PRIMARY KEY'
            AND t.table_schema='{schema}'
            AND t.table_name='{table}';
        """
        cur = self.conn.cursor()
        cur.execute(primary_key_query)
        primary_key_column = cur.fetchone()
        order_clause = ''
        if primary_key_column:
            # not sure what happens if we have a primary key consisting of multiple rows
            # but think/hope we will not get those...
            if len(primary_key_column) > 1:
                log.warning(f'Postgresql Datasource table {self.table} does seem to have a primary key consisting of {len(primary_key_column)} columns, only use first one... proceed with caution')
            order_clause = f' order by {primary_key_column[0]}'
        else:
            log.warning(f'Postgresql Datasource table {self.table} does not seem to have a primary key... proceed with caution')

        cur = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        select_query = f'select * from {self.table} {order_clause}'
        cur.execute(select_query)
        return cur.__iter__()
