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

class VisionAIAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Cloud Vision AI tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Cloud Vision AI.
        """
        if command == 'analyze_image_labels':
            return self._analyze_image_labels(**kwargs)
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by VisionAIAgent.")

    def _analyze_image_labels(self, project_id: str, bucket_name: str, object_name: str, max_labels: int = 10):
        """
        Detects labels in an image stored in a Cloud Storage bucket.
        """
        print(f"VisionAIAgent: Analyzing image '{object_name}' in bucket '{bucket_name}' in project '{project_id}' for labels...")
        try:
            client = vision.ImageAnnotatorClient(credentials=GCPConnector.get_credentials())

            image = vision.Image()
            image.source.image_uri = f'gs://{bucket_name}/{object_name}'

            response = client.label_detection(image=image)
            labels = []
            for label in response.label_annotations:
                labels.append({
                    'Description': label.description,
                    'Score': label.score
                })

            if not labels:
                return {"status": "success", "message": f"No labels detected in image '{object_name}'."}
            
            return {"status": "success", "message": f"Labels detected in image '{object_name}'.", "labels": labels}
        except Exception as e:
            print(f"Error analyzing image for labels: {e}")
            return {"status": "error", "message": str(e)}