"""
Usage / token-tracking endpoints.

Self-service (any authenticated user, own data):
  GET /api/usage              → full usage dict  {date: {model_key: {calls, prompt_tokens, completion_tokens}}}
  GET /api/usage/summary      → aggregate totals per model key, with friendly labels resolved from registry

Admin-only (cross-user aggregation):
  GET /api/usage/all          → summary for every user  {username: summary_dict}
"""
import jwt
from fastapi import APIRouter, HTTPException, Request

from auth_utils import COOKIE_NAME, decode_token, get_user_role
from persona import list_users
import model_registry
import usage_tracker

router = APIRouter(prefix="/api/usage")


def _session_user(request: Request) -> str:
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        raise HTTPException(status_code=401, detail="Not authenticated")
    try:
        return decode_token(token)
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid session")


def _build_label_map(username: str) -> dict[str, str]:
    """Build a map from usage key (backend/model_name) → registered label."""
    label_map: dict[str, str] = {}
    try:
        for m in model_registry.get_all_models(username):
            model_name = m.get("model_name", "")
            label = m.get("label", "")
            host_type = m.get("host_type", "")
            if not model_name or not label:
                continue
            # local models: key is "local/{model_name}"
            if host_type in ("openwebui", "ollama", "openai_compatible"):
                label_map[f"local/{model_name}"] = label
            # cloud Gemini: key is "gemini_api/{model_name}"
            elif host_type == "google":
                label_map[f"gemini_api/{model_name}"] = label
    except Exception:
        pass
    return label_map


def _summarize(data: dict, label_map: dict[str, str] | None = None) -> list[dict]:
    """Collapse date-keyed usage dict into per-model totals, sorted by total tokens desc."""
    totals: dict[str, dict] = {}
    for _date, models in data.items():
        for key, counts in models.items():
            t = totals.setdefault(key, {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0})
            t["calls"]             += counts.get("calls", 0)
            t["prompt_tokens"]     += counts.get("prompt_tokens", 0)
            t["completion_tokens"] += counts.get("completion_tokens", 0)

    result = []
    for key, counts in totals.items():
        entry = {
            "key":              key,
            "label":            (label_map or {}).get(key) or key,
            "calls":            counts["calls"],
            "prompt_tokens":    counts["prompt_tokens"],
            "completion_tokens": counts["completion_tokens"],
            "total_tokens":     counts["prompt_tokens"] + counts["completion_tokens"],
        }
        result.append(entry)

    result.sort(key=lambda x: x["total_tokens"], reverse=True)
    return result


@router.get("")
async def get_usage(request: Request) -> dict:
    """Return the raw daily usage log for the authenticated user."""
    username = _session_user(request)
    return usage_tracker.read_usage(username)


@router.get("/summary")
async def get_usage_summary(request: Request) -> list:
    """Return per-model totals (all time) for the authenticated user, with friendly labels."""
    username = _session_user(request)
    label_map = _build_label_map(username)
    return _summarize(usage_tracker.read_usage(username), label_map)


@router.get("/all")
async def get_all_usage(request: Request) -> dict:
    """Admin: return per-model summary for every user."""
    username = _session_user(request)
    if get_user_role(username) != "admin":
        raise HTTPException(status_code=403, detail="Admin access required")
    result = {}
    for user in list_users():
        label_map = _build_label_map(user)
        result[user] = _summarize(usage_tracker.read_usage(user), label_map)
    return result