"""Runtime audit context shared by user-facing PMFlow AI and LLM logs."""

from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from typing import Iterator


_agent_run_id: ContextVar[str | None] = ContextVar("pmflow_audit_agent_run_id", default=None)
_agent_step_key: ContextVar[str | None] = ContextVar("pmflow_audit_agent_step_key", default=None)


@contextmanager
def bind_agent_audit_context(*, agent_run_id: str | None = None, agent_step_key: str | None = None) -> Iterator[None]:
    run_token = _agent_run_id.set(str(agent_run_id or "").strip() or None) if agent_run_id is not None else None
    step_token = _agent_step_key.set(str(agent_step_key or "").strip() or None) if agent_step_key is not None else None
    try:
        yield
    finally:
        if step_token is not None:
            _agent_step_key.reset(step_token)
        if run_token is not None:
            _agent_run_id.reset(run_token)


def current_agent_audit_context() -> dict[str, str]:
    payload: dict[str, str] = {}
    run_id = str(_agent_run_id.get() or "").strip()
    step_key = str(_agent_step_key.get() or "").strip()
    if run_id:
        payload["agent_run_id"] = run_id[:80]
    if step_key:
        payload["agent_step_key"] = step_key[:120]
    return payload


__all__ = ["bind_agent_audit_context", "current_agent_audit_context"]