from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class Route53Agent(BaseAgent):
    """
    An agent specialized in handling AWS Route 53 tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Route 53.
        """
        if command == 'list_hosted_zones':
            return self._list_hosted_zones()
        elif command == 'list_resource_record_sets':
            return self._list_resource_record_sets(**kwargs)
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by Route53Agent.")

    def _list_hosted_zones(self):
        """Lists all hosted zones in Route 53."""
        print("Route53Agent: Listing hosted zones...")
        try:
            route53_client = AWSConnector.get_client('route53')
            response = route53_client.list_hosted_zones()
            hosted_zones = response.get('HostedZones', [])
            return {"status": "success", "hosted_zones": hosted_zones}
        except Exception as e:
            print(f"Error listing hosted zones: {e}")
            return {"status": "error", "message": str(e)}

    def _list_resource_record_sets(self, hosted_zone_id: str):
        """Lists resource record sets for a given hosted zone ID."""
        print(f"Route53Agent: Listing resource record sets for hosted zone {hosted_zone_id}...")
        try:
            route53_client = AWSConnector.get_client('route53')
            response = route53_client.list_resource_record_sets(HostedZoneId=hosted_zone_id)
            record_sets = response.get('ResourceRecordSets', [])
            return {"status": "success", "resource_record_sets": record_sets}
        except Exception as e:
            print(f"Error listing resource record sets: {e}")
            return {"status": "error", "message": str(e)}