from __future__ import annotations

from typing import Any, Optional
from .transport import patch


class Step:
    def __init__(self, trace_id: str, step_id: str) -> None:
        self.trace_id = trace_id
        self.step_id = step_id

    async def end(
        self,
        *,
        output: Any = None,
        status: str = "success",
        error_message: Optional[str] = None,
        model: Optional[str] = None,
        prompt_tokens: Optional[int] = None,
        completion_tokens: Optional[int] = None,
    ) -> None:
        body: dict[str, Any] = {"output": output, "status": status}
        if error_message:
            body["error_message"] = error_message
        if model:
            body["model"] = model
        if prompt_tokens is not None:
            body["prompt_tokens"] = prompt_tokens
        if completion_tokens is not None:
            body["completion_tokens"] = completion_tokens

        await patch(f"/api/v1/traces/{self.trace_id}/steps/{self.step_id}", body)
