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

class StorageAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Cloud Storage tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Cloud Storage.
        """
        if command == 'list_buckets':
            return self._list_buckets(**kwargs)
        elif command == 'smart_create_bucket':
            bucket_name = kwargs.get('bucket_name', '')
            print(f"You are about to smart-create a Cloud Storage bucket '{bucket_name}'.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_bucket(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create bucket command cancelled by user."}
        elif command == 'delete_bucket':
            bucket_name = kwargs.get('bucket_name', '')
            print(f"WARNING: You are about to delete Cloud Storage bucket '{bucket_name}'. This action is irreversible.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._delete_bucket(**kwargs)
            else:
                return {"status": "cancelled", "message": "Delete bucket command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by StorageAgent.")

    def _list_buckets(self, project_id: str):
        """
        Lists all Cloud Storage buckets in a specified project.
        """
        print(f"StorageAgent: Listing buckets in project '{project_id}'...")
        try:
            client = GCPConnector.get_storage_client()
            buckets = client.list_buckets(project=project_id)
            bucket_names = [bucket.name for bucket in buckets]
            return {"status": "success", "buckets": bucket_names}
        except Exception as e:
            print(f"Error listing buckets: {e}")
            return {"status": "error", "message": str(e)}

    def _smart_create_bucket(self, project_id: str, bucket_name: str, location: str = 'US', storage_class: str = 'STANDARD'):
        """
        Creates a Cloud Storage bucket with sensible defaults.
        """
        print(f"StorageAgent: Smart creating bucket '{bucket_name}' in project '{project_id}', location '{location}'...")
        try:
            client = GCPConnector.get_storage_client()
            bucket = client.bucket(bucket_name)
            bucket.storage_class = storage_class
            bucket.create(location=location, project=project_id)
            bucket.uniform_bucket_level_access = True # Enforce uniform access
            bucket.patch()
            return {"status": "success", "message": f"Bucket '{bucket_name}' created successfully."}
        except Exception as e:
            print(f"Error creating bucket: {e}")
            return {"status": "error", "message": str(e)}

    def _delete_bucket(self, project_id: str, bucket_name: str):
        """
        Deletes a specified Cloud Storage bucket.
        """
        print(f"StorageAgent: Deleting bucket '{bucket_name}' in project '{project_id}'...")
        try:
            client = GCPConnector.get_storage_client()
            bucket = client.bucket(bucket_name)
            bucket.delete(force=True) # force=True deletes non-empty buckets
            return {"status": "success", "message": f"Bucket '{bucket_name}' deleted successfully."}
        except Exception as e:
            print(f"Error deleting bucket: {e}")
            return {"status": "error", "message": str(e)}