from __future__ import annotations

from dataclasses import dataclass, field, replace
import contextvars
from typing import Any, Callable, Iterable
from uuid import uuid4

from .audit import AuditEvent, emit_event, tags_to_strings
from .config import RegulatorConfig, config_scope, get_config


@dataclass
class RunContext:
    run_id: str
    agent_id: str | None
    active_tags: set[Any] = field(default_factory=set)
    context_id: str = field(default_factory=lambda: str(uuid4()))

    def activate(self, tags: Iterable[Any]) -> None:
        self.active_tags.update(tags)

    def audit(self, event: AuditEvent) -> None:
        if event.context_id is None:
            event = replace(event, context_id=self.context_id)
        emit_event(get_config().audit_sink, event)


_CURRENT_CONTEXT: contextvars.ContextVar[RunContext | None] = contextvars.ContextVar(
    "rgltr_context",
    default=None,
)


def get_current_context() -> RunContext | None:
    return _CURRENT_CONTEXT.get()


class context:
    def __init__(
        self,
        *,
        agent_id: str | None = None,
        run_id: str | None = None,
        config: RegulatorConfig | None = None,
    ) -> None:
        self._agent_id = agent_id
        self._run_id = run_id
        self._token: contextvars.Token[RunContext | None] | None = None
        self._config_scope = config_scope(config)

    def _make_context(self) -> RunContext:
        parent = _CURRENT_CONTEXT.get()
        tags = set(parent.active_tags) if parent else set()
        return RunContext(
            run_id=self._run_id or str(uuid4()),
            agent_id=self._agent_id,
            active_tags=tags,
            context_id=str(uuid4()),
        )

    def __enter__(self) -> RunContext:
        self._config_scope.__enter__()
        ctx = self._make_context()
        self._token = _CURRENT_CONTEXT.set(ctx)
        return ctx

    def __exit__(self, *_exc: object) -> None:
        if self._token is not None:
            _CURRENT_CONTEXT.reset(self._token)
        self._config_scope.__exit__(*([None] * 3))

    async def __aenter__(self) -> RunContext:
        self._config_scope.__enter__()
        ctx = self._make_context()
        self._token = _CURRENT_CONTEXT.set(ctx)
        return ctx

    async def __aexit__(self, *_exc: object) -> None:
        if self._token is not None:
            _CURRENT_CONTEXT.reset(self._token)
        self._config_scope.__exit__(*([None] * 3))


def restore_tags_from_messages(
    messages: Iterable[Any],
    rule_lookup: Callable[[str], Iterable[Any] | None],
    *,
    target: RunContext | None = None,
) -> set[Any]:
    ctx = target or get_current_context()
    if ctx is None:
        return set()

    restored: set[Any] = set()
    for message in messages:
        parts = getattr(message, "parts", ())
        for part in parts:
            tool_name = getattr(part, "tool_name", None)
            if not tool_name:
                continue
            activates = rule_lookup(tool_name)
            if activates:
                restored.update(activates)
    if restored:
        ctx.activate(restored)

    if restored:
        ctx.audit(
            AuditEvent(
                event_type="tags_restored",
                tool_name=None,
                agent_id=ctx.agent_id,
                run_id=ctx.run_id,
                active_tags=tags_to_strings(restored),
            )
        )
    return restored
