from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class RekognitionAgent(BaseAgent):
    """
    An agent specialized in handling AWS Rekognition tasks.
    """

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

    def _analyze_image_labels(self, region: str, bucket_name: str, object_key: str, max_labels: int = 10, min_confidence: float = 75.0):
        """
        Detects labels in an image stored in an S3 bucket.
        """
        print(f"RekognitionAgent: Analyzing image '{object_key}' in bucket '{bucket_name}' in region {region} for labels...")
        try:
            rekognition_client = AWSConnector.get_client('rekognition', region_name=region)

            response = rekognition_client.detect_labels(
                Image={'S3Object': {'Bucket': bucket_name, 'Name': object_key}},
                MaxLabels=max_labels,
                MinConfidence=min_confidence
            )

            labels = []
            for label in response['Labels']:
                labels.append({
                    'Name': label['Name'],
                    'Confidence': label['Confidence']
                })

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