from __future__ import annotations

from contextvars import ContextVar
from dataclasses import dataclass, replace

from app_runtime.contracts.trace import TraceContextRecord


@dataclass(frozen=True)
class ActiveTraceContext:
    record: TraceContextRecord
    step: str = ""


class TraceContextStore:
    def __init__(self) -> None:
        self._current: ContextVar[ActiveTraceContext | None] = ContextVar("trace_current", default=None)
        self._stack: ContextVar[tuple[ActiveTraceContext, ...]] = ContextVar("trace_stack", default=())

    def current(self) -> ActiveTraceContext | None:
        return self._current.get()

    def current_trace_id(self) -> str | None:
        active = self.current()
        return active.record.trace_id if active else None

    def push(self, record: TraceContextRecord) -> ActiveTraceContext:
        active = self.current()
        stack = self._stack.get()
        if active is not None:
            self._stack.set(stack + (active,))
        updated = ActiveTraceContext(record=record)
        self._current.set(updated)
        return updated

    def pop(self) -> ActiveTraceContext | None:
        stack = self._stack.get()
        if not stack:
            self._current.set(None)
            return None
        previous = stack[-1]
        self._stack.set(stack[:-1])
        self._current.set(previous)
        return previous

    def set_step(self, step: str) -> ActiveTraceContext | None:
        active = self.current()
        if active is None:
            return None
        updated = replace(active, step=step)
        self._current.set(updated)
        return updated