# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import json
import time
from typing import List,Optional
from google.cloud import iam_admin_v1
from google.cloud.iam_admin_v1 import types
from google.oauth2 import service_account
from google.auth.transport.requests import Request
from google.api_core import exceptions

class ServiceAccountManagerLoggerAdapter(logging.LoggerAdapter):
    """Logger adapter that adds a prefix to all log messages."""
    
    def process(self, msg, kwargs):
        return f"[ServiceAccountManager] {msg}", kwargs

class ServiceAccountManager:
    def __init__(self, project_id: str, logger: logging.Logger, max_retries: int = 3) -> None:
        self.project_id = project_id
        self.client = iam_admin_v1.IAMClient()
        self.logger = ServiceAccountManagerLoggerAdapter(logger, {})
        self.max_retries = max_retries
        self.logger.info(f"Initialized ServiceAccountManager for project: {self.project_id}")

    def _normalize_account_email(self, account_id: str) -> str:
        """
        Normalizes the account identifier to a full email format.
        
        Args:
            account_id (str): The unique identifier or email of the service account.
            
        Returns:
            str: The full service account email address.
        """
        # Handle both account ID and full email formats
        if "@" in account_id and account_id.endswith(".iam.gserviceaccount.com"):
            # account_id is already a full email
            return account_id
        else:
            # account_id is just the account name
            return f"{account_id}@{self.project_id}.iam.gserviceaccount.com"

    def _get_service_accounts(self) -> List[iam_admin_v1.ServiceAccount]:
        """
        Retrieves all service accounts in the specified project.

        Returns:
            List[iam_admin_v1.ServiceAccount]: A list of service account objects.
        """
        request = types.ListServiceAccountsRequest()
        request.name = f"projects/{self.project_id}"

        accounts = self.client.list_service_accounts(request=request)
        self.logger.debug(f"Listed service accounts: {[account.email for account in accounts.accounts]}")
        return list(accounts.accounts)
    
    def _service_account_exists(self, account_id: str) -> bool:
        """
        Checks if a service account with the given account_id exists in the project.

        Args:
            account_id (str): The unique identifier or email of the service account.

        Returns:
            bool: True if the service account exists, False otherwise.
        """
        try:
            self.get_service_account(account_id)
            return True
        except exceptions.NotFound:
            return False
        
    def _service_account_is_enabled(self, account_id: str) -> bool:
        """
        Checks if a service account is enabled.

        Args:
            account_id (str): The unique identifier or email of the service account.

        Returns:
            bool: True if the service account is enabled, False otherwise.
        """
        try:
            service_account = self.get_service_account(account_id)
            return not service_account.disabled
        except exceptions.NotFound:
            self.logger.error(f"Service account {account_id} not found")
            return False

    def create_service_account(self, account_id: str, display_name: Optional[str] = None) -> types.ServiceAccount:
        """
        Creates a service account in the specified project.
        If the service account already exists, returns the existing account (idempotent operation).

        Args:
            account_id (str): The unique identifier for the service account.
            display_name (Optional[str]): A human-readable name for the service account.
        Returns:
            types.ServiceAccount: The created or existing service account object.
        """
        request = types.CreateServiceAccountRequest()
        request.account_id = account_id
        request.name = f"projects/{self.project_id}"

        service_account = types.ServiceAccount()
        service_account.display_name = display_name or account_id
        request.service_account = service_account

        try:
            account = self.client.create_service_account(request=request)

            # Wait for the service account to be created
            delay = 1
            for _ in range(self.max_retries):
                if self._service_account_exists(account_id):
                    break
                time.sleep(delay)
                delay *= 2
            else:
                self.logger.error(f"Service account {account_id} creation timed out after {self.max_retries} retries.")
                raise exceptions.DeadlineExceeded(f"Service account {account_id} creation timed out.")

            self.logger.info(f"Created service account: {account.email}")
            return account
        except exceptions.Conflict:
            existing_account = self.get_service_account(account_id)
            self.logger.info(f"Service account already exists: {existing_account.email}")
            return existing_account
    
    def get_service_account(self, account_id: str) -> types.ServiceAccount:
        """
        Retrieves a service account by its unique identifier or email.

        Args:
            account_id (str): The unique identifier or email of the service account.
        
        Returns:
            types.ServiceAccount: The service account object.
        """
        service_account_email = self._normalize_account_email(account_id)

        request = types.GetServiceAccountRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"

        try:
            service_account = self.client.get_service_account(request=request)
            self.logger.info(f"Retrieved service account: {service_account.email}")
            return service_account
        except exceptions.NotFound:
            self.logger.error(f"Service account {account_id} not found")
            raise
    
    def enable_service_account(self, account_id: str) -> None:
        """
        Enables a service account in the specified project.

        Args:
            account_id (str): The unique identifier or email of the service account to enable.
        """
        service_account_email = self._normalize_account_email(account_id)
        request = types.EnableServiceAccountRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"

        self.client.enable_service_account(request=request)

        # Wait for the service account to be enabled
        delay = 1
        for _ in range(self.max_retries):
            if self._service_account_is_enabled(account_id):
                break
            time.sleep(delay)
            delay *= 2
        else:
            self.logger.error(f"Service account {account_id} enabling timed out after {self.max_retries} retries.")
            raise exceptions.DeadlineExceeded(f"Service account {account_id} enabling timed out.")
        
        self.logger.info(f"Enabled service account: {account_id}")

    def disable_service_account(self, account_id: str) -> None:
        """
        Disables a service account in the specified project.

        Args:
            account_id (str): The unique identifier or email of the service account to disable.
        """
        service_account_email = self._normalize_account_email(account_id)
        request = types.DisableServiceAccountRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"

        self.client.disable_service_account(request=request)

        # Wait for the service account to be disabled
        delay = 1
        for _ in range(self.max_retries):
            if not self._service_account_is_enabled(account_id):
                break
            time.sleep(delay)
            delay *= 2
        else:
            self.logger.error(f"Service account {account_id} disabling timed out after {self.max_retries} retries.")
            raise exceptions.DeadlineExceeded(f"Service account {account_id} disabling timed out.")
        
        self.logger.info(f"Disabled service account: {account_id}")

    def delete_service_account(self, account_id: str) -> None:
        """
        Deletes a service account in the specified project.

        Args:
            account_id (str): The unique identifier or email of the service account to delete.
        """
        service_account_email = self._normalize_account_email(account_id)
        request = types.DeleteServiceAccountRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"

        self.client.delete_service_account(request=request)

        # Wait for the service account to be deleted
        delay = 1
        for _ in range(self.max_retries):
            if not self._service_account_exists(account_id):
                break
            time.sleep(delay)
            delay *= 2
        else:
            self.logger.error(f"Service account {account_id} deletion timed out after {self.max_retries} retries.")
            raise exceptions.DeadlineExceeded(f"Service account {account_id} deletion timed out.")

        self.logger.info(f"Deleted service account: {account_id}")

    def _get_service_account_keys(self, account_id: str) -> List[iam_admin_v1.ServiceAccountKey]:
        """
        Retrieves all keys for the specified service account.

        Args:
            account_id (str): The unique identifier or email of the service account.
        
        Returns:
            List[iam_admin_v1.ServiceAccountKey]: A list of service account key objects.
        """
        service_account_email = self._normalize_account_email(account_id)
        request = types.ListServiceAccountKeysRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"

        response = self.client.list_service_account_keys(request=request)
        self.logger.debug(f"Listed keys for service account: {account_id}")
        return list(response.keys)

    def _service_account_key_exists(self, account_id: str, key_id: str) -> bool:
        """
        Checks if a service account key exists for the specified service account.

        Args:
            account_id (str): The unique identifier or email of the service account.
            key_id (str): The ID of the service account key to check.
        
        Returns:
            bool: True if the key exists, False otherwise.
        """
        keys = self._get_service_account_keys(account_id)
        return any(key.name.split('/')[-1] == key_id for key in keys)
    
    def create_service_account_key(self, account_id: str) -> types.ServiceAccountKey:
        """
        Creates a key for the specified service account.
        Remember the private key ID is only returned once.
        If the service account is disabled, it will be enabled first.
        Includes retry logic to handle service account propagation delays.

        Args:
            account_id (str): The unique identifier or email of the service account.
        
        Returns:
            types.ServiceAccountKey: The created service account key object.
            str: The private key ID of the created key.
        """
        service_account_email = self._normalize_account_email(account_id)
        
        # Retry logic for service account access and key creation
        delay = 1
        for attempt in range(self.max_retries):
            try:
                # Check if service account exists and get its state
                get_request = types.GetServiceAccountRequest()
                get_request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"
                
                service_account = self.client.get_service_account(request=get_request)
                if service_account.disabled:
                    self.logger.info(f"Service account {account_id} is disabled. Enabling it first.")
                    self.enable_service_account(account_id)
                
                # Create the key
                request = types.CreateServiceAccountKeyRequest()
                request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}"
                
                key = self.client.create_service_account_key(request=request)
                
                # Wait for the key to be created and available
                key_delay = 1
                for _ in range(self.max_retries):
                    if self._service_account_key_exists(account_id, key.name.split('/')[-1]):
                        break
                    time.sleep(key_delay)
                    key_delay *= 2
                else:
                    self.logger.error(f"Service account key creation for {account_id} timed out after {self.max_retries} retries.")
                    raise exceptions.DeadlineExceeded(f"Service account key creation for {account_id} timed out.")

                self.logger.info(f"Created service account key for {account_id}")
                return key
                
            except exceptions.NotFound as e:
                if attempt < self.max_retries - 1:
                    self.logger.warning(f"Service account {account_id} not found (attempt {attempt + 1}/{self.max_retries}), retrying in {delay}s. This may be due to propagation delay.")
                    time.sleep(delay)
                    delay *= 2
                else:
                    self.logger.error(f"Service account {account_id} not found after {self.max_retries} attempts")
                    raise
            except Exception as e:
                # For other exceptions, don't retry
                self.logger.error(f"Error creating service account key for {account_id}: {e}")
                raise
        
        # This should not be reached due to the raise in the except block
        raise exceptions.NotFound(f"Service account {account_id} not found after {self.max_retries} attempts")
    
    def delete_service_account_key(self, account_id: str, key_id: str) -> None:
        """
        Deletes a key for the specified service account.

        Args:
            account_id (str): The unique identifier or email of the service account.
            key_id (str): The ID of the key to delete.
        
        Raises:
            exceptions.NotFound: If the key does not exist.
            exceptions.FailedPrecondition: If the key cannot be deleted due to constraints.
        """
        service_account_email = self._normalize_account_email(account_id)
        request = types.DeleteServiceAccountKeyRequest()
        request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}/keys/{key_id}"

        try:
            self.client.delete_service_account_key(request=request)
        except exceptions.NotFound:
            self.logger.warning(f"Service account key {key_id} not found for account: {account_id} (may have been already deleted)")
            raise
        except exceptions.FailedPrecondition as e:
            self.logger.warning(f"Failed to delete service account key {key_id} for account: {account_id}. Error: {e}")
            raise
        except Exception as e:
            self.logger.error(f"Unexpected error deleting service account key {key_id} for account: {account_id}. Error: {e}")
            raise

        # Wait for the key to be deleted
        delay = 1
        for _ in range(self.max_retries):
            if not self._service_account_key_exists(account_id, key_id):
                break
            time.sleep(delay)
            delay *= 2
        else:
            self.logger.error(f"Service account key deletion for {account_id} timed out after {self.max_retries} retries.")
            raise exceptions.DeadlineExceeded(f"Service account key deletion for {account_id} timed out.")

        self.logger.info(f"Deleted service account key: {key_id} for account: {account_id}")

    def test_service_account_key(self, key_data: bytes) -> bool:
        """
        Tests if a service account key is valid by attempting to authenticate and make an API call.
        Includes retry logic to handle key propagation delays.

        Args:
            key_data (bytes): The private key data from the service account key.
        
        Returns:
            bool: True if the key is valid and can authenticate, False otherwise.
        """
        try:
            key_info = json.loads(key_data.decode('utf-8'))
        except json.JSONDecodeError as json_error:
            self.logger.error(f"Invalid JSON in service account key: {json_error}")
            return False

        delay = 1
        for attempt in range(self.max_retries):
            try:
                credentials = service_account.Credentials.from_service_account_info(
                    key_info,
                    scopes=['https://www.googleapis.com/auth/cloud-platform']
                )
                
                request = Request()
                credentials.refresh(request)

                self.logger.info(f"Service account key is valid and can authenticate")
                return True
                    
            except Exception as auth_error:
                if attempt < self.max_retries - 1:  # Don't log on the last attempt
                    delay *= 2
                    self.logger.warning(f"Authentication attempt {attempt + 1} failed (will retry in {delay}s): {auth_error}")
                    time.sleep(delay)
                else:
                    self.logger.error(f"Authentication failed with service account key after {self.max_retries} attempts: {auth_error}")
                    return False
        
        return False