# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
import functools
from inspect import Parameter, Signature
from typing import Any

from aws_lambda_powertools import Tracer

from czc.constants import FEATURE_NAME
from czc.json import serializable

tracer = Tracer(service=FEATURE_NAME)


def init():
    return tracer


def capture_function(
    use_function_name_annotation: bool = True,
    annotations: dict[str, str] = None,
    arg_annotations: set[str] = None,
    skip_args: set[str] = None,
    skip_arg_types: set[type] = None,
    capture_response: bool = True,
):
    def function_decorator(func):
        @tracer.capture_method(capture_response=capture_response)
        @functools.wraps(func)
        def handler_wrapper(*args, **kwargs):
            if use_function_name_annotation:
                tracer.put_annotation('FUNCTION_NAME', func.__name__)
            arg_map = _collect_args(
                Signature.from_callable(func), args, kwargs, (skip_args or set()), (skip_arg_types or set())
            )
            _put_annotations({k: v for k, v in arg_map.items() if arg_annotations and k in arg_annotations})
            _put_annotations(annotations or {})

            tracer.put_metadata('arguments', serializable(_stringify_keys(arg_map)))
            return func(*args, **kwargs)

        return handler_wrapper

    return function_decorator


def _put_annotations(annotations: dict[str, str]):
    for k, v in annotations.items():
        tracer.put_annotation(k, v)


def _stringify_keys(obj: Any) -> Any:
    if isinstance(obj, dict):
        return {str(k): _stringify_keys(v) for k, v in obj.items()}
    if isinstance(obj, list) or isinstance(obj, set) or isinstance(obj, tuple):
        return [_stringify_keys(x) for x in obj]
    return obj


def _collect_args(
    sig: Signature, args: tuple[Any], kwargs: dict[str, Any], skip_args: set[str], skip_types: set[type]
) -> dict[str, Any]:
    params = list(sig.parameters.values())
    positional = params[: len(args)]
    parg_map = {p.name: a for a, p in zip(args, positional)}
    defaults = {p.name: p.default for p in params if p.default != Parameter.empty}

    result = {**parg_map, **defaults, **kwargs}

    return {k: v for k, v in result.items() if not (k in skip_args or type(v) in skip_types)}
