from __future__ import annotations

import threading
from typing import Any

from utils.llm_token_normalize import normalize_stream_id_token

_lock = threading.Lock()
_state_by_stream: dict[str, dict[str, Any]] = {}


def _entry(stream_id: str) -> dict[str, Any]:
    return _state_by_stream.setdefault(
        stream_id,
        {
            "values": {},
            "actions": {},
        },
    )


def set_control_value(stream_id: str, control_id: str, value: Any) -> None:
    sid = normalize_stream_id_token((stream_id or "").strip())
    cid = (control_id or "").strip()
    if not sid or not cid:
        return
    with _lock:
        e = _entry(sid)
        e["values"][cid] = value


def get_control_value(stream_id: str, control_id: str, default: Any = None) -> Any:
    sid = normalize_stream_id_token((stream_id or "").strip())
    cid = (control_id or "").strip()
    if not sid or not cid:
        return default
    with _lock:
        e = _state_by_stream.get(sid)
        if not e:
            return default
        return e.get("values", {}).get(cid, default)


def trigger_action(stream_id: str, action_id: str) -> int:
    sid = normalize_stream_id_token((stream_id or "").strip())
    aid = (action_id or "").strip()
    if not sid or not aid:
        return 0
    with _lock:
        e = _entry(sid)
        actions = e["actions"]
        actions[aid] = int(actions.get(aid, 0)) + 1
        return int(actions[aid])


def get_action_count(stream_id: str, action_id: str) -> int:
    sid = normalize_stream_id_token((stream_id or "").strip())
    aid = (action_id or "").strip()
    if not sid or not aid:
        return 0
    with _lock:
        e = _state_by_stream.get(sid)
        if not e:
            return 0
        return int(e.get("actions", {}).get(aid, 0))


def clear_stream_state(stream_id: str) -> None:
    sid = normalize_stream_id_token((stream_id or "").strip())
    if not sid:
        return
    with _lock:
        _state_by_stream.pop(sid, None)
