"""Role Assignment Manager

Orchestrates automatic role assignment after service publication.
Extracts user information from JWT token and assigns science_module role.

Supports two modes:
1. Direct Keycloak access (requires admin credentials in SDK)
2. SDK Services (credentials stay server-side, more secure)

Author: Balázs Milán (milan.balazs@uni-obuda.hu)
"""

from typing import Optional, Dict, Any
import logging
import jwt
import requests
from datetime import datetime, timezone

from .keycloak_manager import KeycloakManager, get_keycloak_manager
from ..config.keycloak_admin_config import KeycloakAdminConfig
from ..config.docker_config import get_docker_config
from ..exceptions import (
    RoleAssignmentError,
    TokenValidationError,
    BlackAntException
)


logger = logging.getLogger(__name__)


class RoleAssignmentManager:
    """Manager for automatic role assignment after service publication.

    Handles the workflow of:
    1. Extracting user ID from JWT token
    2. Checking if science_module role exists
    3. Assigning role if missing
    4. Audit logging

    Supports two modes:
    - Direct Keycloak: Uses KeycloakManager with admin credentials
    - SDK Services: Delegates to backend, credentials stay server-side

    Args:
        keycloak_manager: Optional KeycloakManager instance
        role_name: Role to assign (default: "science_module")

    Example:
        >>> manager = RoleAssignmentManager()
        >>> success = manager.assign_role_after_publish(
        ...     user_token="eyJhbGciOiJSUzI1NiIs...",
        ...     service_name="my-calculation"
        ... )
    """

    DEFAULT_ROLE = "science_module"

    def __init__(
        self,
        keycloak_manager: Optional[KeycloakManager] = None,
        role_name: str = DEFAULT_ROLE
    ):
        """Initialize role assignment manager.

        Args:
            keycloak_manager: Optional KeycloakManager (creates new if None)
            role_name: Role to assign (default: "science_module")
        """
        self.keycloak_manager = keycloak_manager or get_keycloak_manager()
        self.role_name = role_name
        self.logger = logger

        # Check SDK Services configuration
        try:
            self._docker_config = get_docker_config()
            self._use_sdk_services = (
                self._docker_config.sdk_services.use_for_roles and
                self._docker_config.sdk_services.url
            )
        except Exception:
            self._docker_config = None
            self._use_sdk_services = False

    def extract_user_id_from_token(self, token: str) -> str:
        """Extract Keycloak user UUID from JWT token.

        Args:
            token: JWT token (Bearer token from authentication)

        Returns:
            User UUID string

        Raises:
            TokenValidationError: If token is invalid or user ID not found

        Example:
            >>> user_id = manager.extract_user_id_from_token(token)
            >>> print(user_id)
            'a1b2c3d4-e5f6-7890-abcd-ef1234567890'
        """
        try:
            # Decode JWT without verification (we trust our own tokens)
            # In production, you should verify signature with public key
            decoded = jwt.decode(
                token,
                options={"verify_signature": False}  # Trust our own tokens
            )

            # Extract user ID from 'sub' claim (standard JWT claim)
            user_id = decoded.get('sub')
            if not user_id:
                raise TokenValidationError(
                    "Token does not contain 'sub' claim (user ID)"
                )

            # Additional info for logging
            username = decoded.get('preferred_username', 'unknown')
            email = decoded.get('email', 'unknown')

            self.logger.debug(
                f"Extracted user ID from token: {user_id} "
                f"(username: {username}, email: {email})"
            )

            return user_id

        except jwt.InvalidTokenError as e:
            self.logger.error(f"Invalid JWT token: {e}")
            raise TokenValidationError(f"Invalid JWT token: {e}")
        except Exception as e:
            self.logger.error(f"Failed to extract user ID from token: {e}")
            raise TokenValidationError(f"Token parsing failed: {e}")

    def assign_role_after_publish(
        self,
        user_token: str,
        service_name: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Assign science_module role to user after service publication.

        This is the main entry point called by ServiceRegistry after
        successful Docker image push to registry.

        If SDK Services is configured, delegates to backend.
        Otherwise, uses direct Keycloak access with admin credentials.

        Args:
            user_token: JWT token of the user who published the service
            service_name: Name of the published service
            metadata: Optional metadata for audit logging

        Returns:
            True if role was assigned or already exists, False on failure

        Raises:
            RoleAssignmentError: If assignment fails critically

        Example:
            >>> success = manager.assign_role_after_publish(
            ...     user_token=jwt_token,
            ...     service_name="sentiment-analysis",
            ...     metadata={"image": "env.blackant.app/..."}
            ... )
        """
        # Check if SDK Services should handle role assignment
        if self._use_sdk_services:
            return self._assign_role_via_sdk_services(
                user_token, service_name, metadata
            )

        # Use direct Keycloak access
        return self._assign_role_direct(user_token, service_name, metadata)

    def _assign_role_via_sdk_services(
        self,
        user_token: str,
        service_name: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Assign role via SDK Services endpoint.

        Delegates role assignment to sdk-services backend,
        keeping admin credentials server-side.

        Args:
            user_token: JWT token of the user
            service_name: Name of the published service
            metadata: Optional metadata

        Returns:
            True if successful, False otherwise
        """
        try:
            # Extract user ID from token
            user_id = self.extract_user_id_from_token(user_token)

            assign_url = self._docker_config.sdk_services.roles_assign_url
            self.logger.info(f"Assigning role via SDK Services: {assign_url}")

            # Call SDK Services endpoint
            response = requests.post(
                assign_url,
                json={
                    "user_id": user_id,
                    "role_name": self.role_name
                },
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {user_token}"
                },
                timeout=self._docker_config.sdk_services.timeout,
                verify=False  # For dev environments
            )

            if response.status_code == 200:
                result = response.json()
                if result.get("success"):
                    self.logger.info(
                        f"Role '{self.role_name}' assigned to user {user_id} "
                        f"via SDK Services"
                    )
                    self._audit_log(
                        user_id=user_id,
                        service_name=service_name,
                        action="role_assigned_via_sdk_services",
                        result="success",
                        metadata=metadata
                    )
                    return True
                else:
                    error_msg = result.get("message", "Unknown error")
                    self.logger.error(f"SDK Services role assign failed: {error_msg}")
                    return False
            else:
                self.logger.error(
                    f"SDK Services role assign failed with status {response.status_code}"
                )
                return False

        except TokenValidationError as e:
            self.logger.error(f"Token validation failed: {e}")
            return False
        except Exception as e:
            self.logger.error(f"SDK Services role assign error: {e}")
            return False

    def _assign_role_direct(
        self,
        user_token: str,
        service_name: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Assign role directly via Keycloak Manager.

        Uses local admin credentials from configuration.

        Args:
            user_token: JWT token of the user
            service_name: Name of the published service
            metadata: Optional metadata

        Returns:
            True if successful, False otherwise
        """
        try:
            # Step 1: Extract user ID from JWT token
            self.logger.info(
                f"Starting role assignment for service '{service_name}'"
            )

            user_id = self.extract_user_id_from_token(user_token)

            default = True
            if not default:
                self.role_name=service_name;

                #check if role exists
                if not self.keycloak_manager.role_exists(self.role_name):
                    self.keycloak_manager.create_realm_role(self.role_name)

            # Step 2: Check if user already has the role (idempotent)
            has_role = self.keycloak_manager.has_role(user_id, self.role_name)

            if has_role:
                self.logger.info(
                    f"User {user_id} already has role '{self.role_name}', "
                    f"no assignment needed"
                )
                self._audit_log(
                    user_id=user_id,
                    service_name=service_name,
                    action="role_check",
                    result="already_exists",
                    metadata=metadata
                )
                return True

            # Step 3: Assign the role
            self.logger.info(
                f"Assigning role '{self.role_name}' to user {user_id} "
                f"for service '{service_name}'"
            )

            success = self.keycloak_manager.assign_role_to_user(
                user_id=user_id,
                role_name=self.role_name,
                idempotent=True
            )

            if success:
                self.logger.info(
                    f"Successfully assigned role '{self.role_name}' to user {user_id}"
                )
                self._audit_log(
                    user_id=user_id,
                    service_name=service_name,
                    action="role_assigned",
                    result="success",
                    metadata=metadata
                )
                return True
            else:
                self.logger.warning(
                    f"Role assignment returned False for user {user_id}"
                )
                self._audit_log(
                    user_id=user_id,
                    service_name=service_name,
                    action="role_assignment",
                    result="failed",
                    metadata=metadata
                )
                return False

        except TokenValidationError as e:
            # Token parsing error - log and return False (don't raise)
            self.logger.error(
                f"Token validation failed for service '{service_name}': {e}"
            )
            self._audit_log(
                user_id="unknown",
                service_name=service_name,
                action="role_assignment",
                result="token_error",
                error=str(e),
                metadata=metadata
            )
            return False

        except RoleAssignmentError as e:
            # Keycloak communication error - log and return False
            self.logger.error(
                f"Role assignment failed for service '{service_name}': {e}"
            )
            self._audit_log(
                user_id=user_id if 'user_id' in locals() else "unknown",
                service_name=service_name,
                action="role_assignment",
                result="keycloak_error",
                error=str(e),
                metadata=metadata
            )
            return False

        except Exception as e:
            # Unexpected error - log and return False
            self.logger.exception(
                f"Unexpected error during role assignment for service '{service_name}': {e}"
            )
            self._audit_log(
                user_id=user_id if 'user_id' in locals() else "unknown",
                service_name=service_name,
                action="role_assignment",
                result="unexpected_error",
                error=str(e),
                metadata=metadata
            )
            return False

    def _audit_log(
        self,
        user_id: str,
        service_name: str,
        action: str,
        result: str,
        error: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Create audit log entry for role assignment.

        Logs all role assignment operations for security auditing.

        Args:
            user_id: Keycloak user UUID
            service_name: Service name
            action: Action performed (e.g., "role_assigned")
            result: Result of action (e.g., "success", "failed")
            error: Optional error message
            metadata: Optional additional metadata
        """
        audit_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "user_id": user_id,
            "service_name": service_name,
            "action": action,
            "role": self.role_name,
            "result": result,
        }

        if error:
            audit_entry["error"] = error

        if metadata:
            audit_entry["metadata"] = metadata

        # Log as structured JSON for audit trail
        self.logger.info(
            f"AUDIT: {audit_entry}",
            extra={"audit": audit_entry}
        )

        # TODO: Send to centralized audit logging system
        # (e.g., Elasticsearch, CloudWatch, etc.)


# Singleton instance (optional)
_role_assignment_manager: Optional[RoleAssignmentManager] = None


def get_role_assignment_manager() -> RoleAssignmentManager:
    """Get or create RoleAssignmentManager singleton.

    Returns:
        RoleAssignmentManager instance
    """
    global _role_assignment_manager

    if _role_assignment_manager is None:
        _role_assignment_manager = RoleAssignmentManager()

    return _role_assignment_manager
