"""LumaTrack success-callback for the LiteLLM proxy/SDK. Every completed
LLM call becomes a LumaTrack run event with token usage attached; LumaTrack
prices it (or trusts LiteLLM's own response_cost, stamped "client") and
books it into the value ledger.

Usage: place this file next to your proxy config (or anywhere on
PYTHONPATH), then in the LiteLLM proxy config.yaml:

    litellm_settings:
      callbacks: ["litellm_lumatrack.LumaTrackLogger"]

    environment:
      LUMATRACK_URL: https://your-lumatrack-host
      LUMATRACK_KEY: lmt_...            # an ingest-scope key is enough

Route calls to an automation by putting ``lumatrack_automation: <slug>`` in
the request metadata (per key, per team, or per request). Calls without a
mapping are skipped — LumaTrack never guesses which automation earned what.

Flat-rate traffic: set ``lumatrack_billing: subscription`` in the request
metadata (per call), or ``LUMATRACK_AI_BILLING=subscription`` in the proxy
environment (whole-proxy default) when the traffic rides a flat-rate plan —
LumaTrack then books $0 marginal cost and tracks the API-equivalent as
repricing exposure.
"""

import json
import os
import urllib.request

__all__ = ["build_run_payload", "LumaTrackLogger"]


def build_run_payload(kwargs, response_obj, default_billing=None):
    """Map LiteLLM's logging kwargs + response into a LumaTrack run-event
    payload, or None when the call carries no automation mapping. Pure
    function; unit-tested without litellm installed.

    Token contract: OpenAI-shaped ``prompt_tokens`` INCLUDE cached tokens;
    LumaTrack's ``input_tokens`` are non-cached, so the cached count is
    subtracted out and reported separately.

    Billing basis: ``lumatrack_billing`` in the request metadata wins, else
    ``default_billing`` (the logger passes LUMATRACK_AI_BILLING, for proxies
    whose whole traffic rides a flat-rate subscription). Values pass through
    verbatim; the server validates."""
    metadata = (kwargs.get("litellm_params") or {}).get("metadata") or {}
    automation = metadata.get("lumatrack_automation")
    if not automation:
        return None
    billing = metadata.get("lumatrack_billing") or default_billing
    usage = {}
    if isinstance(response_obj, dict):
        usage = response_obj.get("usage") or {}
    elif response_obj is not None and hasattr(response_obj, "usage"):
        raw = response_obj.usage
        usage = raw if isinstance(raw, dict) else getattr(raw, "__dict__", {}) or {}
    prompt = int(usage.get("prompt_tokens") or 0)
    completion = int(usage.get("completion_tokens") or 0)
    details = usage.get("prompt_tokens_details") or {}
    cached = int((details or {}).get("cached_tokens") or 0)
    status = (kwargs.get("standard_logging_object") or {}).get("status", "success")
    payload = {
        "automation": str(automation),
        "status": "failure" if status == "failure" else "success",
        "source": "litellm",
        "ai": {
            "model": str(kwargs.get("model", "")),
            "input_tokens": max(prompt - cached, 0),
            "output_tokens": completion,
            "cached_tokens": cached,
        },
    }
    response_id = None
    if isinstance(response_obj, dict):
        response_id = response_obj.get("id")
    elif response_obj is not None:
        response_id = getattr(response_obj, "id", None)
    if response_id:
        payload["external_id"] = str(response_id)
    if kwargs.get("response_cost") is not None:
        payload["ai"]["cost"] = str(kwargs["response_cost"])
    if billing:
        payload["ai"]["billing"] = str(billing)
    return payload


def _post(payload):
    base = os.environ.get("LUMATRACK_URL", "").rstrip("/")
    key = os.environ.get("LUMATRACK_KEY", "")
    if not base or not key:
        return
    request = urllib.request.Request(
        f"{base}/api/v1/runs",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"},
        method="POST",
    )
    try:
        urllib.request.urlopen(request, timeout=10)  # noqa: S310
    except Exception:  # never break the caller's LLM traffic over telemetry
        pass


try:  # pragma: no cover - exercised only where litellm is installed
    from litellm.integrations.custom_logger import CustomLogger

    class LumaTrackLogger(CustomLogger):
        def log_success_event(self, kwargs, response_obj, start_time, end_time):
            try:
                payload = build_run_payload(
                    kwargs, response_obj, os.environ.get("LUMATRACK_AI_BILLING")
                )
                if payload:
                    _post(payload)
            except Exception:  # telemetry must never break the caller's traffic
                pass

        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
            try:
                payload = build_run_payload(
                    kwargs, response_obj, os.environ.get("LUMATRACK_AI_BILLING")
                )
                if payload:
                    payload["status"] = "failure"
                    _post(payload)
            except Exception:
                pass

except ImportError:  # LumaTrack's own test suite runs without litellm

    class LumaTrackLogger:  # type: ignore[no-redef]
        pass