# Copyright © 2024, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE.md
import warnings
from os import listdir, PathLike
from os.path import join, isfile
from pathlib import Path

from polaris.network.create.triggers import recreate_network_triggers
from polaris.network.database_connection import get_srid
from polaris.network.starts_logging import logger
from polaris.utils.database.db_utils import commit_and_close
from polaris.utils.database.db_utils import has_table
from polaris.utils.database.migration import Migration, migration_file_pattern
from polaris.utils.database.standard_database import DatabaseType, StandardDatabase
from polaris.utils.optional_deps import check_dependency


class MigrationManager:
    @staticmethod
    def upgrade(db_path: PathLike, db_type: DatabaseType, redo_triggers, max_migration=None):
        migrations_dir = MigrationManager.migration_sql_dir(db_type)
        logger.info(f"Upgrading {db_type} database at location {db_path}")
        with commit_and_close(db_path, spatial=False) as conn:
            old_database = MigrationManager.check_old_database(db_type, conn)
            if not old_database and not MigrationManager.find_applicable_migrations(conn, migrations_dir):
                return

        with commit_and_close(db_path, spatial=True) as conn:
            MigrationManager.apply_all_migrations(conn, migrations_dir, max_migration)

            # Triggers are only available on Supply networks at the moment
            redo_triggers = redo_triggers or (db_type == DatabaseType.Supply)
            if redo_triggers and db_type == DatabaseType.Supply:
                srid = get_srid(conn=conn)
                recreate_network_triggers(conn, srid)

        if old_database:
            MigrationManager.__old_databases(db_type, db_path)

    @staticmethod
    def __old_databases(db_type, db_path):
        check_dependency("shapely")
        from polaris.network.checker.supply_checker import SupplyChecker

        msg = "Database is older than migration manager functionality. Migration may not sufficient."
        warnings.warn(msg)
        logger.error(msg)
        if db_type == DatabaseType.Supply:
            chckr = SupplyChecker(db_path)
            chckr.critical()
            chckr.finish()
            for e in chckr.errors:
                msg = f"Critical error found on the supply: {e}"
                warnings.warn(msg)
                logger.error(msg)

    @staticmethod
    def apply_all_migrations(conn, migrations_dir, max_migration=None):
        StandardDatabase.ensure_required_tables(conn)
        for m in MigrationManager.find_applicable_migrations(conn, migrations_dir, max_migration):
            with conn:
                m.run(conn)

    @staticmethod
    def find_applied_migrations(conn):
        return [e[0] for e in conn.execute("SELECT migration_id from Migrations").fetchall()]

    @staticmethod
    def find_available_migrations(migrations_dir):
        files = [s for s in listdir(migrations_dir) if isfile(join(migrations_dir, s))]
        return sorted([Migration.from_file(join(migrations_dir, f)) for f in files if migration_file_pattern.match(f)])

    @staticmethod
    def find_applicable_migrations(conn, migrations_dir, max_migration=None):
        def in_range(x):
            return max_migration is None or x.migration_id <= max_migration

        applied = Migration.find_applied_migrations(conn)
        available = MigrationManager.find_available_migrations(migrations_dir)
        return [e for e in available if e.migration_id not in applied and in_range(e)]

    @staticmethod
    def migration_sql_dir(db_type):
        dir = Path(__file__).parent.parent.parent
        if db_type == DatabaseType.Supply:
            dir = dir / "network"
        elif db_type == DatabaseType.Demand:
            dir = dir / "demand"
        else:
            raise RuntimeError(f"Unknown db type: {db_type}")
        return dir / "database" / "migrations"

    @staticmethod
    def check_old_database(db_type, conn):
        return not has_table(conn, "Migrations")
