from luigi import Task, Parameter, IntParameter, ListParameter
from tasks.base_tasks import TempTableTask
from tasks.meta import current_session
from tasks.targets import ConstraintExistsTarget
from tasks.targets import PostgresTarget, PostgresFunctionTarget
from lib.logger import get_logger

LOGGER = get_logger(__name__)


class _CountryTask:
    @property
    def _country(self):
        raise NotImplementedError('_CountryTask must define _country()')


class DenormalizedHierarchy(Task, _CountryTask):
    year = IntParameter()

    def requires(self):
        '''
        Subclasses must override this and return a dictionary with the following elements:
            - 'function': The definition of the `GetParentsFunction`
            - 'rel': The `HierarchyChildParentsUnion` defining the hierarchy relations
        '''
        raise NotImplementedError('DenormalizedHierarchy must define requires()')

    def _create_indexes(self, session):
        LOGGER.info('Creating index on {table}'.format(table=self.output().table))
        query = '''
                CREATE INDEX idx_{tablename} ON {table} (geoid, level);
                '''.format(
                    tablename=self.output().tablename,
                    table=self.output().table,
                )
        session.execute(query)

        LOGGER.info('Clustering table {table}'.format(table=self.output().table))
        query = '''
                CLUSTER {table} USING idx_{tablename};
                '''.format(
                    tablename=self.output().tablename,
                    table=self.output().table,
                )
        session.execute(query)

    def run(self):
        session = current_session()

        session.execute('CREATE SCHEMA IF NOT EXISTS {schema}'.format(
            schema=self.output().schema)
        )

        LOGGER.info('Creating table {table}'.format(table=self.output().table))
        query = '''
                CREATE TABLE {output} AS
                SELECT child_id as geoid, child_level as level, parent_names.*
                FROM {input} i,
                    {get_parents_function}(child_id, child_level) parent_names
                '''.format(
                    output=self.output().table,
                    input=self.input()['rel'].table,
                    get_parents_function=self.input()['function'].function,
                )
        session.execute(query)

        self._create_indexes(session)
        session.commit()

    def output(self):
        return PostgresTarget('tiler', '{country}_dz_hierarchy_geonames_{year}'.format(country=self._country,
                                                                                       year=self.year))


class GetParentsFunction(Task, _CountryTask):
    year = IntParameter()

    def version(self):
        return 2

    def requires(self):
        '''
        Subclasses must override this and return a dictionary with the following elements:
            - 'rel': The `HierarchyChildParentsUnion` defining the hierarchy relations
            - 'info': The `HierarchyInfoUnion` defining the geonames
        '''
        raise NotImplementedError('GetParentsFunction must define requires()')

    def run(self):
        input_ = self.input()
        rel_table = input_['rel'].table
        schema = self.output().schema
        function = self.output().function_name

        session = current_session()
        levels_query = '''
                SELECT DISTINCT parent_level
                FROM {input_relations}'''.format(input_relations=rel_table)
        levels = [l[0] for l in session.execute(levels_query).fetchall()]
        levels += ["{}_id".format(l) for l in levels]

        level_types = ', '.join(['"{}" text'.format(l) for l in levels])
        cols_type = """
                "{schema}".{function}_levels
            """.format(schema=schema, function=function)

        session.execute('''
                DROP TYPE IF EXISTS {cols_type} CASCADE
        '''.format(cols_type=cols_type))

        session.execute('''
                CREATE TYPE {cols_type} as ({level_types})
        '''.format(cols_type=cols_type, level_types=level_types))

        query = '''
                CREATE OR REPLACE FUNCTION "{schema}".{function}_json (geoid_p TEXT, level_p TEXT)
                RETURNS JSONB
                AS $$
                DECLARE
                    children JSONB DEFAULT NULL;
                BEGIN
                    WITH RECURSIVE children(child_id, child_level, parent_id, parent_level, geoname) AS (
                        SELECT p.child_id, p.child_level, p.parent_id, p.parent_level, n.geoname
                        FROM {input_relations} p
                        LEFT JOIN {input_geonames} n ON p.parent_id = n.geoid AND p.parent_level = n.level
                        WHERE p.child_id = geoid_p AND p.child_level = level_p
                        UNION ALL
                        SELECT p.child_id, p.child_level, p.parent_id, p.parent_level, n.geoname
                        FROM {input_relations} p
                        INNER JOIN children c ON c.parent_id = p.child_id AND c.parent_level = p.child_level
                        LEFT JOIN {input_geonames} n ON p.parent_id = n.geoid AND p.parent_level = n.level
                    )
                    SELECT ('{{' || string_agg(
                                        '"' || parent_level || '": "' || geoname || '", ' ||
                                        '"' || parent_level || '_id": "' || parent_id || '"',
                                        ',') || '}}')::JSONB
                        INTO children
                        FROM children;

                    RETURN children;
                END
                $$ LANGUAGE plpgsql PARALLEL SAFE;
                '''.format(
                    schema=schema,
                    function=function,
                    input_relations=rel_table,
                    input_geonames=input_['info'].table,
                )

        LOGGER.debug(query)
        session.execute(query)

        cols_query = '''
                CREATE OR REPLACE FUNCTION "{schema}".{function}(geoid TEXT, level TEXT)
                RETURNS {cols_type}
                AS $$
                DECLARE
                    _output {cols_type};
                BEGIN
                    SELECT * FROM jsonb_populate_record(null::{cols_type},
                                                        "{schema}".{function}_json(geoid, level))
                    INTO _output;

                    RETURN _output;
                END
                $$ LANGUAGE plpgsql PARALLEL SAFE;
                '''.format(
            schema=schema,
            function=function,
            level_types=level_types,
            cols_type=cols_type)
        LOGGER.debug(cols_query)
        session.execute(cols_query)

        session.commit()

    def output(self):
        name = 'get_{country}_parents_{year}'.format(country=self._country,
                                                     year=self.year)
        return PostgresFunctionTarget(self.input()['rel'].schema, name)


class Hierarchy(Task, _CountryTask):
    year = IntParameter()

    def requires(self):
        '''
        Subclasses must override this and return a dictionary with the following elements:
            - 'info': The `HierarchyInfoUnion` defining the geonames
            - 'rel': The `HierarchyChildParentsUnion` defining the hierarchy relations
        '''
        raise NotImplementedError('Hierarchy must define requires()')

    def run(self):
        session = current_session()
        input_ = self.input()

        session.execute('ALTER TABLE {rel_table} DROP CONSTRAINT IF EXISTS '
                        '{country}hierarchy_fk_parent'.format(
            rel_table=input_['rel'].qualified_tablename,
            country=self._country, ))

        session.execute('ALTER TABLE {rel_table} ADD CONSTRAINT '
                        '{country}hierarchy_fk_parent '
                        'FOREIGN KEY (child_id, child_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=input_['rel'].qualified_tablename,
            info_table=input_['info'].qualified_tablename,
            country=self._country, ))

        session.execute('ALTER TABLE {rel_table} DROP CONSTRAINT IF EXISTS '
                        '{country}hierarchy_fk_child'.format(
            rel_table=input_['rel'].qualified_tablename,
            country=self._country, ))
        session.execute('ALTER TABLE {rel_table} ADD CONSTRAINT '
                        '{country}hierarchy_fk_child '
                        'FOREIGN KEY (parent_id, parent_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=input_['rel'].qualified_tablename,
            info_table=input_['info'].qualified_tablename,
            country=self._country, ))
        session.commit()

    def output(self):
        table = self.input()['info']
        return ConstraintExistsTarget(table.schema, table.tablename,
                                      '{country}hierarchy_fk_child'.format(country=self._country))


class _YearCountryLevelsTask(_CountryTask):
    year = IntParameter()
    levels = ListParameter(significant=False)


class HierarchyInfoUnion(TempTableTask, _YearCountryLevelsTask):
    def requires(self):
        '''
        Subclasses must override this and return a list of `LevelInfo` (one for each level)
        '''
        raise NotImplementedError('HierarchyInfoUnion must define requires()')

    def _union_query(self, tables, output):
        unions = ['''SELECT geoid, level, string_agg(geoname, ', ') geoname
                     FROM {table}
                     GROUP BY geoid, level'''.format(table=table.qualified_tablename)
                  for table in tables]
        return 'CREATE TABLE {output} AS {unions}'.format(
            output=output,
            unions=' UNION ALL '.join(unions))

    def run(self):
        session = current_session()
        tablename = self.output().qualified_tablename

        session.execute(self._union_query(self.input(), tablename))
        alter_sql = 'ALTER TABLE {tablename} ADD PRIMARY KEY (geoid, level)'
        session.execute(alter_sql.format(tablename=tablename))
        session.commit()


class HierarchyChildParentsUnion(TempTableTask, _YearCountryLevelsTask):
    def requires(self):
        '''
        Subclasses must override this and return a dictionary with the following elements:
            - 'hierarchy': a list of `HierarchyChildParent` (one for each child-parent relation)
        '''
        raise NotImplementedError('HierarchyChildParentsUnion must define requires()')

    def _child_parents(self):
        child_parents = []
        previous = None
        for idx, level in enumerate(self.levels):
            if previous:
                parents = self.levels[idx:]
                if parents:
                    child_parents.append([previous, parents])
            previous = level
        return child_parents

    def _union_query(self, tables, output):
        unions = ['SELECT * FROM {table}'.format(table=table.qualified_tablename)
                  for table in tables]
        return 'CREATE TABLE {output} AS {unions}'.format(
            output=self.output().qualified_tablename,
            unions=' UNION ALL '.join(unions))

    def run(self):