import os
import click
import warnings
import traceback
from dataclasses import asdict

import sqlalchemy
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy_utils.functions import database_exists, create_database

# Alembic migrations
from alembic.script import ScriptDirectory
from alembic import command as alembic_command
from alembic.config import Config as AlembicConfig
from alembic.runtime.migration import MigrationContext

from wurker.settings import conf


def make_engine(db_conf:object=None, echo:bool=False, connect_args:dict={}, pool_recycle:int=600, **kwargs) -> object:
    """See the list below for required elements of kwargs."""

    if db_conf and not kwargs:
        kwargs = {k:v for k,v in asdict(db_conf).items() if v}

    # File URL for sqlite
    if "url" in kwargs and kwargs["url"]:
        return create_engine(
            kwargs["url"],
            connect_args=connect_args,
            execution_options={"isolation_level": "AUTOCOMMIT"}
        )

    # OR build your own DSN and pass it in. Or pass the pieces in and build it here
    if "dsn" in kwargs and kwargs["dsn"]:
        dsn = kwargs["dsn"]
    else:
        req_db_keys = ["dbdriver", "username", "password", "hostname", "port", "database",]
        if any([k not in kwargs or kwargs[k] is None or kwargs[k]=='' for k in req_db_keys]):
            raise Exception("Missing one or more required database settings")

        dsn = "{dbdriver}://{username}:{password}@{hostname}:{port}/{database}".format(**kwargs)

    pool_recycle = -1
    if "dbdriver" in kwargs and kwargs["dbdriver"] == "mysql":
        pool_recycle = pool_recycle
        dsn += '?charset=utf8'

    engine = None
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        engine = create_engine(
            dsn,
            connect_args=connect_args,
            pool_recycle=pool_recycle,
            echo=echo,
            execution_options={"isolation_level": "AUTOCOMMIT"}
        )

    if engine is None:
        raise Exception(f"Unable to connect to config DB: {kwargs['hostname']}")

    # If using SQLite, enable foreign key support via PRAGMA
    if engine.dialect.name == "sqlite":
        @event.listens_for(engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            cursor = dbapi_connection.cursor()
            cursor.execute("PRAGMA foreign_keys=ON")
            cursor.close()

    return engine


def make_session(db_conf:object, **kwargs) -> object:
    """Runs connect() then returns a session object."""

    engine = make_engine(db_conf=db_conf, **kwargs)
    session_factory = sessionmaker(bind=engine)
    Session = scoped_session(session_factory)

    return Session


def schema_exists(db_conf:object, base:object, return_bool:bool=False, **kwargs):
    """Checks that the Wurker schema and all it's tables exist. Raises exception if not."""

    engine = make_engine(db_conf=db_conf, **kwargs)

    try:
        if not database_exists(engine.url):
            if return_bool:
                return False
            raise Exception(f"Database does not exist: '{engine.url}'")
    except Exception as ex:
        if return_bool:
            return False
        raise ex

    for table in base.metadata.sorted_tables:
        if not sqlalchemy.inspect(engine).has_table(table.name):
            if return_bool:
                return False
            raise Exception(f"Table does not exist: '{table}'")

    return True


def is_database_up_to_date(engine, alembic_cfg):
    """Check if the database is already migrated to the latest version."""

    conn = engine.connect()
    context = MigrationContext.configure(conn)
    current_rev = context.get_current_revision()

    script = ScriptDirectory.from_config(alembic_cfg)
    head_rev = script.get_current_head()

    return current_rev == head_rev


def migrate(db_conf: object, cmd: str = "upgrade", revision: str = "") -> bool:
    """Alembic migrate up or down."""

    if not db_conf.dsn and not db_conf.url:
        raise Exception("Migration requires either DBConf.dsn or DBConf.url to be non-empty")

    if cmd not in ["upgrade", "downgrade"]:
        raise Exception(f"Unknown migration command: '{cmd}'")

    revision = revision or ("head" if cmd == "upgrade" else "base")

    alembic_ini = os.path.join(conf.etc_root, "alembic.ini")
    script_location = os.path.join(conf.app_root, "migrations")

    url = db_conf.dsn if db_conf.dsn else db_conf.url

    try:
        engine = create_engine(url)
        if not database_exists(engine.url):
            create_database(engine.url)

        alembic_cfg = AlembicConfig(alembic_ini)
        alembic_cfg.set_main_option("script_location", script_location)
        alembic_cfg.set_main_option("sqlalchemy.url", url)

        if cmd == "upgrade":
            if not is_database_up_to_date(engine, alembic_cfg):
                alembic_command.upgrade(alembic_cfg, revision)
            else:
                click.echo("Database is already up to date. Skipping migration.")
        else:
            alembic_command.downgrade(alembic_cfg, revision)

    except Exception:
        traceback.print_exc()
        return False

    return True


def create_schema(db_conf:object, base:object, **kwargs) -> bool:
    """Installs the database schema if it does not exist."""

    engine = make_engine(db_conf=db_conf, **kwargs)

    if not database_exists(engine.url):
        create_database(engine.url)

    if not database_exists(engine.url):
        return False

    # Schema exists, so now create tables
    base.metadata.create_all(engine)
    return True
