from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class GuardDutyAgent(BaseAgent):
    """
    An agent specialized in handling AWS GuardDuty tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to GuardDuty.
        """
        if command == 'smart_enable_guardduty':
            region = kwargs.get('region', '')
            print(f"You are about to smart-enable GuardDuty in {region}.")
            confirm = input("This will create a GuardDuty detector if one doesn't exist. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_enable_guardduty(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Enable GuardDuty command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by GuardDutyAgent.")

    def _smart_enable_guardduty(self, region: str):
        """
        Enables GuardDuty in the specified region, creating a detector if one doesn't exist.
        """
        print(f"GuardDutyAgent: Smart enabling GuardDuty in region {region}...")
        try:
            gd_client = AWSConnector.get_client('guardduty', region_name=region)

            # Check if a detector already exists
            list_detectors_response = gd_client.list_detectors()
            detector_ids = list_detectors_response['DetectorIds']

            detector_id = None
            if detector_ids:
                detector_id = detector_ids[0]
                print(f"GuardDuty detector '{detector_id}' already exists. Using existing.")
            else:
                print("No GuardDuty detector found. Creating a new one...")
                create_detector_response = gd_client.create_detector(
                    Enable=True,
                    FindingPublishingFrequency='FIFTEEN_MINUTES'
                )
                detector_id = create_detector_response['DetectorId']
                print(f"GuardDuty detector '{detector_id}' created and enabled.")

            return {"status": "success", "message": f"GuardDuty is enabled in region {region} with detector '{detector_id}'.", "detector_id": detector_id}
        except Exception as e:
            print(f"Error during smart GuardDuty enablement: {e}")
            return {"status": "error", "message": str(e)}