from contextvars import ContextVar
from threading import Lock
from typing import TYPE_CHECKING

from contexts.wrapper import RecyclableContextVar

if TYPE_CHECKING:
    from core.datasource.__base.datasource_provider import DatasourcePluginProviderController
    from core.tools.plugin_tool.provider import PluginToolProviderController
    from core.trigger.provider import PluginTriggerProviderController


"""
To avoid race-conditions caused by gunicorn thread recycling, using RecyclableContextVar to replace with
"""
plugin_tool_providers: RecyclableContextVar[dict[str, "PluginToolProviderController"]] = RecyclableContextVar(
    ContextVar("plugin_tool_providers")
)

plugin_tool_providers_lock: RecyclableContextVar[Lock] = RecyclableContextVar(ContextVar("plugin_tool_providers_lock"))

datasource_plugin_providers: RecyclableContextVar[dict[str, "DatasourcePluginProviderController"]] = (
    RecyclableContextVar(ContextVar("datasource_plugin_providers"))
)

datasource_plugin_providers_lock: RecyclableContextVar[Lock] = RecyclableContextVar(
    ContextVar("datasource_plugin_providers_lock")
)

plugin_trigger_providers: RecyclableContextVar[dict[str, "PluginTriggerProviderController"]] = RecyclableContextVar(
    ContextVar("plugin_trigger_providers")
)

plugin_trigger_providers_lock: RecyclableContextVar[Lock] = RecyclableContextVar(
    ContextVar("plugin_trigger_providers_lock")
)