import collections.abc
import contextlib
import contextvars
import functools
import inspect
from collections.abc import Callable
from contextvars import ContextVar
from typing import ParamSpec, TypeVar, NoReturn, Any

from debputy.exceptions import (
    UnhandledOrUnexpectedErrorFromPluginError,
    DebputyRuntimeError,
)
from debputy.packages import SourcePackage, BinaryPackage
from debputy.util import _trace_log, _is_trace_log_enabled

_current_debputy_plugin_cxt_var: ContextVar[str | None] = ContextVar(
    "current_debputy_plugin",
    default=None,
)
_current_debputy_parsing_context: ContextVar[
    tuple[
        dict[tuple[SourcePackage | BinaryPackage, type[Any]], Any],
        SourcePackage | BinaryPackage,
    ]
    | None
] = ContextVar(
    "current_debputy_parsing_context",
    default=None,
)

P = ParamSpec("P")
R = TypeVar("R")


def register_manifest_type_value_in_context(
    value_type: type[Any],
    value: Any,
) -> None:
    context_vars = _current_debputy_parsing_context.get()
    if context_vars is None:
        raise AssertionError(
            "register_manifest_type_value_in_context() was called, but no context was set."
        )
    value_table, context_pkg = context_vars
    if (context_pkg, value_type) in value_table:
        raise AssertionError(
            f"The type {value_type!r} was already registered for {context_pkg}, which the plugin API should have prevented"
        )
    value_table[(context_pkg, value_type)] = value


def begin_parsing_context(
    value_table: dict[tuple[SourcePackage | BinaryPackage, type[Any]], Any],
    context_pkg: SourcePackage,
    func: Callable[P, R],
    *args: P.args,
    **kwargs: P.kwargs,
) -> R:
    context = contextvars.copy_context()
    # Wish we could just do a regular set without wrapping it in `context.run`
    context.run(_current_debputy_parsing_context.set, (value_table, context_pkg))
    assert context.get(_current_debputy_parsing_context) == (value_table, context_pkg)
    return context.run(func, *args, **kwargs)


@contextlib.contextmanager
def with_binary_pkg_parsing_context(
    context_pkg: BinaryPackage,
) -> collections.abc.Iterator[None]:
    context_vars = _current_debputy_parsing_context.get()
    if context_vars is None:
        raise AssertionError(
            "with_binary_pkg_parsing_context() was called, but no context was set."
        )
    value_table, _ = context_vars
    token = _current_debputy_parsing_context.set((value_table, context_pkg))
    try:
        yield
    finally:
        _current_debputy_parsing_context.reset(token)


def current_debputy_plugin_if_present() -> str | None:
    return _current_debputy_plugin_cxt_var.get()


def current_debputy_plugin_required() -> str:
    v = current_debputy_plugin_if_present()
    if v is None:
        raise AssertionError(
            "current_debputy_plugin_required() was called, but no plugin was set."
        )
    return v


def wrap_plugin_code(
    plugin_name: str,
    func: Callable[P, R],
    *,
    non_debputy_exception_handling: bool | Callable[[Exception], NoReturn] = True,
) -> Callable[P, R]:
    if isinstance(non_debputy_exception_handling, bool):

        runner = run_in_context_of_plugin
        if non_debputy_exception_handling:
            runner = run_in_context_of_plugin_wrap_errors

        def _plugin_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            return runner(plugin_name, func, *args, **kwargs)

        functools.update_wrapper(_plugin_wrapper, func)
        return _plugin_wrapper

    def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        try:
            return run_in_context_of_plugin(plugin_name, func, *args, **kwargs)
        except DebputyRuntimeError:
            raise
        except Exception as e:
            non_debputy_exception_handling(e)

    functools.update_wrapper(_wrapper, func)
    return _wrapper


def run_in_context_of_plugin(
    plugin: str,
    func: Callable[P, R],
    *args: P.args,
    **kwargs: P.kwargs,
) -> R:
    context = contextvars.copy_context()
    if _is_trace_log_enabled():
        call_stack = inspect.stack()
        caller: str = "[N/A]"
        for frame in call_stack:
            if frame.filename != __file__:
                try:
                    fname = frame.frame.f_code.co_qualname
                except AttributeError:
                    fname = None
                if fname is None:
                    fname = frame.function
                caller = f"{frame.filename}:{frame.lineno} ({fname})"
                break
        # Do not keep the reference longer than necessary
        del call_stack
        _trace_log(
            f"Switching plugin context to {plugin} at {caller} (from context: {current_debputy_plugin_if_present()})"
        )
    # Wish we could just do a regular set without wrapping it in `context.run`
    context.run(_current_debputy_plugin_cxt_var.set, plugin)
    return context.run(func, *args, **kwargs)


def run_in_context_of_plugin_wrap_errors(
    plugin: str,
    func: Callable[P, R],
    *args: P.args,
    **kwargs: P.kwargs,
) -> R:
    try:
        return run_in_context_of_plugin(plugin, func, *args, **kwargs)
    except DebputyRuntimeError:
        raise
    except Exception as e:
        if plugin != "debputy":
            raise UnhandledOrUnexpectedErrorFromPluginError(
                f"{func.__qualname__} from the plugin {plugin} raised exception that was not expected here."
            ) from e
        else:
            raise AssertionError(
                "Bug in the `debputy` plugin: Unhandled exception."
            ) from e
