from ..base_gcp_agent import BaseGCPAgent
from ..gcp_connector import GCPConnector
from google.cloud import iam_admin_v1

class IAMAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Cloud IAM tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Cloud IAM.
        """
        if command == 'smart_create_service_account':
            sa_id = kwargs.get('service_account_id', '')
            print(f"You are about to smart-create a service account '{sa_id}'.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_service_account(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Service Account command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by IAMAgent.")

    def _smart_create_service_account(self, project_id: str, service_account_id: str, display_name: str = None, role: str = 'roles/viewer'):
        """
        Creates a service account and attaches a specified role.
        """
        print(f"IAMAgent: Smart creating service account '{service_account_id}' in project '{project_id}'...")
        try:
            client = iam_admin_v1.IAMClient(credentials=GCPConnector.get_credentials())

            # 1. Create Service Account
            service_account = iam_admin_v1.ServiceAccount()
            service_account.display_name = display_name if display_name else service_account_id

            request = iam_admin_v1.CreateServiceAccountRequest(
                name=f"projects/{project_id}",
                service_account_id=service_account_id,
                service_account=service_account,
            )

            sa_response = client.create_service_account(request=request)
            service_account_email = sa_response.email
            print(f"Service Account '{service_account_email}' created.")

            # 2. Attach Role to Service Account
            policy_client = iam_admin_v1.IAMPolicyClient(credentials=GCPConnector.get_credentials())
            policy_request = iam_admin_v1.GetIamPolicyRequest(resource=f"projects/{project_id}")
            policy = policy_client.get_iam_policy(request=policy_request)

            # Add the new binding
            binding = iam_admin_v1.Binding()
            binding.role = role
            binding.members = [f"serviceAccount:{service_account_email}"]
            policy.bindings.append(binding)

            set_policy_request = iam_admin_v1.SetIamPolicyRequest(resource=f"projects/{project_id}", policy=policy)
            policy_client.set_iam_policy(request=set_policy_request)
            print(f"Role '{role}' attached to service account '{service_account_email}'.")

            return {"status": "success", "message": f"Service account '{service_account_email}' created and role '{role}' attached.", "service_account_email": service_account_email}
        except Exception as e:
            print(f"Error creating service account or attaching role: {e}")
            return {"status": "error", "message": str(e)}