import dataclasses
import datetime as dt
import difflib
import filecmp
import functools
from pathlib import Path

from asyncpg import Connection
from asyncpg.exceptions import UndefinedTableError

from padmy.logs import logs
from padmy.utils import exec_file, pg_dump, exec_psql_file, insert_many
from .utils import get_files, iter_migration_files, MigrationFile

_GET_LATEST_COMMIT_QUERY = """
SELECT file_ts, file_name
FROM public.migration
ORDER BY applied_at DESC, file_ts DESC
LIMIT 1
"""

# PG_GLOBAL_DUMP_DIR = GLOBAL_TMP / 'matching-pg-backup' / 'dump'
SETUP_FILE = Path(__file__).parent / 'db.sql'


class MigrationError(Exception):
    def __init__(self, msg: str, diff: str):
        self.msg = msg
        self.diff = diff


async def migrate_setup(conn: Connection):
    logs.info('Setting up migration tables')
    await exec_file(conn, SETUP_FILE)
    logs.info('Done')


def _verify_migration(
        database: str,
        schemas: list[str],
        migration_id: str,
        up_file: Path,
        down_file: Path,
        dump_dir: Path,
):
    _before_dump, _after_dump = f'{migration_id}-before.sql', f'{migration_id}-after.sql'

    # Dump before
    pg_dump(database, schemas, dump_path=str(dump_dir / _before_dump),
            options=['-E', 'utf8', '--schema-only'])

    logs.debug(f'Applying {up_file.name}')
    exec_psql_file(database, str(up_file))

    logs.debug(f'Applying {down_file.name}')
    exec_psql_file(database, str(down_file))

    # Dump after
    pg_dump(database, schemas, dump_path=str(dump_dir / _after_dump),
            options=['-E', 'utf8', '--schema-only'])

    _before_dump_file, _after_dump_file = dump_dir / _before_dump, dump_dir / _after_dump

    _no_diff = filecmp.cmp(_before_dump_file, _after_dump_file)

    if not _no_diff:
        diff = difflib.unified_diff(_before_dump_file.read_text().split('\n'), _after_dump_file.read_text().split('\n'),
                                    fromfile=_before_dump_file.name,
                                    tofile=_after_dump_file.name)
        raise MigrationError(f'Difference found for migration: {migration_id}',
                             diff='\n'.join(diff))


def migrate_verify(database: str,
                   schemas: list[str],
                   dump_dir: Path,
                   migration_folder: Path,
                   ):
    """
    Verifies that the up/down migration is correct
    """
    _pg_dump = functools.partial(pg_dump,
                                 database=database,
                                 schemas=schemas)

    migration_files = get_files(migration_folder)
    logs.info(f'Found {len(migration_files)} migration files')

    logs.info('Checking...')
    _down_files: list[MigrationFile] = []
    try:
        for i, (_up_file, _down_file) in enumerate(iter_migration_files(migration_files)):
            _verify_migration(database=database,
                              schemas=schemas,
                              migration_id=_up_file.file_id,
                              up_file=_up_file.path,
                              down_file=_down_file.path,
                              dump_dir=dump_dir)
            # Need to be reapplied for the next tests
            exec_psql_file(database, str(_up_file.path))
            _down_files.append(_down_file)
    except MigrationError as e:
        logs.error(e.msg)
        logs.error(e.diff)
        raise e

    # Restoring to initial state
    logs.info('Restoring to initial state')
    for _down_file in reversed(_down_files):
        exec_psql_file(database, str(_down_file.path))
    # else:
    #     if last_down_file:
    #         exec_psql_file(database, str(last_down_file.path))

    logs.info('Everything ok!')


@dataclasses.dataclass
class LatestMigration:
    timestamp: dt.datetime
    file_name: str

    @classmethod
    def load(cls, data: dict):
        return cls(
            timestamp=data['file_ts'],
            file_name=data['file_name']
        )


class NoSetupTableError(Exception):
    pass


async def _get_latest_migration(conn: Connection):
    try:
        resp = await conn.fetchrow(_GET_LATEST_COMMIT_QUERY)
    except UndefinedTableError as e:
        if 'relation "public.migration" does not exist' in str(e):
            raise NoSetupTableError(
                'Could not find table table "public.migration", did you forget to setup the table'
                ' by running "migration setup" ?'
            )
        else:
            raise
    return LatestMigration.load(dict(resp)) if resp else None


async def migrate_up(conn: Connection,
                     folder: Path,
                     *,
                     nb_migrations: int = -1,
                     ):
    """ Migrate from the latest migration applied available in the database
    to the latest one in the `folder` dir.
    """
    latest_migration = await _get_latest_migration(conn)
    logs.info(f'Latest timestamp: {latest_migration.timestamp}'
              if latest_migration
              else 'No previous migration found')
    migration_files = get_files(folder)

    migrations_to_apply = [
        _up_file for _up_file, _ in iter_migration_files(migration_files)
        if latest_migration is None or (
                _up_file.ts >= latest_migration.timestamp and
                _up_file.path.name != latest_migration.file_name
        )
    ]
    if nb_migrations > 0:
        migrations_to_apply = migrations_to_apply[:nb_migrations]

    if not migrations_to_apply:
        logs.info('No migrations to apply')
        return

    logs.info(f'Found {len(migrations_to_apply)} migrations to apply')
    async with conn.transaction():
        for _migration in migrations_to_apply:
            logs.info(f'Running {_migration.path.name}...')
            try:
                await exec_file(conn, _migration.path)
            except Exception as e:
                logs.error(f'Failed to execute migration_id "{_migration.file_id}"')
                raise e
            commit_data = {
                'file_ts': _migration.ts,
                'file_id': _migration.file_id,
                'migration_type': 'up',
                'file_name': _migration.path.name
            }
            await insert_many(conn, 'public.migration', [commit_data])

    logs.info('Done!')


async def migrate_down(conn: Connection,
                       folder: Path,
                       *,
                       nb_migrations: int = -1,
                       # migration_id: str = None
                       ):
    """
    Rollback `nb_migrations` back (default -1 for all available)
    """
    latest_migration = await _get_latest_migration(conn)

    migration_files = get_files(folder, reverse=True)

    rollback_to_apply = [
        _down_file for _, _down_file in iter_migration_files(migration_files)
        if latest_migration is None or (
                _down_file.ts <= latest_migration.timestamp and
                _down_file.path.name != latest_migration.file_name
        )
    ]

    if nb_migrations > 0:
        rollback_to_apply = rollback_to_apply[:nb_migrations]

    nb_files = len(rollback_to_apply)
    if nb_files == 0:
        logs.info('No rollback files to apply')
        return

    logs.info(f'Found {nb_files} rollback files to apply')
    async with conn.transaction():
        for _rollback in rollback_to_apply:
            logs.info(f'Running {_rollback.path.name}...')
            try:
                await exec_file(conn, _rollback.path)
            except Exception as e:
                logs.error(f'Failed to execute migration_id "{_rollback.file_id}"')
                raise e

            commit_data = {
                'file_ts': _rollback.ts,
                'file_id': _rollback.file_id,
                'migration_type': 'down',
                'file_name': _rollback.path.name
            }
            await insert_many(conn, 'public.migration', [commit_data])

    logs.info('Done!')
