import json
from typing import TYPE_CHECKING

import requests
from aibs_informatics_core.models.aws.core import AWSRegion
from aibs_informatics_core.models.aws.lambda_ import LambdaFunctionName, LambdaFunctionUrl
from aibs_informatics_core.models.base import ModelProtocol
from botocore.exceptions import ClientError
from requests.auth import AuthBase

from aibs_informatics_aws_utils.auth import IamAWSRequestsAuth
from aibs_informatics_aws_utils.core import AWSService, get_client_error_code

if TYPE_CHECKING:  # pragma: no cover
    from mypy_boto3_lambda.type_defs import FileSystemConfigTypeDef
else:
    FileSystemConfigTypeDef = dict


get_lambda_client = AWSService.LAMBDA.get_client


def get_lambda_function_url(
    function_name: LambdaFunctionName | str, region: AWSRegion | None = None
) -> LambdaFunctionUrl | None:
    """Get the function URL for a Lambda function.

    Args:
        function_name (Union[LambdaFunctionName, str]): The name or ARN of the Lambda function.
        region (Optional[AWSRegion]): AWS region. Defaults to None (uses default region).

    Returns:
        The function URL if configured, otherwise None.
    """
    function_name = LambdaFunctionName(function_name)

    lambda_client = get_lambda_client(region=region)

    try:
        response = lambda_client.get_function_url_config(FunctionName=function_name)
    except ClientError as e:
        if get_client_error_code(e) == "ResourceNotFoundException":
            return None
        else:
            raise e
    return LambdaFunctionUrl(response["FunctionUrl"])


def get_lambda_function_file_systems(
    function_name: LambdaFunctionName | str, region: AWSRegion | None = None
) -> list[FileSystemConfigTypeDef]:
    """Get the file system configurations for a Lambda function.

    Args:
        function_name (Union[LambdaFunctionName, str]): The name or ARN of the function.
        region (Optional[AWSRegion]): AWS region. Defaults to None (uses default region).

    Returns:
        List of file system configurations (EFS mount points).
    """
    function_name = LambdaFunctionName(function_name)

    lambda_client = get_lambda_client(region=region)

    response = lambda_client.get_function_configuration(FunctionName=function_name)

    fs_configs = response.get("FileSystemConfigs")

    return fs_configs or []


def call_lambda_function_url(
    function_name: LambdaFunctionName | LambdaFunctionUrl | str,
    payload: ModelProtocol | dict | str | bytes | None = None,
    region: AWSRegion | None = None,
    headers: dict | None = None,
    auth: AuthBase | None = None,
    **request_kwargs,
) -> dict | str | None:
    """Call a Lambda function via its function URL.

    Args:
        function_name (Union[LambdaFunctionName, LambdaFunctionUrl, str]): The function
            name, ARN, or URL to call.
        payload (Optional[Union[ModelProtocol, dict, str, bytes]]): Request payload.
        region (Optional[AWSRegion]): AWS region. Defaults to None.
        headers (Optional[dict]): Optional HTTP headers.
        auth (Optional[AuthBase]): Authentication handler. Defaults to IAM auth.
        **request_kwargs: Additional arguments passed to the requests library.

    Raises:
        ValueError: If the function name/URL is invalid or function not found.

    Returns:
        The response as dict (if JSON) or str, or None on error.
    """
    if LambdaFunctionName.is_valid(function_name):
        function_url = get_lambda_function_url(LambdaFunctionName(function_name), region=region)
        if function_url is None:
            raise ValueError(f"Function {function_name} not found")
        function_url = LambdaFunctionUrl(function_url)
    elif LambdaFunctionUrl.is_valid(function_name):
        function_url = LambdaFunctionUrl(function_name)
    else:
        raise ValueError(f"Invalid function name or url: {function_name}")

    json_payload: str | None = None
    if isinstance(payload, (dict, list)):
        json_payload = json.dumps(payload)
    elif isinstance(payload, ModelProtocol):
        json_payload = json.dumps(payload.to_dict())
    elif isinstance(payload, str):
        json_payload = payload
    elif isinstance(payload, bytes):
        json_payload = payload.decode("utf-8")
    elif payload is None:
        pass
    else:
        raise ValueError(f"Invalid payload type: {type(payload)}")

    if headers is None:
        headers = {}

    if auth is None:
        auth = IamAWSRequestsAuth(service_name="lambda")

    response = requests.request(
        method="POST" if json_payload else "GET",
        url=function_url.base_url + function_url.path,
        json=json_payload,
        params=function_url.query,
        headers=headers,
        auth=auth,
        **request_kwargs,
    )
    if response.ok:
        if response.headers.get("Content-Type") == "application/json":
            return response.json()
        else:
            return response.text
    else:
        response.raise_for_status()
        return None
