from __future__ import annotations

import uuid
from typing import Any

from nees.types import NEESTrace, TraceStep


class TraceRecorder:
    def __init__(self, mode: str, llm: str, enable_trace: bool) -> None:
        self._mode = mode
        self._llm = llm
        self._enable_trace = enable_trace
        self._trace_id = str(uuid.uuid4())
        self._steps: list[TraceStep] = []

    @property
    def trace_id(self) -> str:
        return self._trace_id

    def record(self, name: str, payload: dict[str, Any]) -> None:
        if not self._enable_trace:
            return
        self._steps.append(TraceStep(name=name, payload=payload))

    def build(self, summary: dict[str, Any]) -> NEESTrace:
        return NEESTrace(
            trace_id=self._trace_id,
            mode=self._mode,
            llm=self._llm,
            steps=list(self._steps),
            summary=summary,
        )
