from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import json

class APIGatewayAgent(BaseAgent):
    """
    An agent specialized in handling AWS API Gateway tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to API Gateway.
        """
        if command == 'smart_create_rest_api':
            api_name = kwargs.get('api_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create a REST API '{api_name}' in {region}.")
            confirm = input("This will create a REST API with a /hello resource and mock integration. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_rest_api(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create REST API command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by APIGatewayAgent.")

    def _smart_create_rest_api(self, region: str, api_name: str, stage_name: str = 'dev', path: str = 'hello'):
        """
        Creates a simple REST API with a mock integration and deploys it to a stage.
        """
        print(f"APIGatewayAgent: Smart creating REST API '{api_name}' in region {region}...")
        try:
            apigw_client = AWSConnector.get_client('apigateway', region_name=region)

            # 1. Create REST API
            create_api_response = apigw_client.create_rest_api(
                name=api_name,
                description=f'Smart-created REST API for {api_name}',
                version='1.0',
            )
            api_id = create_api_response['id']
            print(f"REST API '{api_name}' ({api_id}) created.")

            # Get root resource ID
            get_resources_response = apigw_client.get_resources(restApiId=api_id)
            root_resource_id = [r for r in get_resources_response['items'] if r['path'] == '/'][0]['id']

            # 2. Create a resource (e.g., /hello)
            create_resource_response = apigw_client.create_resource(
                restApiId=api_id,
                parentId=root_resource_id,
                pathPart=path
            )
            resource_id = create_resource_response['id']
            print(f"Resource '{path}' ({resource_id}) created.")

            # 3. Create a GET method
            apigw_client.put_method(
                restApiId=api_id,
                resourceId=resource_id,
                httpMethod='GET',
                authorizationType='NONE',
                apiKeyRequired=False
            )
            print(f"GET method on /{path} created.")

            # 4. Set up mock integration
            apigw_client.put_integration(
                restApiId=api_id,
                resourceId=resource_id,
                httpMethod='GET',
                type='MOCK',
                requestTemplates={
                    'application/json': '{"statusCode": 200}'
                }
            )
            print("Mock integration created.")

            # 5. Set up method response
            apigw_client.put_method_response(
                restApiId=api_id,
                resourceId=resource_id,
                httpMethod='GET',
                statusCode='200',
                responseModels={'application/json': 'Empty'}
            )
            print("Method response configured.")

            # 6. Set up integration response
            apigw_client.put_integration_response(
                restApiId=api_id,
                resourceId=resource_id,
                httpMethod='GET',
                statusCode='200',
                selectionPattern='',
                responseTemplates={
                    'application/json': '{"message": "Hello from API Gateway!"}'
                }
            )
            print("Integration response configured.")

            # 7. Deploy API to a stage
            deploy_response = apigw_client.create_deployment(
                restApiId=api_id,
                stageName=stage_name,
                description=f'Deployment for {stage_name} stage'
            )
            print(f"API deployed to stage '{stage_name}'.")

            invoke_url = f'https://{api_id}.execute-api.{region}.amazonaws.com/{stage_name}/{path}'

            return {"status": "success", "message": f"REST API '{api_name}' created and deployed. Invoke URL: {invoke_url}", "api_id": api_id, "invoke_url": invoke_url}
        except Exception as e:
            print(f"Error during smart API Gateway REST API creation: {e}")
            return {"status": "error", "message": str(e)}