"""Payload builders for Opik traces and spans."""

from datetime import datetime, timezone
from typing import Any, Final

from litellm import _logging
from litellm.integrations.opik import utils

from . import types


def build_trace_payload(
    project_name: str,
    trace_id: str,
    response_obj: dict[str, Any],
    start_time: datetime,
    end_time: datetime,
    input_data: Any,
    output_data: Any,
    metadata: dict[str, Any],
    tags: list[str],
    thread_id: str | None,
) -> types.TracePayload:
    """Build a complete trace payload."""
    trace_name: Final = response_obj.get("object", "unknown type")

    return types.TracePayload(
        project_name=project_name,
        id=trace_id,
        name=trace_name,
        start_time=(start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")),
        end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
        input=input_data,
        output=output_data,
        metadata=metadata,
        tags=tags,
        thread_id=thread_id,
    )


def build_span_payload(
    project_name: str,
    trace_id: str,
    parent_span_id: str | None,
    response_obj: dict[str, Any],
    start_time: datetime,
    end_time: datetime,
    input_data: Any,
    output_data: Any,
    metadata: dict[str, Any],
    tags: list[str],
    usage: dict[str, int],
    provider: str | None = None,
    cost: float | None = None,
) -> types.SpanPayload:
    """Build a complete span payload."""
    span_id: Final = utils.create_uuid7()

    model: Final = response_obj.get("model", "unknown-model")
    obj_type: Final = response_obj.get("object", "unknown-object")
    created: Final = response_obj.get("created", 0)
    span_name: Final = f"{model}_{obj_type}_{created}"

    _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id)

    return types.SpanPayload(
        id=span_id,
        project_name=project_name,
        trace_id=trace_id,
        parent_span_id=parent_span_id,
        name=span_name,
        type="llm",
        model=model,
        start_time=(start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")),
        end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
        input=input_data,
        output=output_data,
        metadata=metadata,
        tags=tags,
        usage=usage,
        provider=provider,
        total_cost=cost,
    )