from __future__ import annotations

import re
from contextvars import ContextVar, Token

_SAFE_RUN_ID_RE = re.compile(r"[^A-Za-z0-9_-]+")
_MAX_RUN_ID_LEN = 80
_RUN_CONTEXT: ContextVar[dict[str, str]] = ContextVar("mugi_run_context", default={})


def sanitize_run_id(value: str) -> str:
    raw = str(value or "").strip()
    if not raw:
        return ""
    safe = _SAFE_RUN_ID_RE.sub("", raw)
    return safe[:_MAX_RUN_ID_LEN]


def get_run_context() -> dict[str, str]:
    ctx = _RUN_CONTEXT.get() or {}
    return dict(ctx)


def set_run_context(
    *,
    run_id: str = "",
    thread_id: str = "",
    parent_run_id: str = "",
) -> Token:
    payload: dict[str, str] = {}
    rid = sanitize_run_id(run_id)
    if rid:
        payload["run_id"] = rid
    tid = str(thread_id or "").strip()
    if tid:
        payload["thread_id"] = tid
    parent = sanitize_run_id(parent_run_id)
    if parent:
        payload["parent_run_id"] = parent
    return _RUN_CONTEXT.set(payload)


def clear_run_context(token: Token | None = None) -> None:
    try:
        if token is not None:
            _RUN_CONTEXT.reset(token)
            return
    except Exception:
        pass
    _RUN_CONTEXT.set({})
