from __future__ import annotations

from typing import Any, Optional
from .transport import post, patch
from .step import Step


class Trace:
    def __init__(self, trace_id: str) -> None:
        self.trace_id = trace_id
        self._step_counter = 0

    async def start_step(
        self,
        name: str,
        step_type: str = "llm_call",
        *,
        input: Any = None,
        metadata: Optional[dict[str, Any]] = None,
        model: Optional[str] = None,
    ) -> Step:
        self._step_counter += 1
        body: dict[str, Any] = {
            "step_number": self._step_counter,
            "name": name,
            "type": step_type,
        }
        if input is not None:
            body["input"] = input
        if metadata:
            body["metadata"] = metadata
        if model:
            body["model"] = model

        result = await post(f"/api/v1/traces/{self.trace_id}/steps", body)
        step_id = result.get("step", {}).get("id", "unknown") if result else "unknown"
        return Step(self.trace_id, step_id)

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

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