#!/usr/bin/env python3
"""
Cost guard — API spend tracking, budget enforcement, and auto-cutoff.

Ported from sovereign-org-engine/core/cost_guard.py with identical
semantics: same monthly budget caps, same 80% threshold, same state-file
schema, same month-rollover reset. (The org engine had monthly caps only —
no daily caps existed, so none were invented here.) State lives in
chairman-agent/data/api_costs.json; at cutover you can seed it by copying
sovereign-org-engine/data/api_costs.json so the current month's spend
carries over.

Two intentional differences from the org engine, both stricter:
  - PRICING gained entries for the models the Chairman actually runs
    (claude-sonnet-5, gpt-5.5, deepseek); unknown models still fall back
    to Sonnet pricing like the original.
  - CutoffError lets a caller hard-stop an in-flight tool loop; the org
    engine only checked the flag between conversations.
"""
import json
from datetime import datetime, timezone
from pathlib import Path

COST_FILE = Path(__file__).resolve().parent / "data" / "api_costs.json"

# Per-token pricing (USD per 1K tokens) — org-engine table + Chairman models
PRICING = {
    # Anthropic
    "claude-opus-4-7":          {"input": 0.015,  "output": 0.075},
    "claude-opus-4-6":          {"input": 0.015,  "output": 0.075},
    "claude-opus-4-5":          {"input": 0.015,  "output": 0.075},
    "claude-sonnet-5":          {"input": 0.003,  "output": 0.015},
    "claude-sonnet-4-7":        {"input": 0.003,  "output": 0.015},
    "claude-sonnet-4-6":        {"input": 0.003,  "output": 0.015},
    "claude-sonnet-4-5":        {"input": 0.003,  "output": 0.015},
    "claude-haiku-4-5":         {"input": 0.0008, "output": 0.004},
    "claude-3-5-haiku-latest":  {"input": 0.0008, "output": 0.004},
    # Other providers
    "gpt-4o":                   {"input": 0.0025, "output": 0.01},
    "openrouter/openai/gpt-5.5": {"input": 0.00125, "output": 0.01},
    "sonar-pro":                {"input": 0.003,  "output": 0.015},
    "z-ai/glm-5.2":             {"input": 0.00095, "output": 0.003},
    "openrouter/z-ai/glm-5.2":  {"input": 0.00095, "output": 0.003},
    "deepseek-v4-flash":        {"input": 0.00027, "output": 0.0011},
    "deepseek-v4-pro":          {"input": 0.00055, "output": 0.0022},
    # Local models = free
    "llama3.2":    {"input": 0.0, "output": 0.0},
    "qwen2.5:3b":  {"input": 0.0, "output": 0.0},
    "qwen2.5:7b":  {"input": 0.0, "output": 0.0},
    "qwen2.5:14b": {"input": 0.0, "output": 0.0},
}

# Monthly budget caps (USD) — identical to the org engine
BUDGET = {
    "anthropic": 50.0,
    "openai": 50.0,
    "openrouter": 30.0,
    "perplexity": 20.0,
    "total": 150.0,
}

THRESHOLD = 0.80  # 80% auto-cutoff, same as org engine

PAID_PROVIDERS = ["anthropic", "openai", "openrouter", "perplexity"]


class CutoffError(RuntimeError):
    """Raised by metered callers when the budget cutoff is active."""


def _new_month():
    return {
        "month": datetime.now().strftime("%Y-%m"),
        "anthropic":  {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0},
        "openai":     {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0},
        "perplexity": {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0},
        "ollama":     {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0},
        "cutoff_active": False,
        "cutoff_triggered_at": None,
    }


def _load_costs():
    try:
        if COST_FILE.exists():
            data = json.loads(COST_FILE.read_text())
            if data.get("month") != datetime.now().strftime("%Y-%m"):
                return _new_month()
            return data
    except Exception:
        pass
    return _new_month()


def _save_costs(data):
    COST_FILE.parent.mkdir(parents=True, exist_ok=True)
    COST_FILE.write_text(json.dumps(data, indent=2))


def log_api_call(provider, model, input_tokens, output_tokens):
    """Log an API call, accrue its cost, trip the cutoff at the threshold."""
    data = _load_costs()
    pricing = PRICING.get(model, {"input": 0.003, "output": 0.015})
    cost = (input_tokens / 1000 * pricing["input"]) \
        + (output_tokens / 1000 * pricing["output"])

    if provider not in data:
        data[provider] = {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0}
    data[provider]["calls"] += 1
    data[provider]["input_tokens"] += input_tokens
    data[provider]["output_tokens"] += output_tokens
    data[provider]["cost"] += cost

    total_paid = sum(data[p]["cost"] for p in PAID_PROVIDERS if p in data)
    if total_paid >= BUDGET["total"] * THRESHOLD and not data["cutoff_active"]:
        data["cutoff_active"] = True
        data["cutoff_triggered_at"] = datetime.now(timezone.utc).isoformat()

    _save_costs(data)
    return cost


def is_cutoff_active():
    return _load_costs().get("cutoff_active", False)


def get_cost_report():
    """Full cost report for the current month — org-engine shape."""
    data = _load_costs()
    total_paid = sum(data[p]["cost"] for p in PAID_PROVIDERS if p in data)
    report = {"month": data["month"]}
    for p in PAID_PROVIDERS:
        row = data.get(p, {})
        report[p] = {
            "cost": round(row.get("cost", 0), 4),
            "calls": row.get("calls", 0),
            "budget": BUDGET[p],
            "usage_pct": round(row.get("cost", 0) / BUDGET[p] * 100, 1),
        }
    report["ollama"] = {"cost": 0.0, "calls": data.get("ollama", {}).get("calls", 0)}
    report["total_paid"] = round(total_paid, 4)
    report["total_budget"] = BUDGET["total"]
    report["total_usage_pct"] = round(total_paid / BUDGET["total"] * 100, 1)
    report["cutoff_active"] = data.get("cutoff_active", False)
    return report


def reset_cutoff():
    """Manual human override — /unlock."""
    data = _load_costs()
    data["cutoff_active"] = False
    data["cutoff_triggered_at"] = None
    _save_costs(data)