#!/usr/bin/env python3
"""Minimal, dependency-free live-API helper used by examples, tests and agents.

    from scripts.live import openai_request, anthropic_request, load_env
    st, body, headers = openai_request("GET", "/v1/models")
    st, body, headers = anthropic_request("POST", "/v1/messages", json={...}, beta="prompt-caching-2024-07-31")

- Loads .env from the repo root (never prints keys).
- Every call is appended (sanitized) to reports/live-requests.jsonl.
- Response bodies are returned as parsed JSON when possible, else bytes.
- `stream=True` returns an iterator of decoded lines (SSE) instead of a body.

CLI:  python3 scripts/live.py openai GET /v1/models
      python3 scripts/live.py anthropic POST /v1/messages '{"model":"claude-haiku-4-5-20251001","max_tokens":8,"messages":[{"role":"user","content":"Reply with OK."}]}'
"""
from __future__ import annotations

import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator

ROOT = Path(__file__).resolve().parent.parent
LIVE_LOG = ROOT / "reports" / "live-requests.jsonl"
KEY_RE = re.compile(r"sk-(?:ant-)?[A-Za-z0-9_-]{8,}|xai-[A-Za-z0-9_-]{16,}|AIza[A-Za-z0-9_-]{20,}|AQ\.[A-Za-z0-9_-]{20,}")
XAI_BASE = os.environ.get("XAI_BASE_URL", "https://api.x.ai")
GEMINI_BASE = os.environ.get("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com")
OPENAI_BASE = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com")
ANTHROPIC_BASE = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
ANTHROPIC_VERSION = os.environ.get("ANTHROPIC_VERSION", "2023-06-01")


def load_env() -> None:
    env_file = ROOT / ".env"
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


load_env()


def mask(text: str) -> str:
    text = KEY_RE.sub("***REDACTED-KEY***", text)
    text = re.sub(r"([?&]key=)[A-Za-z0-9._-]{8,}", r"\1***REDACTED***", text)
    return re.sub(r"(?i)(bearer|x-api-key:?|x-goog-api-key:?)\s+[A-Za-z0-9._-]{8,}", r"\1 ***REDACTED***", text)


def now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def log_request(provider: str, method: str, path: str, status: int, est_cost_usd: float = 0.0, note: str = "") -> None:
    LIVE_LOG.parent.mkdir(parents=True, exist_ok=True)
    rec = {"ts": now(), "provider": provider, "method": method, "path": path, "status": status,
           "est_cost_usd": round(est_cost_usd, 6), "note": mask(note)[:300]}
    with LIVE_LOG.open("a") as f:
        f.write(json.dumps(rec) + "\n")


def _request(url: str, method: str, headers: dict[str, str], data: bytes | None, timeout: float,
             stream: bool) -> tuple[int, Any, dict[str, str]]:
    req = urllib.request.Request(url, method=method, headers=headers, data=data)
    try:
        resp = urllib.request.urlopen(req, timeout=timeout)
    except urllib.error.HTTPError as e:
        body = e.read()
        try:
            parsed = json.loads(body)
        except Exception:  # noqa: BLE001
            parsed = body
        return e.code, parsed, dict(e.headers.items())
    hdrs = dict(resp.headers.items())
    if stream:
        def gen() -> Iterator[str]:
            with resp:
                for raw in resp:
                    yield raw.decode("utf-8", "replace").rstrip("\n")
        return resp.status, gen(), hdrs
    with resp:
        body = resp.read()
    ctype = hdrs.get("Content-Type", "")
    if "json" in ctype:
        try:
            return resp.status, json.loads(body), hdrs
        except Exception:  # noqa: BLE001
            return resp.status, body, hdrs
    return resp.status, body, hdrs


def openai_request(method: str, path: str, json_body: Any = None, *, data: bytes | None = None,
                   content_type: str | None = "application/json", extra_headers: dict[str, str] | None = None,
                   timeout: float = 120, stream: bool = False, est_cost_usd: float = 0.0, note: str = "",
                   admin: bool = False) -> tuple[int, Any, dict[str, str]]:
    key = os.environ.get("OPENAI_ADMIN_KEY") if admin else None
    key = key or os.environ["OPENAI_API_KEY"]
    headers = {"Authorization": f"Bearer {key}"}
    if content_type:
        headers["Content-Type"] = content_type
    if extra_headers:
        headers.update(extra_headers)
    if json_body is not None:
        data = json.dumps(json_body).encode()
    st, body, hdrs = _request(OPENAI_BASE + path, method, headers, data, timeout, stream)
    log_request("openai", method, path.split("?")[0], st, est_cost_usd, note)
    return st, body, hdrs


def anthropic_request(method: str, path: str, json_body: Any = None, *, data: bytes | None = None,
                      content_type: str | None = "application/json", beta: str | None = None,
                      extra_headers: dict[str, str] | None = None, timeout: float = 120, stream: bool = False,
                      est_cost_usd: float = 0.0, note: str = "", admin: bool = False) -> tuple[int, Any, dict[str, str]]:
    key = os.environ.get("ANTHROPIC_ADMIN_KEY") if admin else None
    key = key or os.environ["ANTHROPIC_API_KEY"]
    headers = {"x-api-key": key, "anthropic-version": ANTHROPIC_VERSION}
    if content_type:
        headers["Content-Type"] = content_type
    if beta:
        headers["anthropic-beta"] = beta
    if extra_headers:
        headers.update(extra_headers)
    if json_body is not None:
        data = json.dumps(json_body).encode()
    st, body, hdrs = _request(ANTHROPIC_BASE + path, method, headers, data, timeout, stream)
    log_request("anthropic", method, path.split("?")[0], st, est_cost_usd, note)
    return st, body, hdrs


def xai_request(method: str, path: str, json_body: Any = None, *, data: bytes | None = None,
                content_type: str | None = "application/json", extra_headers: dict[str, str] | None = None,
                timeout: float = 120, stream: bool = False, est_cost_usd: float = 0.0, note: str = "",
                management: bool = False) -> tuple[int, Any, dict[str, str]]:
    """xAI REST API (OpenAI-compatible surface + xAI-specific endpoints). management=True targets https://management-api.x.ai
    with XAI_MANAGEMENT_KEY if set (Management API keys are distinct from inference keys)."""
    key = os.environ.get("XAI_MANAGEMENT_KEY") if management else None
    key = key or os.environ["XAI_API_KEY"]
    base = os.environ.get("XAI_MANAGEMENT_BASE_URL", "https://management-api.x.ai") if management else XAI_BASE
    headers = {"Authorization": f"Bearer {key}"}
    if content_type:
        headers["Content-Type"] = content_type
    if extra_headers:
        headers.update(extra_headers)
    if json_body is not None:
        data = json.dumps(json_body).encode()
    st, body, hdrs = _request(base + path, method, headers, data, timeout, stream)
    log_request("xai", method, path.split("?")[0], st, est_cost_usd, note)
    return st, body, hdrs


def gemini_request(method: str, path: str, json_body: Any = None, *, data: bytes | None = None,
                   content_type: str | None = "application/json", extra_headers: dict[str, str] | None = None,
                   timeout: float = 120, stream: bool = False, est_cost_usd: float = 0.0,
                   note: str = "") -> tuple[int, Any, dict[str, str]]:
    """Gemini API (generativelanguage.googleapis.com). `path` includes the version, e.g.
    '/v1beta/models/gemini-2.5-flash:generateContent'. The key is sent as the x-goog-api-key header (never in the URL)."""
    headers = {"x-goog-api-key": os.environ["GEMINI_API_KEY"]}
    if content_type:
        headers["Content-Type"] = content_type
    if extra_headers:
        headers.update(extra_headers)
    if json_body is not None:
        data = json.dumps(json_body).encode()
    st, body, hdrs = _request(GEMINI_BASE + path, method, headers, data, timeout, stream)
    log_request("gemini", method, re.sub(r"models/[^:/?]+", "models/{model}", path.split("?")[0]), st, est_cost_usd, note)
    return st, body, hdrs


def interesting_headers(hdrs: dict[str, str]) -> dict[str, str]:
    """Rate-limit / request-id headers worth recording (never auth)."""
    keep = ("x-ratelimit", "anthropic-ratelimit", "x-goog", "server-timing", "alt-svc-", "x-xai", "retry-after", "x-request-id", "request-id",
            "openai-processing-ms", "openai-version", "openai-organization", "openai-project", "cf-ray",
            "anthropic-organization-id", "x-should-retry")
    return {k: v for k, v in hdrs.items() if k.lower().startswith(keep)}


def save_sanitized(obj: Any, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    text = obj if isinstance(obj, str) else json.dumps(obj, indent=2, ensure_ascii=False, default=str)
    dest.write_text(mask(text))


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print(__doc__)
        sys.exit(1)
    provider, method, path = sys.argv[1], sys.argv[2].upper(), sys.argv[3]
    body = json.loads(sys.argv[4]) if len(sys.argv) > 4 else None
    beta = os.environ.get("ANTHROPIC_BETA")
    t0 = time.time()
    if provider == "openai":
        st, out, hdrs = openai_request(method, path, body)
    elif provider == "xai":
        st, out, hdrs = xai_request(method, path, body)
    elif provider == "gemini":
        st, out, hdrs = gemini_request(method, path, body)
    else:
        st, out, hdrs = anthropic_request(method, path, body, beta=beta)
    print(f"HTTP {st} ({time.time() - t0:.2f}s)", file=sys.stderr)
    print(mask(json.dumps(interesting_headers(hdrs), indent=1)), file=sys.stderr)
    print(mask(json.dumps(out, indent=2, ensure_ascii=False) if not isinstance(out, bytes) else f"<{len(out)} bytes>"))