from typing import Callable, Dict, Optional

from flask import g

from .constants import ReplicaType, PolicyType
from .policies import get_policy


class Replication:
    """Main class which is used to manage requests between replicated databases."""

    def __init__(self, app=None, config=None):
        if not (config is None or isinstance(config, dict)):
            raise ValueError("`config` must be an instance of dict or None")

        self.config = config

        if app is not None:
            self.init_app(app, config)

    def init_app(self, app, config: Optional[Dict] = None):
        """This is used to initialize replication with your app object
        :param app - application object.
        :param config - optional configuration dictionary.
        """
        if not (config is None or isinstance(config, dict)):
            raise ValueError("`config` must be an instance of dict or None")
        app.extensions.setdefault("replicated", self)

        base_config = app.config.copy()
        if self.config:
            base_config.update(self.config)
        if config:
            base_config.update(config)
        config = base_config
        if not config.get('ALLOW_DB_REPLICATION'):
            return

        binds = config.get('SQLALCHEMY_BINDS', {})
        replicas_names = [v.value for v in ReplicaType.__members__.values()]
        if set(replicas_names) != set(binds.keys()):
            raise ValueError(f"SQLALCHEMY_BINDS should contain replicas for {','.join(replicas_names)}.")

        policy_type_name = config.get('DB_POLICY_TYPE', PolicyType.DEFAULT_MASTER.value)
        try:
            policy_type = PolicyType(policy_type_name)
        except ValueError:
            raise ValueError(
                f"Got unexpected value {policy_type_name} for DB_POLICY_TYPE, allowed options are "
                f"{','.join([v.value for v in PolicyType.__members__.values()])}")
        policy = get_policy(policy_type)
        app.before_request(policy.set_replica)
        self.configure_engine_getter(app)

    @staticmethod
    def configure_engine_getter(app):
        db = app.extensions['sqlalchemy'].db
        _get_engine = db.get_engine

        def get_routing_engine(app=app, bind=None):
            if bind is None and "_replica_type" in g:
                bind = g._replica_type.value
            return _get_engine(app, bind)

        db.get_engine = get_routing_engine


def set_replica(type: ReplicaType) -> Callable:
    """Decorator for flask view.
    Use to set specific replica for view regardless of the policy.
    """
    def wrapper(f):
        f.set_replica = type
        return f
    return wrapper