from __future__ import annotations
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Dict, List

from .exports import ToolExport, AgentExport, FlowExport, HttpHandlerExport, McpServerExport, McpToolMap

@dataclass
class Registry:
    tools: Dict[str, ToolExport] = field(default_factory=dict)
    agents: Dict[str, AgentExport] = field(default_factory=dict)
    flows: Dict[str, FlowExport] = field(default_factory=dict)
    http_handlers: List[HttpHandlerExport] = field(default_factory=list)
    mcp_servers: Dict[str, McpServerExport] = field(default_factory=dict)
    mcp_tool_fns: McpToolMap = field(default_factory=dict)

_current: ContextVar[Registry] = ContextVar("aiel_registry")

def get_registry() -> Registry:
    try:
        return _current.get()
    except LookupError:
        r = Registry()
        _current.set(r)
        return r

def reset_registry() -> Registry:
    r = Registry()
    _current.set(r)
    return r
