import logging
import json
from typing import Dict, Any
from http import HTTPStatus
import boto3
from botocore.exceptions import ClientError
import os

REGION = os.environ['AWS_REGION']

logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Replace with the name of your AWS Location Place Index
PLACE_INDEX_NAME = 'helpngoindex'

def ensure_place_index_exists(index_name: str):
    client = boto3.client('location', region_name=REGION)

    # Check if index exists
    try:
        response = client.list_place_indexes()
        if any(entry['IndexName'] == index_name for entry in response['Entries']):
            print(f"✅ Place index '{index_name}' already exists.")
            return
    except ClientError as e:
        print(f"❌ Error listing indexes: {e}")
        return

    # Create the index
    try:
        print(f"🆕 Creating place index '{index_name}' in region '{region}'...")
        client.create_place_index(
            IndexName=index_name,
            DataSource='Here',
            Description='Auto-created place index for geocoding',
            PricingPlan='RequestBasedUsage'
        )
    except ClientError as e:
        print(f"❌ Error creating index: {e}")
        return

    # Wait until ACTIVE
    print("⏳ Waiting for index to become ACTIVE...")
    while True:
        try:
            status = client.describe_place_index(IndexName=index_name)['Status']
            print(f"Status: {status}")
            if status == 'ACTIVE':
                print(f"✅ Place index '{index_name}' is now ACTIVE.")
                break
            time.sleep(2)
        except ClientError as e:
            print(f"❌ Error checking status: {e}")
            break

# Initialize the AWS Location client
location_client = boto3.client('location', region_name=REGION)

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    AWS Lambda handler for processing Bedrock agent requests and geocoding location names using AWS Location Service.
    """
    try:
        action_group = event['actionGroup']
        function = event['function']
        message_version = event.get('messageVersion', 1)
        parameters = event.get('parameters', [])

        parameters_dict = {parameter["name"]: parameter["value"] for parameter in parameters}
        location_name = parameters_dict.get('location_name')
        if not location_name:
            raise KeyError("Missing required parameter: 'location_name'")

        # Use AWS Location Service to geocode the location
        ensure_place_index_exists(PLACE_INDEX_NAME)
        response = location_client.search_place_index_for_text(
            IndexName=PLACE_INDEX_NAME,
            Text=location_name,
            MaxResults=1
        )

        results = response.get("Results", [])
        if not results:
            raise ValueError(f"No coordinates found for '{location_name}'")

        # Extract coordinates [Longitude, Latitude]
        lon, lat = results[0]['Place']['Geometry']['Point']

        # Prepare Bedrock-compatible response
        response_body = {
            'TEXT': {
                'body': f"Coordinates for '{location_name}': Latitude {lat}, Longitude {lon}"
            }
        }

        action_response = {
            'actionGroup': action_group,
            'function': function,
            'functionResponse': {
                'responseBody': response_body
            }
        }

        final_response = {
            'response': action_response,
            'messageVersion': message_version
        }

        logger.info('Success: %s', final_response)
        return final_response

    except KeyError as e:
        logger.error('Missing required field: %s', str(e))
        return {
            'statusCode': HTTPStatus.BAD_REQUEST,
            'body': f'Error: {str(e)}'
        }

    except ClientError as e:
        logger.error('AWS Location error: %s', str(e))
        return {
            'statusCode': HTTPStatus.INTERNAL_SERVER_ERROR,
            'body': f'AWS Location error: {str(e)}'
        }

    except Exception as e:
        logger.error('Unexpected error: %s', str(e))
        return {
            'statusCode': HTTPStatus.INTERNAL_SERVER_ERROR,
            'body': f'Internal server error: {str(e)}'
        }