from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import time

class KinesisAgent(BaseAgent):
    """
    An agent specialized in handling AWS Kinesis tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Kinesis.
        """
        if command == 'smart_create_data_stream':
            stream_name = kwargs.get('stream_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create a Kinesis Data Stream '{stream_name}' in {region}.")
            confirm = input("This will create a new Kinesis Data Stream. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_data_stream(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Kinesis Data Stream command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by KinesisAgent.")

    def _smart_create_data_stream(self, region: str, stream_name: str, shard_count: int = 1):
        """
        Creates a Kinesis Data Stream with a specified number of shards.
        """
        print(f"KinesisAgent: Smart creating data stream '{stream_name}' in region {region}...")
        try:
            kinesis_client = AWSConnector.get_client('kinesis', region_name=region)

            kinesis_client.create_stream(
                StreamName=stream_name,
                ShardCount=shard_count
            )

            # Wait for stream to be active
            print("Waiting for Kinesis Data Stream to become active... (this may take a moment)")
            waiter = kinesis_client.get_waiter('stream_exists')
            waiter.wait(StreamName=stream_name)

            return {"status": "success", "message": f"Kinesis Data Stream '{stream_name}' created successfully with {shard_count} shards.", "stream_name": stream_name}
        except Exception as e:
            print(f"Error during smart Kinesis Data Stream creation: {e}")
            return {"status": "error", "message": str(e)}