from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class OpenSearchAgent(BaseAgent):
    """
    An agent specialized in handling AWS Amazon OpenSearch Service tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to OpenSearch.
        """
        if command == 'create_domain':
            print(f"You are about to create an OpenSearch domain named '{kwargs.get('domain_name')}'.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._create_domain(**kwargs)
            else:
                return {"status": "cancelled", "message": "Create domain command cancelled by user."}
        elif command == 'troubleshoot_domain':
            domain_name = kwargs.get('domain_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to troubleshoot OpenSearch domain '{domain_name}' in {region}.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._troubleshoot_domain(**kwargs)
            else:
                return {"status": "cancelled", "message": "Troubleshoot domain command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by OpenSearchAgent.")

    def _troubleshoot_domain(self, region: str, domain_name: str):
        """
        Runs a series of diagnostic checks on an OpenSearch domain.
        """
        print(f"OpenSearchAgent: Troubleshooting domain '{domain_name}' in region {region}...")
        report = {
            'domain_name': domain_name,
            'region': region,
            'findings': [],
            'recommendations': []
        }
        try:
            os_client = AWSConnector.get_client('opensearch', region_name=region)

            # 1. Describe Domain Status
            response = os_client.describe_domain(DomainName=domain_name)
            domain_status = response['DomainStatus']

            processing = domain_status.get('Processing', False)
            domain_state = "Processing" if processing else "Active"
            report['findings'].append(f"Domain state: '{domain_state}'.")
            if processing:
                report['recommendations'].append("Domain is currently processing updates. Wait for it to become active.")

            if domain_status.get('Deleted', False):
                report['findings'].append("Domain is marked for deletion.")
                report['recommendations'].append("Domain is being deleted. No further action needed unless deletion is stuck.")

            if domain_status.get('Created', False) and not domain_status.get('Endpoint'):
                report['findings'].append("Domain is created but endpoint is not yet available.")
                report['recommendations'].append("Wait for the domain endpoint to become available.")

            if domain_status.get('Endpoint'):
                report['findings'].append(f"Domain endpoint: {domain_status['Endpoint']}")

            # Check health status (if available, requires more advanced monitoring setup)
            # For simplicity, we'll just check the overall status from describe_domain
            if domain_status.get('DomainHealth'): # This field might not always be present or detailed
                report['findings'].append(f"Domain health: {domain_status['DomainHealth']}")

            # Check cluster config
            cluster_config = domain_status.get('ClusterConfig', {})
            report['findings'].append(f"Instance type: {cluster_config.get('InstanceType')}, Count: {cluster_config.get('InstanceCount')}")

            print(f"--- Workflow: Troubleshoot OpenSearch Domain Finished ---")
            return {"status": "success", "report": report}

        except os_client.exceptions.ResourceNotFoundException:
            return {"status": "error", "message": f"OpenSearch domain '{domain_name}' not found in region {region}."}
        except Exception as e:
            print(f"--- Workflow Failed: {e} ---")
            return {"status": "error", "message": f"Troubleshoot OpenSearch Domain workflow failed: {e}"}

    def _create_domain(self, region: str, domain_name: str, engine_version: str, instance_type: str, instance_count: int, volume_size: int):
        """Creates a new OpenSearch Service domain."""
        print(f"OpenSearchAgent: Creating domain '{domain_name}' in region {region}...")
        try:
            os_client = AWSConnector.get_client('opensearch', region_name=region)
            
            response = os_client.create_domain(
                DomainName=domain_name,
                EngineVersion=engine_version,
                ClusterConfig={
                    'InstanceType': instance_type,
                    'InstanceCount': instance_count,
                    'DedicatedMasterEnabled': False, # Default for simplicity
                    'ZoneAwarenessEnabled': False # Default for simplicity
                },
                EBSOptions={
                    'EBSEnabled': True,
                    'VolumeType': 'gp2',
                    'VolumeSize': volume_size
                },
                # In a real-world scenario, you would configure access policies.
                # For this example, we use an open access policy for simplicity. 
                # WARNING: This is not recommended for production environments.
                AccessPolicies='{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Principal": {"AWS": "*"}, "Action": "es:*", "Resource": "arn:aws:es:' + region + ':' + AWSConnector.get_account_id() + ':domain/' + domain_name + '/*"}]}'
            )
            
            domain_arn = response['DomainStatus']['ARN']
            return {"status": "success", "message": f"OpenSearch domain '{domain_name}' creation initiated.", "domain_arn": domain_arn}
        except Exception as e:
            print(f"Error creating OpenSearch domain: {e}")
            return {"status": "error", "message": str(e)}