from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import os

class PollyAgent(BaseAgent):
    """
    An agent specialized in handling AWS Polly tasks.
    """

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

    def _synthesize_speech(self, region: str, text: str, output_s3_bucket: str, output_s3_key: str, voice_id: str = 'Joanna', output_format: str = 'mp3'):
        """
        Synthesizes speech from a given text and saves the audio output to an S3 bucket.
        """
        print(f"PollyAgent: Synthesizing speech for text in region {region} and saving to s3://{output_s3_bucket}/{output_s3_key}...")
        try:
            polly_client = AWSConnector.get_client('polly', region_name=region)
            s3_client = AWSConnector.get_client('s3', region_name=region)

            response = polly_client.synthesize_speech(
                Text=text,
                OutputFormat=output_format,
                VoiceId=voice_id
            )

            # Save the audio stream to a temporary file
            with open('temp_polly_output.mp3', 'wb') as f:
                f.write(response['AudioStream'].read())
            
            # Upload to S3
            s3_client.upload_file('temp_polly_output.mp3', output_s3_bucket, output_s3_key)
            os.remove('temp_polly_output.mp3') # Clean up temporary file

            return {"status": "success", "message": f"Speech synthesized and saved to s3://{output_s3_bucket}/{output_s3_key}."}
        except Exception as e:
            print(f"Error synthesizing speech: {e}")
            return {"status": "error", "message": str(e)}