#!/usr/bin/env python3
"""
token_budget.py — Daily token budget tracker for pipeline orchestration.

Tracks token usage per pipeline per day, enforces daily limits,
and provides a query interface for the orchestrator.

Data: ~/.hermes/pipeline_budget.json
"""

import json
import os
from datetime import datetime, timezone
from pathlib import Path

BUDGET_FILE = Path.home() / ".hermes" / "pipeline_budget.json"
DEFAULT_DAILY_LIMIT = 500_000


def _load() -> dict:
    if BUDGET_FILE.exists():
        try:
            return json.loads(BUDGET_FILE.read_text())
        except (json.JSONDecodeError, OSError):
            pass
    return {}


def _save(data: dict):
    BUDGET_FILE.parent.mkdir(parents=True, exist_ok=True)
    BUDGET_FILE.write_text(json.dumps(data, indent=2))


def today_key() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def get_daily_usage(pipeline: str = None) -> dict:
    """Get token usage for today. If pipeline specified, return just that pipeline."""
    data = _load()
    day = data.get("daily", {}).get(today_key(), {"tokens_used": 0, "pipelines": {}})
    if pipeline:
        return {
            "pipeline": pipeline,
            "tokens_used": day.get("pipelines", {}).get(pipeline, 0),
            "daily_total": day.get("tokens_used", 0),
        }
    return day


def get_remaining(limit: int = DEFAULT_DAILY_LIMIT) -> int:
    """Get remaining token budget for today."""
    usage = get_daily_usage()
    return max(0, limit - usage.get("tokens_used", 0))


def can_afford(tokens: int, limit: int = DEFAULT_DAILY_LIMIT) -> bool:
    """Check if we have budget for a token spend."""
    return get_remaining(limit) >= tokens


def record_usage(pipeline: str, input_tokens: int, output_tokens: int) -> dict:
    """
    Record token usage for a pipeline task.

    Called automatically by the orchestrator after each pipeline task completes.
    Returns the updated daily state.
    """
    total = input_tokens + output_tokens
    data = _load()
    today = today_key()

    daily = data.setdefault("daily", {})
    day = daily.setdefault(today, {"tokens_used": 0, "pipelines": {}})

    day["tokens_used"] = day.get("tokens_used", 0) + total
    pipes = day.setdefault("pipelines", {})
    pipes[pipeline] = pipes.get(pipeline, 0) + total

    # Track breakdown
    breakdown = day.setdefault("breakdown", {})
    pb = breakdown.setdefault(pipeline, {"input": 0, "output": 0, "calls": 0})
    pb["input"] += input_tokens
    pb["output"] += output_tokens
    pb["calls"] += 1

    # Track lifetime stats
    lifetime = data.setdefault("lifetime", {"total_tokens": 0, "total_days": 0})
    lifetime["total_tokens"] = lifetime.get("total_tokens", 0) + total

    _save(data)

    return {
        "pipeline": pipeline,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total": total,
        "daily_used": day["tokens_used"],
        "daily_remaining": get_remaining(),
    }


def get_report() -> str:
    """Generate a human-readable budget report."""
    data = _load()
    today = today_key()
    day = data.get("daily", {}).get(today, {"tokens_used": 0, "pipelines": {}})

    lines = []
    lines.append(f"Token Budget — {today}")
    lines.append(f"  Daily usage: {day.get('tokens_used', 0):,} / {DEFAULT_DAILY_LIMIT:,}")
    lines.append(f"  Remaining: {get_remaining():,}")
    lines.append("")
    lines.append("  Pipelines:")

    breakdown = day.get("breakdown", {})
    for name, stats in sorted(breakdown.items(), key=lambda x: -x[1]["output"]):
        total = stats["input"] + stats["output"]
        lines.append(f"    {name}: {total:,} tokens ({stats['calls']} calls)")

    if not breakdown:
        lines.append("    (no pipelines run today)")

    lifetime = data.get("lifetime", {})
    lines.append("")
    lines.append(f"  Lifetime: {lifetime.get('total_tokens', 0):,} total tokens")

    return "\n".join(lines)


if __name__ == "__main__":
    import sys
    if "--report" in sys.argv:
        print(get_report())
    elif "--remaining" in sys.argv:
        print(get_remaining())
    elif "--can-afford" in sys.argv:
        idx = sys.argv.index("--can-afford")
        tokens = int(sys.argv[idx + 1])
        print("yes" if can_afford(tokens) else "no")
    else:
        print(get_report())