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

class VertexAIAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Vertex AI Workbench tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Vertex AI Workbench.
        """
        if command == 'smart_create_notebook':
            notebook_name = kwargs.get('notebook_name', '')
            print(f"You are about to smart-create a Vertex AI Workbench notebook '{notebook_name}'.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_notebook(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Vertex AI Workbench notebook command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by VertexAIAgent.")

    def _smart_create_notebook(self, project_id: str, location: str, notebook_name: str, machine_type: str = 'n1-standard-1', vm_image_project: str = 'deeplearning-platform-release', vm_image_family: str = 'tf-latest-gpu'):
        """
        Creates a Vertex AI Workbench notebook instance with sensible defaults.
        """
        print(f"VertexAIAgent: Smart creating notebook '{notebook_name}' in project '{project_id}', location '{location}'...")
        try:
            client = notebooks_v1.NotebookServiceClient(credentials=GCPConnector.get_credentials())

            # Configure VM image
            vm_image = notebooks_v1.VmImage()
            vm_image.project = vm_image_project
            vm_image.image_family = vm_image_family

            # Configure instance
            instance = notebooks_v1.Instance()
            instance.vm_image = vm_image
            instance.machine_type = machine_type

            request = notebooks_v1.CreateInstanceRequest(
                parent=f"projects/{project_id}/locations/{location}",
                instance_id=notebook_name,
                instance=instance,
            )

            operation = client.create_instance(request=request)
            print("Waiting for Vertex AI Workbench notebook creation to complete... (this may take several minutes)")
            operation.result() # This blocks until the operation completes

            return {"status": "success", "message": f"Vertex AI Workbench notebook '{notebook_name}' created successfully.", "notebook_name": notebook_name}
        except Exception as e:
            print(f"Error creating Vertex AI Workbench notebook: {e}")
            return {"status": "error", "message": str(e)}