# kybernis/runtime.py
from __future__ import annotations

import uuid
from functools import lru_cache
from typing import Dict, Optional

from .client import KybernisClient
from .context import kyb_run_id, kyb_session_id, kyb_span_id


def _new_id(prefix: str = "") -> str:
    x = uuid.uuid4().hex
    return f"{prefix}{x}" if prefix else x


def ensure_run_id() -> str:
    run = kyb_run_id.get()
    if run:
        return run
    run = _new_id("kyb_")
    kyb_run_id.set(run)
    return run


def ensure_session_id() -> str:
    sess = kyb_session_id.get()
    if sess:
        return sess
    sess = ensure_run_id()  # sane default
    kyb_session_id.set(sess)
    return sess


def ensure_span_id() -> str:
    sp = kyb_span_id.get()
    if sp:
        return sp
    sp = _new_id("sp_")
    kyb_span_id.set(sp)
    return sp


@lru_cache(maxsize=1)
def kyb_client() -> KybernisClient:
    return KybernisClient()


def kyb_headers(existing: Optional[Dict[str, str]] = None) -> Dict[str, str]:
    headers = dict(existing or {})
    headers.setdefault("X-Kyb-Run-Id", ensure_run_id())
    headers.setdefault("X-Kyb-Session-Id", ensure_session_id())
    headers.setdefault("X-Kyb-Span-Id", ensure_span_id())
    return headers
