import io
import json
from typing import Any, TYPE_CHECKING

import boto3

from aws_testing_utils import config

from .logger import log

if TYPE_CHECKING:
    from mypy_boto3_lambda import LambdaClient


class LambdaHandler:
    """Handles interactions with AWS Lambda."""

    lambda_client: 'LambdaClient'

    def __init__(self) -> None:
        self.lambda_client = boto3.client(
            'lambda',
            region_name=config.AWS_REGION,
            endpoint_url=config.LAMBDA_ENDPOINT_URL,
        )

    def invoke(
        self,
        function_name: str,
        payload: dict[str, Any] | None = None,
        *,
        assertion: bool = True,
        expected_status: int = 200,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Invokes a Lambda function and returns the response.

        Args:
            function_name: Lambda function name in AWS.
            payload: Event payload to invoke with.
            assertion: If True, asserts StatusCode == expected_status and no
                FunctionError. Pass False when expecting the lambda to fail.
            expected_status: StatusCode asserted when ``assertion`` is True.
                Defaults to 200 (RequestResponse); async ``Event`` invokes
                return 202, so pass ``expected_status=202`` for those.
            kwargs: Extra keyword arguments forwarded straight to boto3's
                ``invoke`` (e.g. ``InvocationType``, ``LogType``,
                ``ClientContext``, ``Qualifier``), so new boto3 parameters
                need no library release.
        """
        if payload is None:
            payload = {}
        log.info(f'Invoking lambda function: {function_name}')
        # Reserved keys are listed last so they always win over any collision
        # in kwargs — callers can't accidentally override the target function
        # or payload and make function_name misleading.
        invoke_args: dict[str, Any] = {
            **kwargs,
            'FunctionName': function_name,
            'Payload': json.dumps(payload),
        }
        response: dict[str, Any] = dict(self.lambda_client.invoke(**invoke_args))
        log.info(f'Lambda response: {response}')

        # Buffer the payload so callers can read it after this method returns.
        raw_payload = response.pop('Payload', io.BytesIO(b'')).read()
        log.debug(raw_payload)
        response['Payload'] = io.BytesIO(raw_payload)

        # The local Lambda RIE doesn't set X-Amz-Function-Error on unhandled
        # exceptions, so boto3 never adds FunctionError to the response dict.
        # Synthesize it from the payload's errorType field when absent.
        if 'FunctionError' not in response and raw_payload:
            try:
                body = json.loads(raw_payload)
                if isinstance(body, dict) and 'errorType' in body:
                    response['FunctionError'] = 'Unhandled'
            except (json.JSONDecodeError, TypeError, UnicodeDecodeError):
                pass

        if assertion:
            status = response.get('StatusCode')
            assert status == expected_status, (
                f'{function_name} returned StatusCode {status}, '
                f'expected {expected_status}'
            )
            assert 'FunctionError' not in response, (
                f'{function_name} returned FunctionError: '
                f'{response.get("FunctionError")}'
            )
        return response

    def warm_up(self, function_names: str | list[str]) -> None:
        """Fire async invocations to pre-warm lambda containers.

        Args:
            function_names: Single function name or list of function names to warm up.
        """
        if isinstance(function_names, str):
            function_names = [function_names]
        for fn in function_names:
            try:
                log.info(f'Warming up lambda function: {fn}')
                self.lambda_client.invoke(
                    FunctionName=fn,
                    InvocationType='Event',
                    Payload=b'{}',
                )
            except Exception as e:
                log.warning(f'Failed to warm up lambda function {fn}: {e}')