"""Lambda Invoke Hook."""

from datetime import datetime
from os import environ
from typing import Any

from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
from botocore.client import Config

from hooks.lambda_invoke.invocation_response import LambdaInvocationResponse


class OrchLambdaHook(LambdaHook):
    """Lambda Invoke Hook.

    Parameters
    ----------
    function_name : str
        Name of the Lambda function to invoke.
    log_type : str, optional
        Type of logs to return. Defaults to 'Tail'.
    qualifier : str, optional
        Qualifier for the Lambda version or alias. Defaults to 'provisioned'.
    invocation_type : str, optional
        Invocation type (e.g., 'RequestResponse', 'Event').
        Defaults to 'RequestResponse'.
    **kwargs : dict, optional
        Additional keyword arguments passed to
        :class:`airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook`.
        Common options include connection and client configuration such as
        ``aws_conn_id`` (Airflow connection ID).
    """

    def __init__(
            self,
            function_name: str,
            log_type: str = 'Tail',
            qualifier: str = 'provisioned',
            invocation_type: str = 'RequestResponse',
            **kwargs: Any,
    ):
        """Create hook for a Lambda function."""
        self.function_name = function_name
        self.log_type = log_type
        self.qualifier = qualifier
        self.invocation_type = invocation_type

        super(OrchLambdaHook, self).__init__(
            region_name=environ.get('AWS_DEFAULT_REGION', 'us-east-1'),
            config=Config(
                connect_timeout=900,
                read_timeout=900,
                retries={'max_attempts': 0}
            ),
            **kwargs,
        )

    def invoke_lambda(self, payload):
        """Invoke the configured Lambda function."""
        start_time = datetime.now()
        response = super().invoke_lambda(
            function_name=self.function_name,
            log_type=self.log_type,
            invocation_type=self.invocation_type,
            payload=payload,
            qualifier=self.qualifier
        )
        duration = datetime.now() - start_time
        lambda_invocation_response = LambdaInvocationResponse(response)
        print(f'RESPONSE: ({duration}s) ', lambda_invocation_response.dump())
        return lambda_invocation_response