from typing import *
from datetime import date, datetime
from dataclasses import dataclass
from sqlalchemy import text
from sqlalchemy.sql import schema
from . import engine, model
from ... import config, logger
from ...tools import CSTYLE


@dataclass
class PgColumn:
    name: str
    nullable: bool
    data_type: str
    column_type: str
    default: Any
    is_identity: bool
    identity_start: Union[None, int]
    identity_increment: Union[None, int]
    identity_cycle: bool

    def default_sql(self) -> str:
        if self.is_identity:
            return "GENERATED BY DEFAULT AS IDENTITY" \
                   f" (INCREMENT BY {self.identity_increment} START WITH {self.identity_start}" \
                   f"{' CYCLE' if self.identity_cycle else ''})"
        return f'DEFAULT {str(self.default)}'

    def create_sql(self) -> str:
        coldef: List[str] = [
            f'"{self.name}"',
            self.data_type
        ]
        if not self.nullable:
            coldef.append('NOT NULL')
        if self.default is not None or self.is_identity:
            coldef.append(self.default_sql())
        return ' '.join(coldef)


@dataclass
class PgKey:
    name: str
    constraint_name: str
    key_type: str
    is_foreign_key: bool
    referenced_table_name: Union[str, None]
    referenced_columns: List[str]
    on_update: Union[None, str]
    on_delete: Union[None, str]
    col_aliases: List[str]


@dataclass
class PgForeignKey:
    constraint_name: str
    model_idx_columns: List[str]
    ref_idx_columns: List[str]
    ref_table: str
    on_update: Union[None, str]
    on_delete: Union[None, str]


@dataclass
class PgUniqueKey:
    constraint_name: str
    col_aliases: List[str]


@dataclass
class PgIndexKey:
    key_name: str
    col_aliases: List[str]
    column_names: List[str]


@dataclass
class PgTable:
    keys: Dict[str, PgKey]
    columns: Dict[str, PgColumn]
    indexes: Dict[str, PgIndexKey]
    foreign_keys: Dict[str, PgForeignKey]
    unique_keys: Dict[str, PgUniqueKey]
    primary_key: List[str]
    pk_keynames: List[str]


class DatabaseMigration:
    def __init__(
            self,
            pg_schema: str = 'public',
            drop_missing_columns: bool = True,
            drop_missing_tables: bool = True
    ):
        self.cn = None
        self.schema: str = pg_schema
        self.drop_missing_columns: bool = bool(drop_missing_columns)
        self.drop_missing_tables: bool = bool(drop_missing_tables)
        self.current_tables: Dict[str, PgTable] = {}
        self.declared_tables: List[schema.Table] = []
        self.mapped_tables: Dict[str, schema.Table] = {}

    async def migrate(self):
        async with engine.engine.connect() as cn:
            self.cn = cn
            self.current_tables = await self.get_tables_current()
            self.declared_tables = model.OrmModel.metadata.sorted_tables
            self.mapped_tables: Dict[str, schema.Table] = {
                t.name: t for t in self.declared_tables
            }

            migrate_arr: List[str] = []

            for t in self.declared_tables:
                tn: str = t.name
                q: Optional[str] = self.alter_table_sql(t) \
                    if tn in self.mapped_tables \
                    else self.create_table_sql(t)
                if not q:
                    continue
                migrate_arr.append(q)

            if self.drop_missing_tables:
                for tn in self.current_tables.keys():
                    if tn in self.mapped_tables.keys():
                        continue
                    migrate_arr.append(f'DROP TABLE "{tn}"')

            if migrate_arr:
                await self.deconstraint()
                [await self.execute(q, log=True) for q in migrate_arr]
                await self.reconstraint()

            await cn.commit()

    async def execute(self, sql: str, params: Optional[dict] = None, log: bool = False):
        if log:
            logger.debug(f"{CSTYLE['bold']}migrate database{CSTYLE['clear']}: executing:\n{sql}")
        params = params if isinstance(params, dict) else {}
        return await self.cn.execute(text(sql), params)

    async def query(self, sql: str, params: Optional[dict] = None):
        return (await self.execute(sql, params)).fetchall()

    async def get_tables_current(self) -> Dict[str, PgTable]:
        tablenames: List[str] = await self.get_tablenames_current()
        return {tn: await self.get_table_current(tn) for tn in tablenames}

    async def get_tablenames_current(self) -> List[str]:
        sql: str = 'SELECT "table_name" AS table_name FROM "information_schema"."tables" WHERE "table_schema" = :ts'
        return [r['table_name'] for r in (await self.query(sql, {'ts': self.schema}))]

    async def get_table_current(self, tablename: str) -> Optional[PgTable]:
        """ Reads the table structure from DBMS of the Model. Reads attributes definitions and keys,
        including indexes, primary keys, foreign keys, unique keys, etc.

        :param tablename:
            The existing database table name.
        :type tablename:
            str

        :returns:
            The existing database table structure.
        :rtype:
            dict
        """

        table: PgTable = PgTable(
            keys={},
            columns={},
            indexes={},
            foreign_keys={},
            unique_keys={},
            primary_key=[],
            pk_keynames=[]
        )

        query: str
        dbr: list

        # Selecting attributes definitions from INFORMATION_SCHEMA, basing on PostgreSQL documentation
        # "34.15. The Information Schema.columns"
        query = (
            'SELECT "column_name", "column_default", "is_nullable", "data_type", "character_maximum_length",'
            ' "character_octet_length", "numeric_precision", "numeric_scale", "datetime_precision", "udt_name",'
            ' "is_identity", "identity_start", "identity_increment", "identity_cycle"'
            ' FROM "information_schema"."columns" WHERE "table_schema"=:ts AND "table_name"=:tn'
        )
        dbr = (await self.query(query, {'ts': self.schema, 'tn': tablename}))
        if not dbr:
            return None

        for r in dbr:
            column_name: str = r['column_name']
            column_def: PgColumn = PgColumn(
                name=column_name,
                nullable=True if str(r['is_nullable']).upper() == "YES" else False,
                data_type=str(r['data_type']).lower(),
                column_type=str(r['udt_name']).upper(),
                default=r['column_default'],
                is_identity=str(r['is_identity']).upper() == 'YES',
                identity_start=int(r['identity_start']) if r['identity_start'] else None,
                identity_increment=int(r['identity_increment']) if r['identity_increment'] else None,
                identity_cycle=str(r['identity_cycle']).upper() == 'YES',
            )
            if column_def.data_type == 'character varying':
                character_maximum_length: Optional[int] = r['character_maximum_length']
                if character_maximum_length:
                    column_def.data_type += f"({r['character_maximum_length']})"
            elif column_def.data_type == 'numeric':
                numeric_precision: Optional[int] = r['numeric_precision']
                numeric_scale: Optional[int] = r['numeric_scale']
                numeric: List[str] = []
                if numeric_precision:
                    numeric.append(str(numeric_precision))
                if numeric_scale:
                    numeric.append(str(numeric_scale))
                if numeric:
                    column_def.data_type += f"({', '.join(numeric)})"

            table.columns[column_name] = column_def

        # Selecting informations about declared keys usage, basing on PostgreSQL documentation
        # "34.30 The Information Schema.key_column_usage"
        query = (
            'SELECT tc."constraint_name" AS constraint_name,'
            ' kcu."column_name" AS column_name,'
            ' tc."constraint_type" AS constraint_type'
            ' FROM "information_schema"."table_constraints" tc'
            ' JOIN "information_schema"."key_column_usage" kcu'
            ' ON kcu."constraint_name" = tc."constraint_name"'
            ' WHERE tc."constraint_schema"=:ts AND tc."table_name"=:tn'
            ' AND UPPER(tc."constraint_type") IN (\'PRIMARY KEY\', \'FOREIGN KEY\', \'UNIQUE\')'
            ' ORDER BY kcu."constraint_name", kcu."position_in_unique_constraint", kcu."ordinal_position"'
        )
        for r in (await self.query(query, {'ts': self.schema, 'tn': tablename})):
            key_name: str = str(r['constraint_name'])
            key_type: str = str(r['constraint_type']).upper()
            if key_type == 'PRIMARY KEY' and r['column_name']:
                # PRIMARY KEY part
                table.primary_key.append(r['column_name'])
                table.pk_keynames.append(r['constraint_name'])
                continue
            if key_name in table.keys:
                # Composite key part
                table.keys[key_name].col_aliases.append(r['column_name'])
                continue
            # Defining a new key
            table.keys[key_name] = PgKey(
                name=key_name,
                constraint_name=key_name,
                key_type=key_type,
                is_foreign_key=(key_type == "FOREIGN KEY"),
                referenced_table_name=None,
                referenced_columns=[],
                on_update=None,
                on_delete=None,
                col_aliases=[r['column_name'], ]
            )

        # Additionally selecting information about referential constraints - foreign keys. Doing it
        # by selecting from REFERENTIAL_CONSTRAINTS basing on PostgreSQL documentation
        # "34.32 The Information Schema.referential_constraints"
        # Note that PostgreSQL gives no posibility to filter referential_constraints table againts
        # specific table name and returns all constraints for the database at a time.
        query = (
            'SELECT "constraint_name", "update_rule", "delete_rule" FROM'
            ' "information_schema"."referential_constraints" WHERE'
            ' "constraint_schema" = :ts'
        )
        for r in (await self.query(query, {'ts': self.schema})):
            fk_name: str = r['constraint_name']
            if fk_name not in table.keys:
                continue
            table.keys[fk_name].is_foreign_key = True
            table.keys[fk_name].on_update = r['update_rule'] or None
            table.keys[fk_name].on_delete = r['delete_rule'] or None

        # Handling got keys
        for name in table.keys:
            # We about to define three general types of keys:
            # a) PRIMARY KEY
            # b) FOREIGN KEY
            # c) UNIQUE KEY
            # INDEX KEY (or just a KEY in terms of PostgreSQL) are stored
            # separatelly, so will not be found in general KEYS structure.
            # The first type (a) - the PRIMARY KEY (even composite)
            # detects in the procedure above and not stored in the general
            # `keys` variable, so this is not a point to work with PK here.
            key: PgKey = table.keys[name]
            is_fk: bool = key.is_foreign_key
            if is_fk:
                # Due PostgreSQL does not stores information about foreign key referenced
                # table and columns in the information schema with key column usage - we have
                # to query for this information additionaly againts every FK constraint.
                query = (
                    'SELECT "table_name", "column_name"'
                    ' FROM "information_schema"."constraint_column_usage"'
                    ' WHERE "constraint_name"=:cn'
                )
                dbr = (await self.query(query, {'cn': key.constraint_name}))
                if dbr:
                    ref_tablename: Optional[str] = None
                    ref_columns: List[str] = []
                    for r in dbr:
                        ref_tablename = str(r['table_name'])
                        ref_columns.append(str(r['column_name']))
                    key.referenced_table_name = ref_tablename
                    key.referenced_columns = ref_columns
                table.foreign_keys[name] = PgForeignKey(
                    constraint_name=key.constraint_name,
                    model_idx_columns=key.col_aliases,
                    ref_idx_columns=key.referenced_columns,
                    ref_table=key.referenced_table_name,
                    on_update=key.on_update,
                    on_delete=key.on_delete
                )
            else:
                table.unique_keys[name] = PgUniqueKey(
                    constraint_name=key.constraint_name,
                    col_aliases=key.col_aliases
                )

        # Selecting information about INDEX KEYs, basing on internal PostgreSQL functionality.
        # Too bad, but this information is not accessible from standard schema (from
        # INFORMATION SCHEMA), so we have to use PostgreSQL specific way only.
        query = (
            "SELECT ix.relname AS relname,"
            " regexp_replace(pg_get_indexdef(indexrelid), '.*\\((.*)\\)', '\\1') AS idx_columns"
            " FROM pg_index i"
            " JOIN pg_class t ON t.oid = i.indrelid"
            " JOIN pg_class ix ON ix.oid = i.indexrelid"
            " WHERE t.relname = :tn"
        )
        for r in (await self.query(query, {'tn': tablename})):
            name: str = r['relname']
            # If the index key is already registered in the `keys` varitable - skipping
            # due this is not an index key.
            if name in table.keys or name in table.pk_keynames:
                continue
            table.indexes[name] = PgIndexKey(
                key_name=name,
                col_aliases=[],
                column_names=[]
            )
            idx_columns: List[str] = r['idx_columns'].split(',')
            for cn in idx_columns:
                table.indexes[name].col_aliases.append(cn)
                table.indexes[name].column_names.append(cn)

        # Accumulating the resulting information about model's table definition
        return table

    async def deconstraint(self) -> None:
        for tn, table in self.current_tables.items():
            constraints: list = \
                list(table.foreign_keys.keys()) \
                + list(table.unique_keys.keys())
            for cn in constraints:
                await self.execute(
                    f'ALTER TABLE "{tn}" DROP CONSTRAINT "{cn}"',
                )

    async def reconstraint(self) -> None:
        for table in self.declared_tables:
            tn: str = table.name

            # Foreign keys
            fkcs: Set[schema.ForeignKeyConstraint] = table.foreign_key_constraints
            if not fkcs:
                continue
            for fkc in fkcs:
                fks: List[schema.ForeignKey] = fkc.elements
                fk_columns: List[schema.Column] = fkc.columns
                fk_refcolumns: List[schema.Column] = [
                    fk.column for fk in fks
                ]
                rt: schema.Table = fkc.referred_table
                rtn: str = rt.name
                cn: str = '_'.join([tn] + [c.name for c in fk_columns] + ['fkey'])
                q_columns: str = ','.join([
                    f'"{c.name}"' for c in fk_columns
                ])
                q_refcolumns: str = ','.join([
                    f'"{c.name}"' for c in fk_refcolumns
                ])
                q_rulesarr: List[str] = []
                if fkc.ondelete:
                    q_rulesarr.append(f"ON DELETE {fkc.ondelete}")
                if fkc.onupdate:
                    q_rulesarr.append(f"ON UPDATE {fkc.onupdate}")
                q_rules: str = ' '.join(q_rulesarr)
                await self.execute(
                    f'ALTER TABLE "{tn}" ADD CONSTRAINT "{cn}"'
                    f' FOREIGN KEY ({q_columns})'
                    f' REFERENCES "{rtn}" ({q_refcolumns}) {q_rules}'
                )

            # Unique keys
            ucs: List[schema.UniqueConstraint] = [
                c for c in table.constraints if isinstance(c, schema.UniqueConstraint)
            ]
            for uc in ucs:
                cn: str = '_'.join([tn] + [c.name for c in uc.columns] + ['key'])
                cns: str = ','.join(f'"{c.name}"' for c in uc.columns)
                await self.execute(
                    f'ALTER TABLE "{tn}" ADD CONSTRAINT "{cn}"'
                    f' UNIQUE ({cns})'
                )

    def create_table_sql(self, table: schema.Table) -> str:
        tn: str = table.name
        qc: List[str] = [self.column_create_sql(c) for c in table.columns]
        pk: schema.PrimaryKeyConstraint = table.primary_key
        pk_columns: str = ','.join([f'"{c.name}"' for c in pk.columns])
        qi: str = ',\n'.join(
            [f"  {c}" for c in qc]
            + [f'  PRIMARY KEY ({pk_columns})']
        )
        return f'CREATE TABLE "{tn}" (\n{qi}\n)'

    def alter_table_sql(self, table: schema.Table) -> Optional[str]:
        tn: str = table.name
        current: PgTable = self.current_tables.get(tn, None)
        if not current:
            return self.create_table_sql(table)

        mapped_columnnames: List[str] = []
        alter_arr: List[str] = []
        for c in table.columns:
            cn: str = c.name
            mapped_columnnames.append(cn)

            if cn in current.columns:
                cc: PgColumn = current.columns[cn]
                declared_create_sql: str = self.column_create_sql(c)
                current_create_sql: str = cc.create_sql()
                if declared_create_sql == current_create_sql:
                    continue
                alter_arr.append(self.column_alter_sql(c))
            else:
                alter_arr.append(f'ADD COLUMN {self.column_create_sql(c)}')

        if self.drop_missing_columns:
            for cn in current.columns.keys():
                if cn in mapped_columnnames:
                    continue
                alter_arr.append(f'DROP COLUMN "{cn}"')

        if alter_arr:
            return f'ALTER TABLE "{tn}"\n  ' + ',\n  '.join(alter_arr)

    def column_create_sql(self, c: schema.Column) -> str:
        coldef: List[str] = [
            f'"{c.name}"',
            self.column_sql_type(c)
        ]

        if not c.nullable:
            coldef.append('NOT NULL')

        column_default: Any = self.column_default(c)
        if column_default is not None:
            if isinstance(column_default, str) and 'IDENTITY' in column_default:
                coldef.append(column_default)
            else:
                coldef.append(f'DEFAULT {column_default}')

        return ' '.join(coldef)

    def column_alter_sql(self, c: schema.Column) -> str:
        if c.identity is not None:
            raise NotImplementedError(
                "Automatic modifying the existing IDENTITY columns (auto-increment primary keyed columns)"
                " has not been implemented yet! Do NOT alter primary key columns!!!"
            )

        column_default: Any = self.column_default(c)

        alters: List[str] = list()
        alters.append(f'TYPE {self.column_sql_type(c)}')
        if c.primary_key and c.identity is not None:
            pass  # TODO
        elif column_default is None:
            alters.append(f'DROP DEFAULT')
        else:
            alters.append(f'SET DEFAULT {column_default}')
        alters.append('SET NOT NULL' if not c.nullable else 'DROP NOT NULL')

        return ', '.join([
            f'ALTER COLUMN "{c.name}" {a}' for a in alters
        ])

    def column_default(self, c: schema.Column) -> Any:
        if c.primary_key and c.identity is not None:
            # If this column is a primary key and serial (auto-increment)
            identity: schema.Identity = c.identity
            return "GENERATED BY DEFAULT AS IDENTITY" \
                   f" (INCREMENT BY {identity.increment} START WITH {identity.start}"\
                   f"{' CYCLE' if identity.cycle else ''})"

        declared_default: Optional[schema.ColumnDefault] = c.default
        if declared_default is None:
            return None

        default_arg: Any = declared_default.arg

        if declared_default.is_callable:
            # If the default value of the column is callable function
            if default_arg == datetime.now:
                # If the callable is NOW - we can set if on the RDBMS side
                return "now()"
            if default_arg == date.today:
                # If the callable is TODAY - we can set if on the RDBMS side
                return "CURRENT_DATE"
            return None

        if declared_default.is_scalar:
            value: Any = declared_default.arg
            if isinstance(value, bool):
                return 'true' if value else 'false'
            if isinstance(value, str):
                return "'" + str(value).replace("'", "\\'") + "'::character varying"
            return declared_default.arg

        return None

    def column_sql_type(self, c: schema.Column) -> str:
        typedef: str = str(c.type).lower()

        if typedef.startswith('varchar'):
            typedef = typedef.replace('varchar', 'character varying')

        elif typedef == 'datetime':
            typedef = 'timestamp'
            if c.type.timezone:
                typedef += ' with time zone'
            else:
                typedef += ' without time zone'

        return typedef

    async def drop_all(self) -> None:
        async with engine.engine.connect() as cn:
            self.cn = cn
            self.current_tables = await self.get_tables_current()
            await self.deconstraint()
            [await self.execute(
                f'DROP TABLE "{tn}"', log=True
            ) for tn in self.current_tables.keys()]
            await cn.commit()


async def migrate() -> None:
    await (DatabaseMigration(
        drop_missing_tables=bool(config.DATABASE.get('migrate', {}).get('drop_missing_tables')),
        drop_missing_columns=bool(config.DATABASE.get('migrate', {}).get('drop_missing_columns'))
    )).migrate()


async def dropall() -> None:
    await (DatabaseMigration()).drop_all()

