import analytics
import asyncio
import json
import os
import copy

from typing import Callable
from base64 import b64decode as decode
from functools import wraps

from relevanceai.package_utils.config import CONFIG
from relevanceai.package_utils.json_encoder import json_encoder
from relevanceai.package_utils.fire_and_forget import fire_and_forget


def is_tracking_enabled():
    if CONFIG.is_field("mixpanel.is_tracking_enabled", CONFIG.config):
        return CONFIG.get_field("mixpanel.is_tracking_enabled", CONFIG.config)


def get_json_size(json_obj):
    # Returns it in bytes
    return len(json.dumps(json_obj).encode("utf-8")) / 1024


TRANSIT_ENV_VAR = "_IS_ANALYTICS_IN_TRANSIT"


def track(func: Callable):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if os.getenv(TRANSIT_ENV_VAR) == "TRUE":
            return func(*args, **kwargs)

        os.environ[TRANSIT_ENV_VAR] = "TRUE"

        try:
            if is_tracking_enabled():

                @fire_and_forget
                def send_analytics():
                    properties = {}
                    if "firebase_uid" in kwargs:
                        user_id = kwargs["firebase_uid"]
                    elif hasattr(args[0], "firebase_uid"):
                        user_id = args[0].firebase_uid
                    else:
                        user_id = "firebase_uid_not_detected"

                    event_name = f"pysdk-{func.__name__}"

                    # kwargs.update(dict(zip(func.__code__.co_varnames, args)))
                    # all_kwargs = copy.deepcopy(kwargs)
                    # all_kwargs = all_kwargs.update(
                    #     dict(zip(func.__code__.co_varnames, args))
                    # )

                    # if hasattr(func, "__code__"):
                    additional_args = dict(zip(func.__code__.co_varnames, args))

                    if "dataset_id" in kwargs:
                        properties["dataset_id"] = kwargs["dataset_id"]

                    elif "dataset_id" in additional_args:
                        properties["dataset_id"] = additional_args["dataset_id"]

                    elif hasattr(args[0], "dataset_id"):
                        properties["dataset_id"] = args[0].dataset_id

                    elif "self" in additional_args:
                        if hasattr(additional_args["self"], "dataset_id"):
                            properties["dataset_id"] = additional_args[
                                "self"
                            ].dataset_id

                    full_properties = copy.deepcopy(properties)
                    full_properties.update(
                        {
                            "additional_args": additional_args,
                            "args": args,
                            "kwargs": kwargs,
                        }
                    )
                    if user_id is not None:
                        # TODO: Loop through the properties and remove anything
                        # greater than 5kb
                        if (
                            "insert" in event_name
                            or "upsert" in event_name
                            or "update" in event_name
                            or "fit" in event_name
                            or "predict" in event_name
                            or get_json_size(
                                json_encoder(full_properties, force_string=True)
                            )
                            > 30
                        ):
                            response = analytics.track(
                                user_id=user_id,
                                event=event_name,
                                properties=json_encoder(properties, force_string=True),
                            )
                        else:
                            response = analytics.track(
                                user_id=user_id,
                                event=event_name,
                                properties=json_encoder(
                                    full_properties, force_string=True
                                ),
                            )

                send_analytics()
                # asyncio.ensure_future(send_analytics())
        except Exception as e:
            pass
        try:
            return func(*args, **kwargs)
        finally:
            os.environ[TRANSIT_ENV_VAR] = "FALSE"

    return wrapper


def track_event_usage(event_name: str):
    EVENT_NAME = event_name
    write_key = CONFIG.get_field("mixpanel.write_key", CONFIG.config)
    analytics.write_key = decode(write_key).decode("utf-8")

    def track(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if os.getenv(TRANSIT_ENV_VAR) == "TRUE":
                return func(*args, **kwargs)
            os.environ[TRANSIT_ENV_VAR] = "TRUE"
            try:
                if is_tracking_enabled():

                    async def send_analytics():
                        user_id = "OPEN_SOURCE_USER"
                        event_name = f"pysdk-{EVENT_NAME}"

                        if user_id is not None:
                            # TODO: Loop through the properties and remove anything
                            # greater than 5kb
                            response = analytics.track(
                                user_id=user_id,
                                event=event_name,
                            )

                    asyncio.ensure_future(send_analytics())
            except Exception as e:
                pass
            try:
                return func(*args, **kwargs)
            finally:
                os.environ[TRANSIT_ENV_VAR] = "FALSE"

        return wrapper

    return track


def identify(func: Callable):
    def wrapper(*args, **kwargs):
        try:
            if is_tracking_enabled():
                user_id = args[0].firebase_uid
                region = args[0].region
                traits = {
                    "region": region,
                }
                if user_id is not None:
                    analytics.identify(user_id, traits)
        except Exception as e:
            pass

        return func(*args, **kwargs)

    return wrapper
