import threading
import uuid
from functools import wraps
from typing import Any, Callable, Hashable, Tuple

from cachetools import cached, keys
from flask import current_app, g
from flask.ctx import AppContext
from loguru import logger

g_thread_local_store = threading.local()


def thread_specific_cache(*args: Any, **kwargs: Any) -> Callable:
    """
    A version of `cached()` that is local to a single thread. In other words,
    cache entries will only be valid in the thread where they were created.

    Args:
        *args: Will be forwarded to `cached()`.
        **kwargs: Will be forwarded to `cached()`.

    Returns:
        The wrapped function.

    """

    def _key_func(*args_: Any, **kwargs_: Any) -> Tuple[Hashable, ...]:
        base_hash = keys.hashkey(*args_, **kwargs_)

        if hasattr(g_thread_local_store, "thread_id"):
            # We already gave this thread a unique ID. Use that.
            thread_id = g_thread_local_store.thread_id
        else:
            # Give this thread a new unique ID.
            thread_id = uuid.uuid4().hex
            g_thread_local_store.thread_id = thread_id

        return (thread_id,) + base_hash

    return cached(*args, **kwargs, key=_key_func)


def thread_with_app_context(to_wrap: Callable) -> Callable:
    """
    Decorator that wraps the entry point to a thread and injects the current
    app context from Flask. This is useful when we want to use multiple
    threads to handle a single request.

    When using this wrapped function, `current_app.app_context()` should be
    passed as the first argument when initializing the thread.

    Args:
        to_wrap: The function to wrap.

    Returns:
        The wrapped function.

    """

    @wraps(to_wrap)
    def _run_with_context(
        app_context: AppContext | None, *args: Any, **kwargs: Any
    ) -> Any:
        if app_context is None:
            # Do nothing.
            return to_wrap(*args, **kwargs)

        with app_context:
            return to_wrap(*args, **kwargs)

    return _run_with_context


def thread_context() -> AppContext | None:
    """
    Pushes a new app context for a thread that is being spawned to handle the
    current request. Will copy all the global data from the current context.

    Returns:
        The new context, or None if no context is active.

    """
    # Copy global data.
    global_data = {}
    try:
        for key in g:
            global_data[key] = g.get(key)
    except (TypeError, RuntimeError):
        # Context is not initialized (TypeError) or no request/app context
        # is active (RuntimeError). Don't change anything.
        logger.debug(
            "No active request context, skipping Flask g copy for child thread"
        )

    try:
        context = current_app.app_context()
    except RuntimeError:
        # Context is not initialized.
        logger.debug("No current app context, not passing to thread.")
        return None

    with context:
        for key, value in global_data.items():
            setattr(g, key, value)

    return context
