Rootdata

"""Cost tracking helper for skill subprocesses.

Skills that call sc-proxy via plain `requests` need to:
  1. Tag every paid call with a SC-CALLER-ID that ties it back to the user
     turn that triggered the skill (so the agent's per-turn cost summary
     shows the cost in the right cost card).
  2. After each call, parse the sc-proxy response headers
     (`X-Credits-Used`, `X-Credits-Api-Type`) and write a row to the cost
     ledger that the agent reads back when it builds the SSE
     `cost_summary` event.

This file is intentionally zero-dependency (stdlib only) so it can be
dropped into any skill folder without coupling to starchild-clawd internals.

Env vars consumed (set by the agent before dispatching the bash subprocess):
  - STARCHILD_TOOL_CALLER_ID  — opaque tag for the current tool call
  - STARCHILD_USER_TURN_ID    — uuid of the current user turn
  - STARCHILD_COST_LEDGER_DIR — optional override for ledger directory

When env vars are absent (e.g. running the script outside an agent), the
helpers degrade gracefully: caller-id falls back to a synthetic string so
the call still goes through, and ledger writes still happen for audit but
the user-turn reader will skip them.
"""
from __future__ import annotations

import fcntl
import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional
from urllib.parse import urlparse


_DEFAULT_LEDGER_DIR = "/data/.starchild/cost_ledger"

# Allowlisted request payload keys we forward into the ledger row's
# `details` field. MUST stay in sync with starchild-clawd's
# core/http_client._record_cost_to_ledger allowlist — anything not in
# that allowlist won't be picked up by the agent and won't render in
# the frontend cost card.
_PAYLOAD_ALLOWLIST = (
    # Identity
    "model", "provider",
    # Image geometry
    "aspect_ratio", "quality", "resolution", "image_size", "size",
    # Video / motion
    "duration", "duration_s", "fps", "motion_strength",
    # Quantity
    "n", "count",
    # Generation knobs
    "seed", "steps", "guidance_scale", "cfg_scale", "strength",
    "scheduler", "sampler",
    # Reference / mode hints
    "image_to_image", "image_to_video", "use_reference", "reference_count",
)


def caller_headers(extra: Optional[Dict[str, str]] = None,
                   tool_default: str = "skill") -> Dict[str, str]:
    """Return an HTTP-headers dict with SC-CALLER-ID filled in.

    Resolution order:
      1. `extra["SC-CALLER-ID"]` (case-insensitive) — caller wins.
      2. STARCHILD_TOOL_CALLER_ID env (set by the agent)
      3. Synthetic `f"{tool_default}:{int(time.time())}"` — tags the call so
         charges are attributable to *some* identifier even when the agent
         didn't inject one (standalone CLI runs, tests, cron).
    """
    merged: Dict[str, str] = dict(extra or {})
    has_caller = any(k.lower() == "sc-caller-id" for k in merged)
    if not has_caller:
        cid = os.environ.get("STARCHILD_TOOL_CALLER_ID") \
            or f"{tool_default}:{int(time.time())}"
        merged["SC-CALLER-ID"] = cid
    return merged


def record_response(response,
                    request_url: str,
                    request_payload: Optional[Dict[str, Any]] = None,
                    api_type_hint: Optional[str] = None) -> None:
    """Inspect a sc-proxy response and append a ledger row when paid.

    Best-effort. Silently no-ops when:
      - response carries no X-Credits-Used / X-Credits-Api-Type
      - cost is 0 or unparseable
      - file write fails
    Never raises — must not break a real request flow.
    """
    try:
        headers = getattr(response, "headers", None) or {}
        used = headers.get("X-Credits-Used") or headers.get("x-credits-used")
        api_type = (headers.get("X-Credits-Api-Type")
                    or headers.get("x-credits-api-type")
                    or api_type_hint)
        if not used or not api_type:
            return
        try:
            cost_f = float(used)
        except (TypeError, ValueError):
            return
        if cost_f <= 0:
            return

        turn_id = os.environ.get("STARCHILD_USER_TURN_ID") or ""
        caller_id = os.environ.get("STARCHILD_TOOL_CALLER_ID") or ""

        host = ""
        try:
            host = urlparse(request_url).netloc or ""
        except Exception:
            pass

        details: Dict[str, Any] = {}
        if isinstance(request_payload, dict):
            for k in _PAYLOAD_ALLOWLIST:
                v = request_payload.get(k)
                if v not in (None, "", []):
                    details[k] = v

        # fal.ai puts the model in the URL path, not the body.
        if "model" not in details and api_type == "falai":
            try:
                path = urlparse(request_url).path or ""
                model_path = path.lstrip("/")
                if "/requests/" in model_path:
                    model_path = model_path.split("/requests/", 1)[0]
                if model_path and not model_path.startswith("requests/"):
                    details["model"] = model_path
                    details["provider"] = "fal"
            except Exception:
                pass

        _append_ledger(
            turn_id=turn_id,
            caller_id=caller_id,
            api_type=api_type,
            cost_usd=cost_f,
            url_host=host,
            details=details or None,
        )
    except Exception:
        # Never let cost tracking break the actual request.
        pass


def _ledger_dir() -> Path:
    base = os.environ.get("STARCHILD_COST_LEDGER_DIR") or _DEFAULT_LEDGER_DIR
    p = Path(base)
    try:
        p.mkdir(parents=True, exist_ok=True)
    except OSError:
        p = Path("/tmp/starchild_cost_ledger")
        p.mkdir(parents=True, exist_ok=True)
    return p


def _today_path() -> Path:
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    return _ledger_dir() / f"{today}.jsonl"


def _derive_tool(caller_id: str, api_type: str) -> str:
    """Match starchild-clawd's _derive_tool_from_caller fallback."""
    if not caller_id:
        return api_type or "unknown"
    # chat:{sid}/tool:{name} → name
    if "/tool:" in caller_id:
        return caller_id.rsplit("/tool:", 1)[-1] or api_type
    # skill:{name} | job:{id} | video:{ts}
    head = caller_id.split(":", 1)[0]
    return head or api_type or "unknown"


def _append_ledger(*, turn_id: str, caller_id: str, api_type: str,
                   cost_usd: float, url_host: str,
                   details: Optional[Dict[str, Any]]) -> None:
    row = {
        "ts": round(time.time(), 3),
        "turn_id": turn_id,
        "caller_id": caller_id,
        "tool": _derive_tool(caller_id, api_type),
        "api_type": api_type or "unknown",
        "cost_usd": round(cost_usd, 8),
        "url_host": url_host or "",
    }
    if details:
        row["details"] = details

    line = json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n"
    path = _today_path()
    try:
        with open(path, "ab") as f:
            try:
                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
            except OSError:
                pass
            try:
                f.write(line.encode("utf-8"))
                f.flush()
                try:
                    os.fsync(f.fileno())
                except OSError:
                    pass
            finally:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
                except OSError:
                    pass
    except OSError:
        pass

"""
RootData skill exports (script-mode).

Usage:
    python3 - <<'EOF'
    import sys
    sys.path.insert(0, "/data/workspace/skills/rootdata")
    from exports import rd_search, rd_hot_index
    print(rd_search(query="berachain")[:1])
    print(rd_hot_index(days=1)[:3])
    EOF
"""

import os
import sys
from pathlib import Path

import requests

# Make _cost_track importable when this script is invoked from any CWD
# (mirrors the pattern from the video skill).
_HERE = Path(__file__).resolve().parent
if str(_HERE) not in sys.path:
    sys.path.insert(0, str(_HERE))

from _cost_track import caller_headers, record_response  # noqa: E402

BASE = "https://api.rootdata.com/open/skill"
TIMEOUT = 30


def _headers(language: str = "en"):
    key = os.environ.get("ROOTDATA_SKILL_KEY", "").strip()
    if not key:
        raise RuntimeError(
            "ROOTDATA_SKILL_KEY is not set. Call rd_init_key() and persist the returned api_key first."
        )
    base = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "language": language,
    }
    # Wrap with SC-CALLER-ID so paid calls are attributed to the user turn.
    return caller_headers(base, tool_default="rootdata")


def _post(path: str, body: dict, language: str = "en"):
    url = f"{BASE}/{path}"
    r = requests.post(
        url,
        headers=_headers(language=language),
        json=body,
        timeout=TIMEOUT,
    )
    # Record cost ledger row before raising — even error responses may have
    # been billed by the proxy. record_response() silently no-ops on non-paid.
    record_response(r, request_url=url, request_payload=body)
    r.raise_for_status()
    data = r.json()
    if isinstance(data, dict) and data.get("result") != 200:
        # upstream returns {result: <code>, error: ...} even on some non-200 cases
        raise RuntimeError(f"RootData API error: {data}")
    return data.get("data") if isinstance(data, dict) and "data" in data else data


def rd_init_key():
    """Initialize anonymous RootData key via /init. Returns {'api_key', 'message'}"""
    r = requests.post(f"{BASE}/init", json={}, timeout=TIMEOUT)
    r.raise_for_status()
    return r.json()


def rd_search(query: str, precise_x_search: bool = False, language: str = "en"):
    """Search entities by keyword. type: 1=project, 2=institution, 3=person."""
    return _post(
        "ser_inv",
        {"query": query, "precise_x_search": bool(precise_x_search)},
        language=language,
    )


def rd_id_map(type: int, language: str = "en"):
    """Get all IDs by type: 1=project, 2=institution, 3=person."""
    return _post("id_map", {"type": int(type)}, language=language)


def rd_project_detail(
    project_id: int = None,
    contract_address: str = None,
    include_investors: bool = True,
    language: str = "en",
):
    """Project detail by project_id or contract_address."""
    body = {"include_investors": bool(include_investors)}
    if project_id is not None:
        body["project_id"] = int(project_id)
    if contract_address:
        body["contract_address"] = contract_address
    if "project_id" not in body and "contract_address" not in body:
        raise ValueError("rd_project_detail requires project_id or contract_address")
    return _post("get_item", body, language=language)


def rd_funding_rounds(
    page: int = 1,
    page_size: int = 20,
    project_id: int = None,
    start_time: str = None,
    end_time: str = None,
    min_amount: float = None,
    max_amount: float = None,
    language: str = "en",
):
    """Funding rounds (past 365 days; max 3 investors per round)."""
    body = {"page": int(page), "page_size": int(page_size)}
    if project_id is not None:
        body["project_id"] = int(project_id)
    if start_time is not None:
        body["start_time"] = start_time
    if end_time is not None:
        body["end_time"] = end_time
    if min_amount is not None:
        body["min_amount"] = min_amount
    if max_amount is not None:
        body["max_amount"] = max_amount
    return _post("get_fac", body, language=language)


def rd_hot_index(days: int = 1, language: str = "en"):
    """Trending projects. days: 1=today, 7=this week."""
    if days not in (1, 7):
        raise ValueError("rd_hot_index days must be 1 or 7")
    return _post("hot_index", {"days": days}, language=language)


def rd_job_changes(
    recent_joinees: bool = True,
    recent_resignations: bool = True,
    language: str = "en",
):
    """Recent hires/departures. Max 20 entries per category."""
    return _post(
        "job_changes",
        {
            "recent_joinees": bool(recent_joinees),
            "recent_resignations": bool(recent_resignations),
        },
        language=language,
    )