import os
import logging
from flask_appbuilder.security.manager import AUTH_OAUTH
from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride

log = logging.getLogger(__name__)

AUTH_TYPE = AUTH_OAUTH
AUTH_USER_REGISTRATION = True
AUTH_ROLES_SYNC_AT_LOGIN = True
AUTH_USER_REGISTRATION_ROLE = "Viewer"

# Keycloak OIDC configuration
KEYCLOAK_HOST = "{{ .Env.KEYCLOAK_HOST }}"
KEYCLOAK_REALM = "{{ .Env.KEYCLOAK_REALM }}"
OIDC_ISSUER = f"https://{KEYCLOAK_HOST}/realms/{KEYCLOAK_REALM}"

# OAuth Providers configuration
OAUTH_PROVIDERS = [{
    'name': 'keycloak',
    'icon': 'fa-key',
    'token_key': 'access_token',
    'remote_app': {
        'client_id': os.environ.get('AIRFLOW_OAUTH_CLIENT_ID', ''),
        'client_secret': os.environ.get('AIRFLOW_OAUTH_CLIENT_SECRET', ''),
        'server_metadata_url': f'{OIDC_ISSUER}/.well-known/openid-configuration',
        'api_base_url': f'{OIDC_ISSUER}/protocol/openid-connect',
        'access_token_url': f'{OIDC_ISSUER}/protocol/openid-connect/token',
        'authorize_url': f'{OIDC_ISSUER}/protocol/openid-connect/auth',
        'request_token_url': None,
        'client_kwargs': {
            'scope': 'openid profile email'
        }
    }
}]

# Role mappings from Keycloak to Airflow
AUTH_ROLES_MAPPING = {
    "airflow_admin": ["Admin"],
    "airflow_op": ["Op"],
    "airflow_user": ["User"],
    "airflow_viewer": ["Viewer"]
}

# Security Manager Override
class KeycloakSecurityManager(FabAirflowSecurityManagerOverride):
    """Custom Security Manager for Keycloak integration"""

    def __init__(self, appbuilder):
        super().__init__(appbuilder)

SECURITY_MANAGER_CLASS = KeycloakSecurityManager