from ..base_gcp_agent import BaseGCPAgent
from ..gcp_connector import GCPConnector
from google.cloud import functions_v1
from google.cloud.functions_v1.types import CloudFunction, SourceRepository
import time
import zipfile
import os

class FunctionsAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Cloud Functions tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Cloud Functions.
        """
        if command == 'smart_create_function':
            function_name = kwargs.get('function_name', '')
            print(f"You are about to smart-create a Cloud Function '{function_name}'.")
            confirm = input("This will create a new Cloud Function. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_function(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Cloud Function command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by FunctionsAgent.")

    def _smart_create_function(self, project_id: str, region: str, function_name: str, runtime: str = 'python39', entry_point: str = 'hello_world', trigger_topic: str = None):
        """
        Creates a Cloud Function with sensible defaults.
        """
        print(f"FunctionsAgent: Smart creating function '{function_name}' in project '{project_id}', region '{region}'...")
        try:
            client = GCPConnector.get_functions_client()
            pubsub_client = GCPConnector.get_pubsub_client()

            # 1. Create a Pub/Sub topic if trigger_topic is provided
            if trigger_topic:
                topic_path = pubsub_client.topic_path(project_id, trigger_topic)
                try:
                    pubsub_client.get_topic(request={"topic": topic_path})
                    print(f"Pub/Sub topic '{trigger_topic}' already exists. Using existing.")
                except Exception:
                    pubsub_client.create_topic(request={"name": topic_path})
                    print(f"Pub/Sub topic '{trigger_topic}' created.")

            # 2. Prepare source code (simple 'hello world')
            source_code = """
import functions_framework

@functions_framework.http
def hello_world(request):
    return 'Hello World!'
"""
            zip_file_name = f'{function_name}-source.zip'
            with zipfile.ZipFile(zip_file_name, 'w') as zf:
                zf.writestr('main.py', source_code)
            print(f"Created source code zip: {zip_file_name}")

            # 3. Upload source code to Cloud Storage (requires a bucket)
            storage_client = GCPConnector.get_storage_client()
            bucket_name = f'{project_id}-cloudfunctions'
            try:
                bucket = storage_client.get_bucket(bucket_name)
            except Exception:
                bucket = storage_client.create_bucket(bucket_name, location=region)
                print(f"Created Cloud Storage bucket '{bucket_name}' for functions.")
            
            blob = bucket.blob(f'{function_name}/{zip_file_name}')
            blob.upload_from_filename(zip_file_name)
            os.remove(zip_file_name) # Clean up local zip file
            source_archive_url = f'gs://{bucket_name}/{function_name}/{zip_file_name}'
            print(f"Uploaded source code to: {source_archive_url}")

            # 4. Create Cloud Function
            parent = f"projects/{project_id}/locations/{region}"
            function = CloudFunction()
            function.name = f"{parent}/functions/{function_name}"
            function.entry_point = entry_point
            function.runtime = runtime
            function.source_archive_url = source_archive_url
            function.available_memory_mb = 128
            function.timeout = '60s'

            if trigger_topic:
                function.event_trigger.event_type = 'google.cloud.pubsub.topic.v1.messagePublished'
                function.event_trigger.resource = topic_path
            else:
                function.https_trigger.security_level = functions_v1.HttpsTrigger.SecurityLevel.ALLOW_UNAUTHENTICATED

            request = functions_v1.CreateFunctionRequest(
                parent=parent,
                cloud_function=function,
                function_id=function_name,
            )

            operation = client.create_function(request=request)
            print("Waiting for Cloud Function creation to complete... (this may take a few minutes)")
            operation.result() # This blocks until the operation completes

            return {"status": "success", "message": f"Cloud Function '{function_name}' created successfully.", "function_name": function_name}
        except Exception as e:
            print(f"Error creating Cloud Function: {e}")
            return {"status": "error", "message": str(e)}