from __future__ import annotations

from dataclasses import dataclass
from typing import Callable


@dataclass(slots=True)
class ExecutionContext:
    session: dict
    run_budget: dict
    get_steps: Callable[[], list[dict]]
    get_memory: Callable[[], list[dict]]

    def remaining_tool_calls(self) -> int:
        return max(0, int(self.run_budget.get("max_tool_calls") or 0) - int(self.run_budget.get("tool_calls") or 0))

    def has_tool_budget(self, count: int = 1) -> bool:
        return self.remaining_tool_calls() >= max(1, int(count or 1))

    def consume_tool_call(self, count: int = 1) -> None:
        self.run_budget["tool_calls"] = int(self.run_budget.get("tool_calls") or 0) + max(1, int(count or 1))

    def set_trace(self, trace: list[dict] | None) -> None:
        self.session["last_tool_trace"] = list(trace or [])

    def set_result(
        self,
        goal: str,
        result: str,
        *,
        verdict: str = "DONE",
        completed: bool = True,
        trace: list[dict] | None = None,
        step_meta: dict | None = None,
        replace_steps: bool = True,
    ) -> None:
        if trace is not None:
            self.set_trace(trace)
        entry = {
            "step": goal,
            "result": result,
            "verdict": verdict,
        }
        meta = dict(step_meta or {})
        if meta:
            entry.update({
                "tool_name": meta.get("tool_name", ""),
                "tool_result": meta.get("tool_result", ""),
                "tool_details": dict(meta.get("tool_details") or {}),
                "deterministic_result": meta.get("deterministic_result", True),
            })
        if replace_steps:
            self.session["steps"] = [entry]
        self.session["final_report"] = result
        self.session["completed"] = completed