"""ARIS Budget — AI cost control and tracking."""

import json
import os
import time
from pathlib import Path
from typing import Any, Optional

MODEL_COSTS = {
    "claude-opus": {"input": 15.0, "output": 75.0},
    "claude-sonnet": {"input": 3.0, "output": 15.0},
    "claude-haiku": {"input": 0.25, "output": 1.25},
    "gpt-4o": {"input": 2.5, "output": 10.0},
    "gpt-4o-mini": {"input": 0.15, "output": 0.6},
    "gemini-pro": {"input": 0.0, "output": 0.0},
    "gemini-flash": {"input": 0.0, "output": 0.0},
    "grok": {"input": 5.0, "output": 15.0},
    "ollama": {"input": 0.0, "output": 0.0},
}


class BudgetTracker:
    """Track and control AI spending."""

    def __init__(self, base_dir: Optional[str] = None) -> None:
        self._base = Path(base_dir or os.path.expanduser("~/.aris/budget"))
        self._base.mkdir(parents=True, exist_ok=True)
        self._config_file = self._base / "config.json"
        self._usage_file = self._base / "usage.json"
        self._config: dict[str, Any] = self._load_config()
        self._usage: list[dict[str, Any]] = self._load_usage()

    def _load_config(self) -> dict[str, Any]:
        if self._config_file.exists():
            try:
                return json.loads(self._config_file.read_text())
            except json.JSONDecodeError:
                pass
        return {
            "monthly_limit_usd": 0, "daily_limit_usd": 0,
            "alert_threshold_pct": 80, "block_on_exceed": False,
            "subscriptions": {},
        }

    def _save_config(self) -> None:
        self._config_file.write_text(json.dumps(self._config, indent=2))

    def _load_usage(self) -> list[dict[str, Any]]:
        if self._usage_file.exists():
            try:
                return json.loads(self._usage_file.read_text())
            except json.JSONDecodeError:
                pass
        return []

    def _save_usage(self) -> None:
        self._usage_file.write_text(json.dumps(self._usage))

    def set_limit(self, monthly: float = 0, daily: float = 0,
                  block_on_exceed: bool = False) -> dict[str, Any]:
        self._config["monthly_limit_usd"] = monthly
        self._config["daily_limit_usd"] = daily
        self._config["block_on_exceed"] = block_on_exceed
        self._save_config()
        return self._config

    def add_subscription(self, name: str, cost: float, features: list[str]) -> None:
        self._config["subscriptions"][name] = {"monthly_cost": cost, "features": features}
        self._save_config()

    def record_usage(self, model: str, tokens_in: int, tokens_out: int,
                     task: str = "") -> dict[str, Any]:
        costs = MODEL_COSTS.get(model, {"input": 5.0, "output": 15.0})
        cost = (tokens_in / 1_000_000 * costs["input"] +
                tokens_out / 1_000_000 * costs["output"])
        record = {"model": model, "tokens_in": tokens_in, "tokens_out": tokens_out,
                  "cost_usd": round(cost, 6), "task": task, "timestamp": time.time()}
        self._usage.append(record)
        self._save_usage()
        return {"cost": record["cost_usd"], "within_budget": self._check_budget()}

    def _check_budget(self) -> bool:
        ml = self._config["monthly_limit_usd"]
        if ml > 0 and self.get_spend("month") >= ml:
            return False
        dl = self._config["daily_limit_usd"]
        if dl > 0 and self.get_spend("day") >= dl:
            return False
        return True

    def get_spend(self, period: str = "month") -> float:
        cutoffs = {"day": 86400, "week": 604800, "month": 2592000}
        cutoff = time.time() - cutoffs.get(period, 2592000)
        return round(sum(r["cost_usd"] for r in self._usage if r["timestamp"] > cutoff), 4)

    def get_report(self) -> dict[str, Any]:
        by_model: dict[str, float] = {}
        for r in self._usage:
            by_model[r["model"]] = by_model.get(r["model"], 0) + r["cost_usd"]
        sub_total = sum(s["monthly_cost"] for s in self._config["subscriptions"].values())
        monthly = self.get_spend("month")
        return {
            "api_spend": {"today": self.get_spend("day"), "this_week": self.get_spend("week"),
                          "this_month": monthly},
            "subscriptions": {"monthly_total": sub_total, "details": self._config["subscriptions"]},
            "total_monthly": round(monthly + sub_total, 2),
            "by_model": {k: round(v, 4) for k, v in sorted(by_model.items(), key=lambda x: -x[1])},
            "limits": {"monthly": self._config["monthly_limit_usd"],
                       "daily": self._config["daily_limit_usd"],
                       "block_on_exceed": self._config["block_on_exceed"]},
            "within_budget": self._check_budget(),
        }
