from sqlalchemy.exc import IntegrityError, MultipleResultsFound, NoResultFound
from sqlalchemy.orm import sessionmaker
from werkzeug.security import check_password_hash, generate_password_hash

from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import (
    INVALID_STATE,
    RESOURCE_ALREADY_EXISTS,
    RESOURCE_DOES_NOT_EXIST,
)
from mlflow.server.auth.db import utils as dbutils
from mlflow.server.auth.db.models import (
    SqlExperimentPermission,
    SqlGatewayEndpointPermission,
    SqlGatewayModelDefinitionPermission,
    SqlGatewaySecretPermission,
    SqlRegisteredModelPermission,
    SqlScorerPermission,
    SqlUser,
)
from mlflow.server.auth.entities import (
    ExperimentPermission,
    GatewayEndpointPermission,
    GatewayModelDefinitionPermission,
    GatewaySecretPermission,
    RegisteredModelPermission,
    ScorerPermission,
    User,
)
from mlflow.server.auth.permissions import _validate_permission
from mlflow.store.db.utils import _get_managed_session_maker, create_sqlalchemy_engine_with_retry
from mlflow.utils.uri import extract_db_type_from_uri
from mlflow.utils.validation import _validate_password, _validate_username


class SqlAlchemyStore:
    def init_db(self, db_uri):
        self.db_uri = db_uri
        self.db_type = extract_db_type_from_uri(db_uri)
        self.engine = create_sqlalchemy_engine_with_retry(db_uri)
        dbutils.migrate_if_needed(self.engine, "head")
        SessionMaker = sessionmaker(bind=self.engine)
        self.ManagedSessionMaker = _get_managed_session_maker(SessionMaker, self.db_type)

    def authenticate_user(self, username: str, password: str) -> bool:
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username)
                return check_password_hash(user.password_hash, password)
            except MlflowException:
                return False

    def create_user(self, username: str, password: str, is_admin: bool = False) -> User:
        _validate_username(username)
        _validate_password(password)
        pwhash = generate_password_hash(password)
        with self.ManagedSessionMaker() as session:
            try:
                user = SqlUser(username=username, password_hash=pwhash, is_admin=is_admin)
                session.add(user)
                session.flush()
                return user.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"User (username={username}) already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                ) from e

    @staticmethod
    def _get_user(session, username: str) -> SqlUser:
        try:
            return session.query(SqlUser).filter(SqlUser.username == username).one()
        except NoResultFound:
            raise MlflowException(
                f"User with username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple users with username={username}",
                INVALID_STATE,
            )

    def has_user(self, username: str) -> bool:
        with self.ManagedSessionMaker() as session:
            return session.query(SqlUser).filter(SqlUser.username == username).first() is not None

    def get_user(self, username: str) -> User:
        with self.ManagedSessionMaker() as session:
            return self._get_user(session, username).to_mlflow_entity()

    def list_users(self) -> list[User]:
        with self.ManagedSessionMaker() as session:
            users = session.query(SqlUser).all()
            return [u.to_mlflow_entity() for u in users]

    def update_user(
        self, username: str, password: str | None = None, is_admin: bool | None = None
    ) -> User:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username)
            if password is not None:
                pwhash = generate_password_hash(password)
                user.password_hash = pwhash
            if is_admin is not None:
                user.is_admin = is_admin
            return user.to_mlflow_entity()

    def delete_user(self, username: str):
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username)
            session.delete(user)

    def create_experiment_permission(
        self, experiment_id: str, username: str, permission: str
    ) -> ExperimentPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlExperimentPermission(
                    experiment_id=experiment_id, user_id=user.id, permission=permission
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Experiment permission (experiment_id={experiment_id}, username={username}) "
                    f"already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                )

    def _get_experiment_permission(
        self, session, experiment_id: str, username: str
    ) -> SqlExperimentPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlExperimentPermission)
                .filter(
                    SqlExperimentPermission.experiment_id == experiment_id,
                    SqlExperimentPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Experiment permission with experiment_id={experiment_id} and "
                f"username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple experiment permissions with experiment_id={experiment_id} "
                f"and username={username}",
                INVALID_STATE,
            )

    def get_experiment_permission(self, experiment_id: str, username: str) -> ExperimentPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_experiment_permission(
                session, experiment_id, username
            ).to_mlflow_entity()

    def list_experiment_permissions(self, username: str) -> list[ExperimentPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlExperimentPermission)
                .filter(SqlExperimentPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_experiment_permission(
        self, experiment_id: str, username: str, permission: str
    ) -> ExperimentPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_experiment_permission(session, experiment_id, username)
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_experiment_permission(self, experiment_id: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_experiment_permission(session, experiment_id, username)
            session.delete(perm)

    def create_registered_model_permission(
        self, name: str, username: str, permission: str
    ) -> RegisteredModelPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlRegisteredModelPermission(
                    name=name, user_id=user.id, permission=permission
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Registered model permission (name={name}, username={username}) "
                    f"already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                )

    def _get_registered_model_permission(
        self, session, name: str, username: str
    ) -> SqlRegisteredModelPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlRegisteredModelPermission)
                .filter(
                    SqlRegisteredModelPermission.name == name,
                    SqlRegisteredModelPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Registered model permission with name={name} and username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple registered model permissions with name={name} "
                f"and username={username}",
                INVALID_STATE,
            )

    def get_registered_model_permission(
        self, name: str, username: str
    ) -> RegisteredModelPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_registered_model_permission(session, name, username).to_mlflow_entity()

    def list_registered_model_permissions(self, username: str) -> list[RegisteredModelPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlRegisteredModelPermission)
                .filter(SqlRegisteredModelPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_registered_model_permission(
        self, name: str, username: str, permission: str
    ) -> RegisteredModelPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_registered_model_permission(session, name, username)
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_registered_model_permission(self, name: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_registered_model_permission(session, name, username)
            session.delete(perm)

    def rename_registered_model_permissions(self, old_name: str, new_name: str):
        with self.ManagedSessionMaker() as session:
            perms = (
                session.query(SqlRegisteredModelPermission)
                .filter(SqlRegisteredModelPermission.name == old_name)
                .all()
            )
            for perm in perms:
                perm.name = new_name

    def create_scorer_permission(
        self, experiment_id: str, scorer_name: str, username: str, permission: str
    ) -> ScorerPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlScorerPermission(
                    experiment_id=experiment_id,
                    scorer_name=scorer_name,
                    user_id=user.id,
                    permission=permission,
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Scorer permission (experiment_id={experiment_id}, scorer_name={scorer_name}, "
                    f"username={username}) already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                ) from e

    def _get_scorer_permission(
        self, session, experiment_id: str, scorer_name: str, username: str
    ) -> SqlScorerPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlScorerPermission)
                .filter(
                    SqlScorerPermission.experiment_id == experiment_id,
                    SqlScorerPermission.scorer_name == scorer_name,
                    SqlScorerPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Scorer permission with experiment_id={experiment_id}, "
                f"scorer_name={scorer_name}, and username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple scorer permissions with experiment_id={experiment_id}, "
                f"scorer_name={scorer_name}, and username={username}",
                INVALID_STATE,
            )

    def get_scorer_permission(
        self, experiment_id: str, scorer_name: str, username: str
    ) -> ScorerPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_scorer_permission(
                session, experiment_id, scorer_name, username
            ).to_mlflow_entity()

    def list_scorer_permissions(self, username: str) -> list[ScorerPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlScorerPermission)
                .filter(SqlScorerPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_scorer_permission(
        self, experiment_id: str, scorer_name: str, username: str, permission: str
    ) -> ScorerPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_scorer_permission(session, experiment_id, scorer_name, username)
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_scorer_permission(self, experiment_id: str, scorer_name: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_scorer_permission(session, experiment_id, scorer_name, username)
            session.delete(perm)

    def delete_scorer_permissions_for_scorer(self, experiment_id: str, scorer_name: str):
        with self.ManagedSessionMaker() as session:
            session.query(SqlScorerPermission).filter(
                SqlScorerPermission.experiment_id == experiment_id,
                SqlScorerPermission.scorer_name == scorer_name,
            ).delete()

    def create_gateway_secret_permission(
        self, secret_id: str, username: str, permission: str
    ) -> GatewaySecretPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlGatewaySecretPermission(
                    secret_id=secret_id, user_id=user.id, permission=permission
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Gateway secret permission (secret_id={secret_id}, username={username}) "
                    f"already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                ) from e

    def _get_gateway_secret_permission(
        self, session, secret_id: str, username: str
    ) -> SqlGatewaySecretPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlGatewaySecretPermission)
                .filter(
                    SqlGatewaySecretPermission.secret_id == secret_id,
                    SqlGatewaySecretPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Gateway secret permission with secret_id={secret_id} and "
                f"username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple gateway secret permissions with secret_id={secret_id} "
                f"and username={username}",
                INVALID_STATE,
            )

    def get_gateway_secret_permission(
        self, secret_id: str, username: str
    ) -> GatewaySecretPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_gateway_secret_permission(
                session, secret_id, username
            ).to_mlflow_entity()

    def list_gateway_secret_permissions(self, username: str) -> list[GatewaySecretPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlGatewaySecretPermission)
                .filter(SqlGatewaySecretPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_gateway_secret_permission(
        self, secret_id: str, username: str, permission: str
    ) -> GatewaySecretPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_secret_permission(session, secret_id, username)
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_gateway_secret_permission(self, secret_id: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_secret_permission(session, secret_id, username)
            session.delete(perm)

    def delete_gateway_secret_permissions_for_secret(self, secret_id: str):
        with self.ManagedSessionMaker() as session:
            session.query(SqlGatewaySecretPermission).filter(
                SqlGatewaySecretPermission.secret_id == secret_id,
            ).delete()

    def create_gateway_endpoint_permission(
        self, endpoint_id: str, username: str, permission: str
    ) -> GatewayEndpointPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlGatewayEndpointPermission(
                    endpoint_id=endpoint_id, user_id=user.id, permission=permission
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Gateway endpoint permission (endpoint_id={endpoint_id}, username={username}) "
                    f"already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                ) from e

    def _get_gateway_endpoint_permission(
        self, session, endpoint_id: str, username: str
    ) -> SqlGatewayEndpointPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlGatewayEndpointPermission)
                .filter(
                    SqlGatewayEndpointPermission.endpoint_id == endpoint_id,
                    SqlGatewayEndpointPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Gateway endpoint permission with endpoint_id={endpoint_id} and "
                f"username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple gateway endpoint permissions with endpoint_id={endpoint_id} "
                f"and username={username}",
                INVALID_STATE,
            )

    def get_gateway_endpoint_permission(
        self, endpoint_id: str, username: str
    ) -> GatewayEndpointPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_gateway_endpoint_permission(
                session, endpoint_id, username
            ).to_mlflow_entity()

    def list_gateway_endpoint_permissions(self, username: str) -> list[GatewayEndpointPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlGatewayEndpointPermission)
                .filter(SqlGatewayEndpointPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_gateway_endpoint_permission(
        self, endpoint_id: str, username: str, permission: str
    ) -> GatewayEndpointPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_endpoint_permission(session, endpoint_id, username)
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_gateway_endpoint_permission(self, endpoint_id: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_endpoint_permission(session, endpoint_id, username)
            session.delete(perm)

    def delete_gateway_endpoint_permissions_for_endpoint(self, endpoint_id: str):
        with self.ManagedSessionMaker() as session:
            session.query(SqlGatewayEndpointPermission).filter(
                SqlGatewayEndpointPermission.endpoint_id == endpoint_id,
            ).delete()

    def create_gateway_model_definition_permission(
        self, model_definition_id: str, username: str, permission: str
    ) -> GatewayModelDefinitionPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            try:
                user = self._get_user(session, username=username)
                perm = SqlGatewayModelDefinitionPermission(
                    model_definition_id=model_definition_id, user_id=user.id, permission=permission
                )
                session.add(perm)
                session.flush()
                return perm.to_mlflow_entity()
            except IntegrityError as e:
                raise MlflowException(
                    f"Gateway model definition permission "
                    f"(model_definition_id={model_definition_id}, username={username}) "
                    f"already exists. Error: {e}",
                    RESOURCE_ALREADY_EXISTS,
                ) from e

    def _get_gateway_model_definition_permission(
        self, session, model_definition_id: str, username: str
    ) -> SqlGatewayModelDefinitionPermission:
        try:
            user = self._get_user(session, username=username)
            return (
                session.query(SqlGatewayModelDefinitionPermission)
                .filter(
                    SqlGatewayModelDefinitionPermission.model_definition_id == model_definition_id,
                    SqlGatewayModelDefinitionPermission.user_id == user.id,
                )
                .one()
            )
        except NoResultFound:
            raise MlflowException(
                f"Gateway model definition permission with "
                f"model_definition_id={model_definition_id} and username={username} not found",
                RESOURCE_DOES_NOT_EXIST,
            )
        except MultipleResultsFound:
            raise MlflowException(
                f"Found multiple gateway model definition permissions with "
                f"model_definition_id={model_definition_id} and username={username}",
                INVALID_STATE,
            )

    def get_gateway_model_definition_permission(
        self, model_definition_id: str, username: str
    ) -> GatewayModelDefinitionPermission:
        with self.ManagedSessionMaker() as session:
            return self._get_gateway_model_definition_permission(
                session, model_definition_id, username
            ).to_mlflow_entity()

    def list_gateway_model_definition_permissions(
        self, username: str
    ) -> list[GatewayModelDefinitionPermission]:
        with self.ManagedSessionMaker() as session:
            user = self._get_user(session, username=username)
            perms = (
                session.query(SqlGatewayModelDefinitionPermission)
                .filter(SqlGatewayModelDefinitionPermission.user_id == user.id)
                .all()
            )
            return [p.to_mlflow_entity() for p in perms]

    def update_gateway_model_definition_permission(
        self, model_definition_id: str, username: str, permission: str
    ) -> GatewayModelDefinitionPermission:
        _validate_permission(permission)
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_model_definition_permission(
                session, model_definition_id, username
            )
            perm.permission = permission
            return perm.to_mlflow_entity()

    def delete_gateway_model_definition_permission(self, model_definition_id: str, username: str):
        with self.ManagedSessionMaker() as session:
            perm = self._get_gateway_model_definition_permission(
                session, model_definition_id, username
            )
            session.delete(perm)

    def delete_gateway_model_definition_permissions_for_model_definition(
        self, model_definition_id: str
    ):
        with self.ManagedSessionMaker() as session:
            session.query(SqlGatewayModelDefinitionPermission).filter(
                SqlGatewayModelDefinitionPermission.model_definition_id == model_definition_id,
            ).delete()
