""" Handles the creation of the tables described in 'core/schema.py'.
    All tables and objects are placed in the 'chemrecon' schema.
    Should ensure admin access of the connection before attempting to create tables.
    If tables already exist, verify that they are identical.
    This module has the following expectations of the database being connected to:
    - Database named 'chemrecon_db'
    - Initial login with 'postgres' user
    - Afterwards, login with either local_docker or local_docker_dev
    -
    -
    -
    -
"""
import argparse
import builtins
import os
from enum import Enum
from typing import get_origin, get_args
from types import GenericAlias

import psycopg as pg
import psycopg.errors
import psycopg.sql as sql

import chemrecon.schema as schema
from chemrecon import Params

from chemrecon.database import params
from chemrecon.database.connect import postgres_connect
from chemrecon.utils.set_cwd import set_cwd

from chemrecon.schema import entrytypes, relationtypes, InverseRelation, Relation

# Command-line-arguments
parser = argparse.ArgumentParser()
# Whether to set up local or public users
parser.add_argument('--params', action = 'store')
parser.add_argument('--devparams', action = 'store')

# Custom params
parser.add_argument('--host', action = 'store')
parser.add_argument('--port', action = 'store')
parser.add_argument('--database', action = 'store')
parser.add_argument('--username', action = 'store')
parser.add_argument('--password', action = 'store')


def main():
    print(f'Running ChemRecon database initialisation script.')
    args = vars(parser.parse_args())

    print(f'\nARGS\n')
    print(args)

    if 'params' not in args:
        raise ValueError('Must provide parameters file')
    if 'devparams' not in args:
        raise ValueError('Must provide dev parameters file.')

    # Connection params
    connection_params = Params(
        connection_title = 'CLI',
        db_name = args['database'],
        db_host = args['host'],
        db_port = args['port'],
        username = args['username'],
        password = args['password']
    )

    print(os.getcwd())

    # User parameters
    with open(args['params']) as f:
        lines = f.readlines()
        params_user = Params(
            connection_title = 'CLI-user',
            db_name = lines[0].strip(),
            username = lines[1].strip(),
            password = lines[2].strip(),
            db_host = lines[3].strip(),
            db_port = lines[4].strip()
        )
    with open(args['devparams']) as f:
        lines = f.readlines()
        params_dev = Params(
            connection_title = 'CLI-user',
            db_name = lines[0].strip(),
            username = lines[1].strip(),
            password = lines[2].strip(),
            db_host = lines[3].strip(),
            db_port = lines[4].strip()
        )

    conn = postgres_connect(
        connection_params,
        initialise_enums = False
    )

    # Initialize
    initialise_database(
        conn,
        dev_parameters = params_dev,
        pub_parameters = params_user,
    )


def initialise_database(
        conn: pg.Connection,
        dev_parameters: Params,
        pub_parameters: Params,
        only_create_users: bool = False
):
    # Login with the 'postgres' user to create the DB.
    # Create the database, schema and accounts from an otherwise empty database

    # Get connection parameters
    set_cwd()

    print(f'\nCreating database...')

    with conn.cursor() as c:

        # Check if DB is already populated
        # TODO

        # Load parameters
        # ------------------------------------------------------------------------------------------------------
        db_user_dev = dev_parameters.username
        db_pass_dev = dev_parameters.password

        db_user_pub = pub_parameters.username
        db_pass_pub = pub_parameters.password

        # Create users if not exist
        try:

            # Create group for public users
            # ----------------------------------------------------------------------------------------------------------

            c.execute(sql.SQL(
                'CREATE GROUP {public_group};'
            ).format(
                public_group = sql.Identifier('public_group')
            ))

            # Create users/roles
            # ----------------------------------------------------------------------------------------------------------
            c.execute(sql.SQL(
                'CREATE USER {dev_username} WITH PASSWORD {dev_password};'
            ).format(
                dev_username = sql.Identifier(db_user_dev),
                dev_password = db_pass_dev,
            ))
            c.execute(sql.SQL(
                'CREATE USER {pub_username} WITH PASSWORD {pub_password} IN GROUP {public_group};'
            ).format(
                pub_username = sql.Identifier(db_user_pub),
                pub_password = db_pass_pub,
                public_group = sql.Identifier('public_group')
            ))

        except psycopg.errors.DuplicateObject as e:
            # Already exists, reset
            conn.commit()
            print(f' -> Users already exist, skipping.')
            pass

        # If set to only prepare for restoring with psql, exit here
        # --------------------------------------------------------------------------------------------------------------
        if only_create_users:
            conn.commit()
            return

        # Create schemata
        # --------------------------------------------------------------------------------------------------------------
        c.execute(sql.SQL(  # 'meta' schema for dev purposes
            'CREATE SCHEMA {meta_schema} AUTHORIZATION {dev_username};'
        ).format(
            meta_schema = sql.Identifier('meta'),
            dev_username = sql.Identifier(db_user_dev)
        ))
        c.execute(sql.SQL(  # 'chemrecon' schema for data
            'CREATE SCHEMA {chemrecon_schema} AUTHORIZATION {dev_username};'
        ).format(
            chemrecon_schema = sql.Identifier('chemrecon'),
            dev_username = sql.Identifier(db_user_dev)
        ))

        # Create the 'meta' tables
        # --------------------------------------------------------------------------------------------------------------
        c.execute(sql.SQL("""
                          CREATE TABLE IF NOT EXISTS meta.procedures
                          (
                              procedure_name
                              TEXT,
                              procedure_version
                              INT,
                              data_source_version
                              TEXT,
                              finished_time
                              TIMESTAMP,
                              progress
                              INT
                          );
                          """))
        conn.commit()

        # Create and adapt enums
        # --------------------------------------------------------------------------------------------------------------
        print(f'Creating enums')
        for enum in schema.enums.enum_register:
            # Create enum in db
            c.execute(sql.SQL(
                'CREATE TYPE {chemrecon}.{enum_name} AS ENUM ({enum_list});'
            ).format(
                chemrecon = sql.Identifier('chemrecon'),
                enum_name = sql.Identifier(enum.__name__),
                enum_list = sql.SQL(', ').join(
                    e.name for e in enum
                )
            ))
        conn.commit()

        # Create tables
        # --------------------------------------------------------------------------------------------------------------
        # Entry tables
        # --------------------------------------------------------------------------------------------------------------
        for entrytype in entrytypes:
            table = sql.Identifier(entrytype.get_table_name())
            column_sqls: list[sql.Composable] = list()

            # Compute the columns
            for i, col in enumerate(entrytype.get_columns()):
                col: schema.db_object.Column
                colname = sql.Identifier(col.name)
                coltype = get_postgres_type(col.col_type)
                col_sql = sql.SQL(
                    '{colname} {coltype} {opt_primary_key}'
                ).format(
                    colname = colname,
                    coltype = coltype if not col.serial else sql.SQL('SERIAL'),
                    opt_primary_key = sql.SQL('PRIMARY KEY') if i == 0 else sql.SQL('')
                )
                column_sqls.append(col_sql)

            # Create the table
            q = sql.SQL(
                'CREATE TABLE IF NOT EXISTS {table} ({cols});'
            ).format(
                table = table,
                cols = sql.SQL(', ').join(column_sqls),
            )
            print(f' - {q.as_string()}')
            c.execute(q)

            # Create index and unique constraint on index columns
            if len(entrytype.get_index_indices()) > 0:
                index_cols: list[sql.Composable] = list()
                for col in entrytype.get_index_columns():
                    if col.index_hash:
                        # If column has potential to be long, index the md5 value of the column instead
                        index_cols.append(sql.SQL('md5({colname})').format(colname = sql.Identifier(col.name)))
                    else:
                        index_cols.append(sql.Identifier(col.name))
                q = sql.SQL(
                    'CREATE UNIQUE INDEX IF NOT EXISTS {index_name} ON {table} ({index_cols});'
                ).format(
                    index_name = sql.Identifier(f'{entrytype.get_table_name()}_index'),
                    constraint_name = sql.Identifier(f'{entrytype.get_table_name()}_constraint'),
                    table = table,
                    index_cols = sql.SQL(', ').join(index_cols)
                )
                print(f' - {q.as_string()}')
                c.execute(q)

        # Relation Tables
        # --------------------------------------------------------------------------------------------------------------
        for relationtype in relationtypes:
            if issubclass(relationtype, InverseRelation):
                # Inverse relations are handles in the next step
                continue

            table = relationtype.get_table_name()
            column_sqls: list[sql.Composable] = get_column_sql_terms(relationtype)

            # Table query
            q = sql.SQL(
                'CREATE TABLE {rel_table_name} ({cols} {optional_check});'
            ).format(
                rel_table_name = sql.Identifier(table),
                cols = sql.SQL(', ').join(column_sqls),
                optional_check = (  # If relation is from/to the same table, ensure invariant id_1 <= id_2
                    sql.SQL(', CHECK (recon_id_1 <= recon_id_2)') if relationtype.symmetric
                    else sql.SQL('')
                )
            )
            print(f' - {q.as_string()}')
            c.execute(q)

            # Create unique constraint on (recon_id_1, recon_id_2, *index cols)
            index_cols: list[sql.Composable] = [
                sql.Identifier(col.name) for col in relationtype.get_index_columns()
            ]
            q = sql.SQL(
                'CREATE UNIQUE INDEX {index_name} ON {table} ({index_cols})'
            ).format(
                index_name = sql.Identifier(f'{relationtype.get_table_name()}_index'),
                table = sql.Identifier(relationtype.get_table_name()),
                index_cols = sql.SQL(', ').join(index_cols)
            )
            print(f' - {q.as_string()}')
            c.execute(q)

        # Create views for inverse relations
        # --------------------------------------------------------------------------------------------------------------
        for relationtype in relationtypes:
            if not issubclass(relationtype, InverseRelation):
                # Inverse relations are handles in the next step
                continue

            # Sanity checks
            assert relationtype.inverse_main_relation.target_entrytype is relationtype.source_entrytype
            assert relationtype.inverse_main_relation.source_entrytype is relationtype.target_entrytype

            # Create table as view
            table = relationtype.get_table_name()
            q = sql.SQL(""" 
                CREATE VIEW {view_name} AS
                    SELECT recon_id_1 AS recon_id_2, recon_id_2 AS recon_id_1, {attr_cols}
                    FROM {parent_table_name}
                    ;
            """).format(
                view_name = sql.Identifier(table),
                parent_table_name = sql.Identifier(relationtype.inverse_main_relation.get_table_name()),
                attr_cols = sql.SQL(', ').join(
                    sql.Identifier(attr_col.name) for attr_col in relationtype.get_attribute_columns()
                )
            )
            print(f' - {q.as_string()}')
            c.execute(q)

        # Create corresponding views for ALL relations
        # --------------------------------------------------------------------------------------------------------------
        for relationtype in relationtypes:
            # Create view corresponding to this relation (relation with adjacent entries)
            table = relationtype.get_table_name()
            q = sql.SQL("""
                        CREATE VIEW {view_name} AS
                        SELECT rel.recon_id_1, rel.recon_id_2, {rel_attr_cols} {t1_attr_cols}, {t2_attr_cols}
                        FROM {rel_table} rel
                                 INNER JOIN {table_1} t1 ON t1.recon_id = rel.recon_id_1
                                 INNER JOIN {table_2} t2 ON t2.recon_id = rel.recon_id_2
                        ;
                        """).format(
                view_name = sql.Identifier(f'{table}_v'),
                rel_table = sql.Identifier(table),
                table_1 = sql.Identifier(relationtype.source_entrytype.get_table_name()),
                table_2 = sql.Identifier(relationtype.target_entrytype.get_table_name()),
                rel_attr_cols = sql.SQL(', ').join(
                    sql.SQL('rel.') + sql.Identifier(c.name)
                    + sql.SQL(' ') + sql.Identifier(f'rel_{c.name}')
                    for c in relationtype.get_attribute_columns()
                ) + (sql.SQL(', ') if len(relationtype.get_attribute_columns()) > 0 else sql.SQL('')),
                t1_attr_cols = sql.SQL(', ').join(
                    sql.SQL('t1.') + sql.Identifier(c.name)
                    + sql.SQL(' ') + sql.Identifier(f't1_{c.name}')
                    for c in relationtype.source_entrytype.get_columns(include_recon_id = False)
                ),
                t2_attr_cols = sql.SQL(', ').join(
                    sql.SQL('t2.') + sql.Identifier(c.name)
                    + sql.SQL(' ') + sql.Identifier(f't2_{c.name}')
                    for c in relationtype.target_entrytype.get_columns(include_recon_id = False)
                ),
            )

            print(f' -> {relationtype.get_table_name()}')
            print(f'   - {q.as_string()}')
            c.execute(q)

        # Grant read/write permissions to 'dev' user on all tables
        # --------------------------------------------------------------------------------------------------------------
        for schema_name in {'meta', 'chemrecon'}:
            c.execute(sql.SQL(
                """ GRANT SELECT, INSERT, UPDATE, DELETE 
                    ON ALL TABLES IN SCHEMA {schema_name} 
                    TO {dev_username};
                    GRANT USAGE, SELECT 
                    ON ALL SEQUENCES IN SCHEMA {schema_name} 
                    TO {dev_username};
                """
            ).format(
                schema_name = sql.Identifier(schema_name),
                dev_username = sql.Identifier(db_user_dev)
            ))

        # Grant USAGE, UPDATE on sequences
        # TODO

        # Grant read permissions to 'public' group on all tables
        # --------------------------------------------------------------------------------------------------------------
        for schema_name in {'meta', 'chemrecon'}:
            c.execute(sql.SQL(
                """ GRANT USAGE ON SCHEMA {schema_name} TO {pub_username};
                    GRANT SELECT 
                    ON ALL TABLES IN SCHEMA {schema_name} 
                    TO {pub_username};
                    GRANT SELECT
                    ON ALL SEQUENCES IN SCHEMA {schema_name}
                    TO {pub_username};
                """
            ).format(
                schema_name = sql.Identifier(schema_name),
                pub_username = sql.Identifier('public_group')
            ))

        # Grant USAGE on sequences
        # TODO

        # Commit and finalise
        # --------------------------------------------------------------------------------------------------------------
        print('Database initialized.')
        conn.commit()


# Utility functions
# ----------------------------------------------------------------------------------------------------------------------
def get_column_sql_terms(relationtype: type[Relation], only_attrs: bool = False) -> list[sql.Composable]:
    """ Generate SQL statemetns to create each column
    """
    column_sqls: list[sql.Composable] = list()

    # Source and target columns
    if not only_attrs:
        column_sqls.append(sql.SQL(
            'recon_id_1 int references {table_1}(recon_id)'
        ).format(
            table_1 = sql.Identifier(relationtype.source_entrytype.get_table_name()),
        ))
        column_sqls.append(sql.SQL(
            'recon_id_2 int references {table_2}(recon_id)'
        ).format(
            table_2 = sql.Identifier(relationtype.target_entrytype.get_table_name()),
        ))

    # Attribute columns
    for attrcol in relationtype.get_attribute_columns():
        column_sqls.append(sql.SQL(
            '{colname} {coltype}'
        ).format(
            colname = sql.Identifier(attrcol.name),
            coltype = get_postgres_type(attrcol.col_type)
        ))

    return column_sqls

def get_postgres_type(t: type) -> sql.Composable:
    # Get the postgres type as an SQL string of the given type
    match t:
        case builtins.str:
            return sql.SQL('text')
        case builtins.int:
            return sql.SQL('integer')
        case builtins.float:
            return sql.SQL('float8')
        case builtins.bool:
            return sql.SQL('boolean')
        case GenericAlias():
            # Subscripted type, e.g. list[int] -> 'integer[]'
            origin_type = get_origin(t)
            match origin_type:
                case builtins.list:
                    # List gets translated into postgres array
                    subscript = get_args(t)[0]
                    subscript_type = get_postgres_type(subscript)
                    q =  sql.SQL('{subscript_sql} ARRAY').format(
                        subscript_sql = subscript_type
                    )
                    return q
                case _:
                    raise NotImplementedError('Implement creating table from whatever type this  is.')
        case _:
            # If it is an enum
            if issubclass(t, Enum):
                return sql.SQL('{chemrecon}.{enumname}').format(
                    chemrecon = sql.Identifier('chemrecon'),
                    enumname = sql.Identifier(t.__name__)
                )
            else:
                raise NotImplementedError('!')


# Script functionality
if __name__ == '__main__':
    main()
