"""Trace normalization — guarantee canonical shape regardless of backend.

Called at the storage boundary so every consumer sees the same fields.
Does NOT fill missing optional fields (model, tokens, cost) — those are
genuinely absent for non-LLM spans. Only guarantees required fields exist
with correct types.
"""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from typing import Any


_VALID_KINDS = frozenset({
    "llm", "tool", "retrieval", "agent", "chain",
    "embedding", "guardrail", "custom",
})

_VALID_STATUSES = frozenset({"ok", "error"})


def normalize_trace(raw: dict[str, Any]) -> dict[str, Any]:
    """Guarantee canonical trace shape.

    Required fields get defaults if missing. Optional fields pass through.
    Spans are normalized individually.
    """
    trace = dict(raw)  # shallow copy

    # Required fields
    if not trace.get("trace_id"):
        trace["trace_id"] = str(uuid.uuid4())
    if not trace.get("project"):
        trace["project"] = "unknown"
    if not trace.get("started_at"):
        trace["started_at"] = _now_iso()
    else:
        trace["started_at"] = _normalize_ts(trace["started_at"])
    if trace.get("status") not in _VALID_STATUSES:
        trace["status"] = "ok"

    # Normalize ended_at if present
    if trace.get("ended_at"):
        trace["ended_at"] = _normalize_ts(trace["ended_at"])

    # Ensure spans is a list
    spans = trace.get("spans")
    if not isinstance(spans, list):
        trace["spans"] = []
    else:
        trace["spans"] = [normalize_span(s) for s in spans]

    # Ensure tags is a list if present
    tags = trace.get("tags")
    if tags is not None and not isinstance(tags, list):
        trace["tags"] = []

    return trace


def normalize_span(raw: dict[str, Any]) -> dict[str, Any]:
    """Guarantee canonical span shape.

    Required fields get defaults if missing. LLM-specific fields pass through.
    """
    span = dict(raw)  # shallow copy

    # Required fields
    if not span.get("span_id"):
        span["span_id"] = str(uuid.uuid4())
    if "parent_id" not in span:
        span["parent_id"] = None
    if not span.get("name"):
        span["name"] = span.get("kind", "unknown")
    if span.get("kind") not in _VALID_KINDS:
        span["kind"] = "custom"
    if not span.get("started_at"):
        span["started_at"] = _now_iso()
    else:
        span["started_at"] = _normalize_ts(span["started_at"])
    if span.get("status") not in _VALID_STATUSES:
        span["status"] = "ok"

    # Normalize ended_at if present
    if span.get("ended_at"):
        span["ended_at"] = _normalize_ts(span["ended_at"])

    # Ensure numeric fields are correct type if present
    for field in ("input_tokens", "output_tokens"):
        val = span.get(field)
        if val is not None:
            try:
                span[field] = int(val)
            except (ValueError, TypeError):
                del span[field]

    cost = span.get("cost")
    if cost is not None:
        try:
            span["cost"] = float(cost)
        except (ValueError, TypeError):
            del span["cost"]

    return span


def _normalize_ts(ts: Any) -> str:
    """Normalize a timestamp to ISO-8601 with Z suffix."""
    if isinstance(ts, datetime):
        s = ts.isoformat()
        return s.replace("+00:00", "Z") if "+00:00" in s else s
    if isinstance(ts, str):
        # Ensure Z suffix consistency
        return ts.replace("+00:00", "Z") if "+00:00" in ts else ts
    return str(ts)


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
