# Copyright (c) 2025 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 inspect
import threading
import typing
from datetime import timedelta
from functools import cache, wraps
from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Set, Tuple, Type, TypeVar, Union
from uuid import uuid4

import cachetools
from boto3.dynamodb.conditions import Key
from cachetools import keys
from toolz import get_in, keyfilter
from typeguard import TypeCheckError, check_type

from czc import dynamodb as dyn
from czc import feature_data, logging, type_util
from czc.helpers import coalesce
from czc.string import get_md5
from czc.type_util import deserializable

KEYSPACE = 'fast_distributed_cache'
TOKEN_KEY = 'token'
ORGANIZATION_KEY = 'organization'
RESULT_KEY = 'result'
ORGANIZATION_NONE = 'none'
PERSISTENCE_TTL = timedelta(days=7)

logger = logging.Logger()

T = TypeVar('T')


@deserializable()
class _CachedResult(NamedTuple):
    token: str
    parameters: Optional[dict]
    data: Any


result_store = feature_data.for_type(_CachedResult)


@cache
def _dynamo_table():
    table_name = dyn.get_table_name()
    return dyn.table(table_name)


def _get_return_type(func: Callable[..., T]) -> Optional[Type[T]]:
    hints = typing.get_type_hints(func)
    return hints.get('return')


def _safe_coerce(return_type: Optional[Type[T]], data: Any) -> Optional[T]:
    def is_return_type() -> bool:
        try:
            check_type(data, return_type)
        except TypeCheckError:
            return False
        return True

    if return_type is None or is_return_type():
        return data

    try:
        return type_util.coerce(return_type, data)
    except Exception:
        logger.warning(
            f'Failed to coerce data to {return_type}. Ignoring cache hit. '
            f'Could be stale data incompatible with the latest deployed code.',
            exc_info=True,
        )
        return None


def _get_function_args(func: Callable, args: Tuple, kwargs: Dict[str, Any]) -> Dict[str, Any]:
    sig = inspect.signature(func)
    bound_args = sig.bind(*args, **kwargs)
    bound_args.apply_defaults()
    return bound_args.arguments


def _compose_token_sk(organization_id: str) -> str:
    return dyn.compose_key(ORGANIZATION_KEY, organization_id, TOKEN_KEY)


def _decompose_token_sk(sk: str) -> str:
    return list(dyn.decompose_key(sk))[1]


def _get_cache_token(name: str, organization_id: str, invalidate: bool = False) -> str:
    token_key = {'PK': dyn.prefix(KEYSPACE, name), 'SK': dyn.compose_key(ORGANIZATION_KEY, organization_id, TOKEN_KEY)}

    def set_new():
        token = str(uuid4())
        dyn.put_item(
            _dynamo_table(),
            Item=dict(**token_key, data=token),
        )
        return token

    if invalidate:
        return set_new()

    item = _dynamo_table().get_item(Key=token_key, ConsistentRead=True)
    token = get_in(['Item', 'data'], item)
    if token is None:
        return set_new()

    return token


def _validate_name(name: str) -> None:
    if dyn.COMPOSITE_CHAR in name:
        raise ValueError(f'Cache name must not contain the character: {dyn.COMPOSITE_CHAR}')


def cached(
    name: str,
    memcache_maxsize: int = 1000,
    key=keys.hashkey,
    save_params: Union[Set[str], Literal['*'], None] = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    _validate_name(name)

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        cache = cachetools.LRUCache(maxsize=memcache_maxsize)
        lock = threading.Lock()

        @wraps(func)
        def wrapper(*args, **kwargs):
            params = _get_function_args(func, args, kwargs)
            organization_id = coalesce(
                params.get('organization_id'), params.get('cz_organization_id'), params.get('organizationId'), ORGANIZATION_NONE
            )
            obj_key = get_md5(repr(key(*args, **kwargs)))
            result_key = (organization_id, KEYSPACE, name, obj_key, RESULT_KEY)
            store_param_keys = set(params.keys()) if save_params == '*' else save_params
            store_params = keyfilter(lambda k: k in store_param_keys, params) if store_param_keys else None

            @cachetools.cached(
                cache=cache,
                key=lambda t, a, k: keys.hashkey(t, key(*a, **k)),
                lock=lock,
            )
            def cached_func(token: str, args, kwargs):
                conn = feature_data.connect(ttl=PERSISTENCE_TTL)
                cached_result = result_store.get(conn, *result_key)
                if cached_result and cached_result.token == token:
                    if cached_result.data is None:
                        return None

                    return_type = _get_return_type(func)
                    typed_result = _safe_coerce(return_type, cached_result.data)
                    if typed_result is not None:
                        return typed_result

                result = func(*args, **kwargs)
                result_store.set(conn, *result_key, data=_CachedResult(token=token, parameters=store_params, data=result))
                return result

            token = _get_cache_token(name, organization_id)
            return cached_func(token, args, kwargs)

        return wrapper

    return decorator


def invalidate(name: str, organization_id: Optional[str] = None) -> None:
    _validate_name(name)

    def get_orgs():
        if organization_id:
            return [organization_id]

        cond = Key('PK').eq(dyn.prefix(KEYSPACE, name)) & Key('SK').begins_with(dyn.compose_key(ORGANIZATION_KEY, ''))
        items = dyn.query(_dynamo_table(), dict(KeyConditionExpression=cond, ConsistentRead=True))
        return [_decompose_token_sk(x['SK']) for x in items]

    for org in get_orgs():
        _get_cache_token(name, org, invalidate=True)
