# retoor <retoor@molodetz.nl>

from __future__ import annotations

import contextvars
from dataclasses import dataclass, field
from typing import Any, Optional

_active_state: contextvars.ContextVar = contextvars.ContextVar(
    "devii_agent_state", default=None
)
_eval_depth: contextvars.ContextVar = contextvars.ContextVar(
    "devii_eval_depth", default=0
)

MAX_EVAL_DEPTH = 2


@dataclass
class AgentState:
    plan: Optional[dict[str, Any]] = None
    reflections: list[dict[str, str]] = field(default_factory=list)
    iteration: int = 0
    verified: bool = False
    gate_triggered: bool = False
    mutations: set[str] = field(default_factory=set)


def get_state() -> Optional[AgentState]:
    return _active_state.get()


def set_state(state: AgentState) -> contextvars.Token:
    return _active_state.set(state)


def reset_state(token: contextvars.Token) -> None:
    _active_state.reset(token)


def record_mutation(name: str) -> None:
    state = _active_state.get()
    if state is not None:
        state.mutations.add(name)


def get_eval_depth() -> int:
    return _eval_depth.get()


def set_eval_depth(value: int) -> contextvars.Token:
    return _eval_depth.set(value)


def reset_eval_depth(token: contextvars.Token) -> None:
    _eval_depth.reset(token)
