import os
import time
import json
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["RADIUM_API_KEY"],
    base_url="https://api.radium.cloud/v1",
)

# Pricing (per 1M tokens) — populate from https://radium.cloud/pricing
PRICING = {
    "tycho-1.0": {"prompt": 0.00, "completion": 0.00},
    "clarke-1.0": {"prompt": 0.00, "completion": 0.00},
    "hal-1.0": {"prompt": 0.00, "completion": 0.00},
}


@dataclass
class RequestLog:
    timestamp: str
    model: str
    latency_ms: float
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    estimated_cost_usd: float
    status: str  # "success" | "error"
    error_type: str | None


def estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> float:
    rates = PRICING.get(model, PRICING["clarke-1.0"])
    return (
        (prompt_tokens / 1_000_000) * rates["prompt"]
        + (completion_tokens / 1_000_000) * rates["completion"]
    )


def logged_chat(messages, model="clarke-1.0", **kwargs) -> tuple[str, RequestLog]:
    """
    Wraps chat.completions.create with logging.
    Returns (content, log_entry).
    """
    start = time.perf_counter()
    status = "success"
    error_type = None
    usage = None

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            **kwargs,
        )
        content = response.choices[0].message.content
        usage = response.usage
    except Exception as e:
        status = "error"
        error_type = type(e).__name__
        content = ""

    latency = (time.perf_counter() - start) * 1000

    log = RequestLog(
        timestamp=datetime.now(timezone.utc).isoformat(),
        model=model,
        latency_ms=latency,
        prompt_tokens=usage.prompt_tokens if usage else 0,
        completion_tokens=usage.completion_tokens if usage else 0,
        total_tokens=usage.total_tokens if usage else 0,
        estimated_cost_usd=estimate_cost(
            model,
            usage.prompt_tokens if usage else 0,
            usage.completion_tokens if usage else 0,
        ) if usage else 0.0,
        status=status,
        error_type=error_type,
    )
    return content, log


# --- Run it ---
logs: list[RequestLog] = []

for prompt in [
    "Summarize the theory of relativity in one paragraph.",
    "List 3 benefits of exercise.",
    "Explain recursion with a Python example.",
]:
    content, log = logged_chat(
        messages=[{"role": "user", "content": prompt}],
        model="clarke-1.0",
        max_tokens=256,
    )
    logs.append(log)
    print(f"[{log.model}] {log.status} — {log.total_tokens} tokens, ${log.estimated_cost_usd:.6f}, {log.latency_ms:.0f}ms")

# Export to JSONL
with open("usage_logs.jsonl", "w") as f:
    for log in logs:
        f.write(json.dumps(asdict(log)) + "\n")

# Aggregate summary
if logs:
    total_cost = sum(l.estimated_cost_usd for l in logs)
    avg_latency = sum(l.latency_ms for l in logs) / len(logs)
    total_tokens = sum(l.total_tokens for l in logs)
    error_rate = sum(1 for l in logs if l.status == "error") / len(logs)

    print(f"\nSummary: {len(logs)} requests, {total_tokens} tokens, ${total_cost:.4f}, avg {avg_latency:.0f}ms, {error_rate:.0%} errors")

# Pseudocode — wire into your existing metrics client
from prometheus_client import Counter, Histogram

llm_requests = Counter("llm_requests_total", "Total LLM requests", ["model", "status"])
llm_tokens = Counter("llm_tokens_total", "Total tokens", ["model", "type"])
llm_latency = Histogram("llm_latency_seconds", "Request latency", ["model"])

for log in logs:
    llm_requests.labels(model=log.model, status=log.status).inc()
    llm_tokens.labels(model=log.model, type="prompt").inc(log.prompt_tokens)
    llm_tokens.labels(model=log.model, type="completion").inc(log.completion_tokens)
    llm_latency.labels(model=log.model).observe(log.latency_ms / 1000)