# kybernis/spans.py
from __future__ import annotations

import time
import uuid
from contextlib import contextmanager
from typing import Any, Dict, Optional

from .client import KybernisClient
from .context import kyb_run_id, kyb_span_id, kyb_parent_span_id
from .runtime import kyb_client

KSP_VERSION = "ksp.v1"


def _now_rfc3339() -> str:
    # validator.go uses RFC3339 (no nanos requirement enforced there)
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _new_span_id() -> str:
    return "span_" + uuid.uuid4().hex[:16]


def _uuid_or_empty(value: Any) -> str:
    raw = str(value or "").strip()
    if not raw:
        return ""
    try:
        return str(uuid.UUID(raw))
    except Exception:
        return ""


def emit_ksp_event(
    *,
    client: KybernisClient,
    event_type: str,
    run_id: str,
    payload: Dict[str, Any],
) -> None:
    """
    Emits a KSP span event through KybernisClient.emit_event.
    We send type=event_type and payload=flat KSP fields.
    """
    # Ensure KSP-required fields exist (validator.go will also fill)
    p = dict(payload)
    p.setdefault("ksp_version", KSP_VERSION)
    p.setdefault("event_type", event_type)
    p.setdefault("ts", _now_rfc3339())
    p.setdefault("run_id", run_id)

    # IMPORTANT: client.emit_event body is:
    # {run_id, parent_id, type, payload}
    # API expects outer parent_id to be UUID. KSP parent_span_id can remain
    # a non-UUID span token in payload.
    outer_parent_id = _uuid_or_empty(p.get("parent_span_id") or "")
    client.emit_event(run_id, event_type, p, parent_id=outer_parent_id)


def span_start(
    *,
    client: KybernisClient,
    run_id: str,
    span_type: str,
    name: str,
    tenant_id: str = "",
    span_id: Optional[str] = None,
    parent_span_id: str = "",
    attributes: Optional[Dict[str, Any]] = None,
    attempt: int = 1,
    idempotency_key: Optional[str] = None,
    input_hash: Optional[str] = None,
) -> str:
    sid = span_id or _new_span_id()
    payload: Dict[str, Any] = {
        "tenant_id": tenant_id or "public",   # validator warns/fixes if missing
        "span_id": sid,
        "parent_span_id": parent_span_id or None,
        "span_type": span_type,
        "name": name,
        "attempt": int(attempt) if attempt else 1,
        "attributes": attributes or {},
    }
    if idempotency_key:
        payload["idempotency_key"] = idempotency_key
    if input_hash:
        payload["input_hash"] = input_hash

    emit_ksp_event(client=client, event_type="SPAN_START", run_id=run_id, payload=payload)
    return sid


def span_end(
    *,
    client: KybernisClient,
    run_id: str,
    span_id: str,
    tenant_id: str = "",
    status: str = "ok",
    latency_ms: int = 0,
    parent_span_id: str = "",
    attributes: Optional[Dict[str, Any]] = None,
    error_code: Optional[str] = None,
    error_message: Optional[str] = None,
    output_hash: Optional[str] = None,
    cost_micros: Optional[int] = None,
    usage: Optional[Dict[str, Any]] = None,
) -> None:
    payload: Dict[str, Any] = {
        "tenant_id": tenant_id or "public",
        "span_id": span_id,
        "parent_span_id": parent_span_id or None,
        "status": status,
        "latency_ms": int(latency_ms) if latency_ms >= 0 else 0,
        "attributes": attributes or {},
    }
    if error_code:
        payload["error_code"] = error_code
    if error_message:
        payload["error_message"] = error_message
    if output_hash:
        payload["output_hash"] = output_hash
    if cost_micros is not None:
        payload["cost_micros"] = int(cost_micros)
    if usage is not None:
        payload["usage"] = usage

    emit_ksp_event(client=client, event_type="SPAN_END", run_id=run_id, payload=payload)


@contextmanager
def span(
    *,
    client: Optional[KybernisClient] = None,
    span_type: str,
    name: str,
    tenant_id: str = "",
    attributes: Optional[Dict[str, Any]] = None,
    attempt: int = 1,
    idempotency_key: Optional[str] = None,
    input_hash: Optional[str] = None,
):
    """
    ContextVar-aware span:
    - run_id from kyb_run_id (must be set by agent patch) else generates a local one.
    - parent span from kyb_span_id (current span becomes child).
    - sets kyb_span_id for the duration of the block.
    """
    c = client or kyb_client()

    run_id = kyb_run_id.get()
    if not run_id:
        raise RuntimeError("kyb_run_id is not set; cannot start span before run is bound")
    parent = kyb_span_id.get() or ""
    sid = _new_span_id()

    start_ts = time.time()

    # set span context for nested spans
    tok_parent = kyb_parent_span_id.set(parent or None)
    tok_span = kyb_span_id.set(sid)

    try:
        span_start(
            client=c,
            run_id=run_id,
            tenant_id=tenant_id,
            span_id=sid,
            parent_span_id=parent,
            span_type=span_type,
            name=name,
            attributes=attributes,
            attempt=attempt,
            idempotency_key=idempotency_key,
            input_hash=input_hash,
        )
        yield sid
        latency = int((time.time() - start_ts) * 1000)
        span_end(
            client=c,
            run_id=run_id,
            tenant_id=tenant_id,
            span_id=sid,
            parent_span_id=parent,
            status="ok",
            latency_ms=latency,
        )
    except Exception as exc:
        latency = int((time.time() - start_ts) * 1000)
        span_end(
            client=c,
            run_id=run_id,
            tenant_id=tenant_id,
            span_id=sid,
            parent_span_id=parent,
            status="error",
            latency_ms=latency,
            error_message=str(exc),
            attributes={"error": str(exc)},
        )
        raise
    finally:
        kyb_span_id.reset(tok_span)
        kyb_parent_span_id.reset(tok_parent)
