#!/usr/bin/env python3
"""
Genesis Weekly Usage Tracker — PostToolUse Hook
================================================
Tracks Claude Max plan usage against weekly budget to prevent
unexpected stalls from hitting the weekly limit.

Architecture:
  - Reads cost_usd from data/context_state/current.json (written by statusline)
  - Appends to data/usage/weekly_ledger.jsonl (one entry per session snapshot)
  - Computes rolling 7-day spend from ledger
  - Injects warning into additionalContext when approaching weekly limit

Budget assumptions (configurable below):
  - 20x Max plan: ~$140/week equivalent
  - 5x Max plan:  ~$35/week equivalent
  - Daily safe pace = weekly_budget / 7

Thresholds:
  60% of weekly budget → WARN
  80% of weekly budget → ALERT (slow down)
  95% of weekly budget → CRITICAL (stop heavy work)

Author: Genesis Command Centre
"""

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

# ============================================================================
# CONFIGURATION — adjust to match active plan
# ============================================================================

GENESIS_ROOT = Path("/mnt/e/genesis-system")
STATE_DIR = GENESIS_ROOT / "data" / "context_state"
USAGE_DIR = GENESIS_ROOT / "data" / "usage"
LEDGER_FILE = USAGE_DIR / "weekly_ledger.jsonl"
BUDGET_FILE = USAGE_DIR / "budget_config.json"

# Default budget config (overridden by budget_config.json)
DEFAULT_CONFIG = {
    "active_plan": "20x_max",
    "weekly_budget_usd": 140.0,   # 20x Max plan equivalent
    "warn_pct": 60,
    "alert_pct": 80,
    "critical_pct": 95,
    "check_interval_calls": 10,   # Check every N tool calls
}

# How often to run (every N PostToolUse calls)
COUNTER_FILE = STATE_DIR / "usage_tracker_counter.json"


# ============================================================================
# CONFIG & COUNTER
# ============================================================================

def load_config() -> dict:
    try:
        if BUDGET_FILE.exists():
            with open(BUDGET_FILE) as f:
                return {**DEFAULT_CONFIG, **json.load(f)}
    except Exception:
        pass
    return DEFAULT_CONFIG


def get_counter() -> int:
    try:
        if COUNTER_FILE.exists():
            with open(COUNTER_FILE) as f:
                return json.load(f).get("count", 0)
    except Exception:
        pass
    return 0


def set_counter(count: int):
    try:
        STATE_DIR.mkdir(parents=True, exist_ok=True)
        with open(COUNTER_FILE, "w") as f:
            json.dump({"count": count}, f)
    except Exception:
        pass


# ============================================================================
# COST READING
# ============================================================================

def get_current_cost() -> dict:
    """Read current session cost from statusline state file."""
    try:
        state_file = STATE_DIR / "current.json"
        if state_file.exists():
            with open(state_file) as f:
                state = json.load(f)
                return {
                    "cost_usd": state.get("cost_usd", 0.0),
                    "session_id": state.get("session_id", "unknown"),
                    "timestamp": state.get("timestamp", datetime.now(timezone.utc).isoformat()),
                    "context_pct": state.get("used_percentage", 0),
                }
    except Exception:
        pass
    return {"cost_usd": 0.0, "session_id": "unknown", "timestamp": "", "context_pct": 0}


# ============================================================================
# LEDGER
# ============================================================================

def append_to_ledger(cost_data: dict):
    """Append current cost snapshot to weekly ledger."""
    try:
        USAGE_DIR.mkdir(parents=True, exist_ok=True)
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "session_id": cost_data["session_id"],
            "cost_usd": cost_data["cost_usd"],
            "context_pct": cost_data["context_pct"],
            "interface": os.environ.get("CLAUDE_INTERFACE", "claude_code"),
        }
        with open(LEDGER_FILE, "a") as f:
            f.write(json.dumps(entry) + "\n")
    except Exception:
        pass


def get_weekly_spend() -> float:
    """
    Calculate total spend over the last 7 days from the ledger.
    Uses max cost_usd per session_id (session cost is cumulative,
    so we take the peak reading per session to avoid double-counting).
    """
    if not LEDGER_FILE.exists():
        return 0.0

    cutoff = datetime.now(timezone.utc) - timedelta(days=7)
    session_peaks = {}  # session_id -> max cost seen

    try:
        with open(LEDGER_FILE) as f:
            for line in f:
                try:
                    entry = json.loads(line.strip())
                    ts = datetime.fromisoformat(entry["timestamp"].replace("Z", "+00:00"))
                    if ts < cutoff:
                        continue
                    sid = entry.get("session_id", "unknown")
                    cost = entry.get("cost_usd", 0.0)
                    if sid not in session_peaks or cost > session_peaks[sid]:
                        session_peaks[sid] = cost
                except Exception:
                    continue
    except Exception:
        return 0.0

    return sum(session_peaks.values())


def get_today_spend() -> float:
    """Calculate today's spend (since midnight UTC)."""
    if not LEDGER_FILE.exists():
        return 0.0

    today_start = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
    session_peaks = {}

    try:
        with open(LEDGER_FILE) as f:
            for line in f:
                try:
                    entry = json.loads(line.strip())
                    ts = datetime.fromisoformat(entry["timestamp"].replace("Z", "+00:00"))
                    if ts < today_start:
                        continue
                    sid = entry.get("session_id", "unknown")
                    cost = entry.get("cost_usd", 0.0)
                    if sid not in session_peaks or cost > session_peaks[sid]:
                        session_peaks[sid] = cost
                except Exception:
                    continue
    except Exception:
        return 0.0

    return sum(session_peaks.values())


# ============================================================================
# USAGE REPORT
# ============================================================================

def build_usage_report(config: dict) -> dict:
    """Build usage report with spend vs budget."""
    weekly_budget = config["weekly_budget_usd"]
    weekly_spend = get_weekly_spend()
    today_spend = get_today_spend()
    daily_safe = weekly_budget / 7

    weekly_pct = (weekly_spend / weekly_budget * 100) if weekly_budget > 0 else 0
    remaining_weekly = max(0, weekly_budget - weekly_spend)

    # Days until reset (weekly reset = 7 days from oldest entry in window)
    # Approximate: show remaining budget as days at current daily pace
    daily_avg = weekly_spend / 7 if weekly_spend > 0 else 0
    days_remaining_at_pace = (remaining_weekly / daily_avg) if daily_avg > 0 else 7

    return {
        "weekly_budget": weekly_budget,
        "weekly_spend": round(weekly_spend, 2),
        "weekly_pct": round(weekly_pct, 1),
        "remaining_weekly": round(remaining_weekly, 2),
        "today_spend": round(today_spend, 2),
        "daily_safe_limit": round(daily_safe, 2),
        "days_remaining_at_pace": round(days_remaining_at_pace, 1),
        "plan": config["active_plan"],
        "status": (
            "CRITICAL" if weekly_pct >= config["critical_pct"] else
            "ALERT" if weekly_pct >= config["alert_pct"] else
            "WARN" if weekly_pct >= config["warn_pct"] else
            "OK"
        )
    }


def write_usage_summary(report: dict):
    """Write human-readable usage summary for statusline / desktop use."""
    try:
        USAGE_DIR.mkdir(parents=True, exist_ok=True)
        summary_file = USAGE_DIR / "weekly_summary.json"
        with open(summary_file, "w") as f:
            json.dump({**report, "updated_at": datetime.now(timezone.utc).isoformat()}, f, indent=2)
    except Exception:
        pass


# ============================================================================
# MAIN
# ============================================================================

def main():
    try:
        hook_input = json.loads(sys.stdin.read())
    except Exception:
        print(json.dumps({}))
        return

    config = load_config()

    # Throttle: only run every N tool calls
    count = get_counter() + 1
    if count < config["check_interval_calls"]:
        set_counter(count)
        print(json.dumps({}))
        return
    set_counter(0)

    # Get current cost and log it
    cost_data = get_current_cost()
    if cost_data["cost_usd"] > 0:
        append_to_ledger(cost_data)

    # Build report
    report = build_usage_report(config)
    write_usage_summary(report)

    # Inject context warning if needed
    result = {}
    status = report["status"]

    if status == "CRITICAL":
        result["additionalContext"] = (
            f"🚨 WEEKLY BUDGET CRITICAL: ${report['weekly_spend']:.2f} / ${report['weekly_budget']:.2f} "
            f"({report['weekly_pct']:.0f}% used). Only ${report['remaining_weekly']:.2f} remaining. "
            f"STOP heavy Claude Max work immediately. Use OpenRouter/MiniMax/Kimi swarm only. "
            f"Today: ${report['today_spend']:.2f} (safe limit: ${report['daily_safe_limit']:.2f}/day)."
        )
    elif status == "ALERT":
        result["additionalContext"] = (
            f"⚠️ WEEKLY BUDGET ALERT: ${report['weekly_spend']:.2f} / ${report['weekly_budget']:.2f} "
            f"({report['weekly_pct']:.0f}% used). ${report['remaining_weekly']:.2f} left. "
            f"Shift heavy work to OpenRouter swarm. Daily safe limit: ${report['daily_safe_limit']:.2f}. "
            f"Today spent: ${report['today_spend']:.2f}."
        )
    elif status == "WARN":
        result["additionalContext"] = (
            f"💛 WEEKLY BUDGET WARN: ${report['weekly_spend']:.2f} / ${report['weekly_budget']:.2f} "
            f"({report['weekly_pct']:.0f}% used). Pace is ${report['today_spend']:.2f} today "
            f"vs ${report['daily_safe_limit']:.2f} safe daily limit."
        )

    print(json.dumps(result))


if __name__ == "__main__":
    main()