"""Usage / utilization endpoints and helpers (Claude + Codex rate limits)."""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import os
import subprocess
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from aiohttp import web

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Caching globals
# ---------------------------------------------------------------------------

_oauth_cache: dict | None = None
_oauth_cached_at: float = 0.0
_OAUTH_CACHE_TTL = 300.0  # 5 minutes — Anthropic rate limits the OAuth usage endpoint
_openai_usage_cache: dict | None = None
_openai_usage_cached_at: float = 0.0
_OPENAI_USAGE_CACHE_TTL = 60.0
_dotenv_cache: dict[str, str] | None = None
_dotenv_cached_at: float = 0.0
_DOTENV_CACHE_TTL = 60.0

USAGE_URL = "https://api.anthropic.com/api/oauth/usage"
KEYCHAIN_SERVICE = "Claude Code-credentials"

# Utilization history — stored in SQLite via the store
# _pending_util_snapshot is used to queue a snapshot for async storage
_pending_util_snapshot: dict | None = None


def _record_util_snapshot(windows: list[dict]) -> None:
    """Queue a utilization snapshot for storage (called from sync context)."""
    global _pending_util_snapshot
    five_hour = next((w for w in windows if w["name"] == "five_hour"), None)
    seven_day = next((w for w in windows if w["name"] == "seven_day"), None)
    if not five_hour and not seven_day:
        return

    from datetime import datetime
    _pending_util_snapshot = {
        "timestamp": datetime.now(UTC).isoformat(),
        "five_hour": round(five_hour["utilization"] * 100, 1) if five_hour else None,
        "seven_day": round(seven_day["utilization"] * 100, 1) if seven_day else None,
    }


OPENAI_USAGE_URL = "https://api.openai.com/v1/organization/usage/completions"
OPENAI_COSTS_URL = "https://api.openai.com/v1/organization/costs"
OPENAI_USAGE_WINDOW_DAYS = 7


def _read_keychain_token() -> str | None:
    """Read Claude Code OAuth access token from macOS Keychain or Linux credentials file."""
    # Try Linux credentials file first (~/.claude/.credentials.json)
    from pathlib import Path
    creds_path = Path.home() / ".claude" / ".credentials.json"
    if creds_path.exists():
        try:
            data = json.loads(creds_path.read_text())
            token = data.get("claudeAiOauth", {}).get("accessToken")
            if token:
                return token
        except (json.JSONDecodeError, KeyError, OSError):
            pass

    # Fall back to macOS Keychain
    try:
        result = subprocess.run(
            ["security", "find-generic-password", "-s", KEYCHAIN_SERVICE, "-w"],
            capture_output=True, text=True, timeout=10,
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None
    if result.returncode != 0:
        return None
    try:
        data = json.loads(result.stdout)
        return data["claudeAiOauth"]["accessToken"]
    except (json.JSONDecodeError, KeyError):
        return None


def _load_dotenv_values() -> dict[str, str]:
    """Load simple KEY=VALUE pairs from the local .env file."""
    global _dotenv_cache, _dotenv_cached_at

    now = time.monotonic()
    if _dotenv_cache is not None and (now - _dotenv_cached_at) < _DOTENV_CACHE_TTL:
        return _dotenv_cache

    values: dict[str, str] = {}
    for filename in (".env",):
        path = os.path.join(os.getcwd(), filename)
        if not os.path.exists(path):
            continue
        try:
            with open(path, encoding="utf-8") as handle:
                for raw_line in handle:
                    line = raw_line.strip()
                    if not line or line.startswith("#"):
                        continue
                    if line.startswith("export "):
                        line = line[7:].strip()
                    if "=" not in line:
                        continue
                    key, value = line.split("=", 1)
                    key = key.strip()
                    value = value.strip()
                    if not key:
                        continue
                    if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
                        value = value[1:-1]
                    elif " #" in value:
                        value = value.split(" #", 1)[0].rstrip()
                    values[key] = value
        except OSError:
            continue

    _dotenv_cache = values
    _dotenv_cached_at = now
    return values


def _config_value(name: str) -> str | None:
    value = os.environ.get(name)
    if value:
        return value
    return _load_dotenv_values().get(name)


def _config_csv(name: str) -> list[str]:
    raw = _config_value(name)
    if not raw:
        return []
    return [item.strip() for item in raw.split(",") if item.strip()]


def _codex_sessions_root() -> Path:
    return Path.home() / ".codex" / "sessions"


def _unix_timestamp_to_iso(ts: Any) -> str | None:
    if ts is None:
        return None
    try:
        value = float(ts)
    except (TypeError, ValueError):
        return None
    return datetime.fromtimestamp(value, tz=UTC).isoformat().replace("+00:00", "Z")


def _remaining_percent(window: dict[str, Any] | None) -> float | None:
    if not window:
        return None
    try:
        used = float(window.get("used_percent"))
    except (TypeError, ValueError):
        return None
    return max(0.0, min(100.0, 100.0 - used))


def _serialize_codex_limit_window(window: dict[str, Any] | None) -> dict[str, Any]:
    if not window:
        return {
            "available": False,
            "used_percent": None,
            "remaining_percent": None,
            "window_minutes": None,
            "resets_at": None,
        }

    return {
        "available": True,
        "used_percent": float(window.get("used_percent") or 0.0),
        "remaining_percent": _remaining_percent(window),
        "window_minutes": window.get("window_minutes"),
        "resets_at": _unix_timestamp_to_iso(window.get("resets_at")),
    }


def _latest_codex_rate_limits_event(max_files: int = 24) -> dict[str, Any] | None:
    root = _codex_sessions_root()
    if not root.exists():
        return None

    try:
        files = sorted(
            (path for path in root.rglob("*.jsonl") if path.is_file()),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
    except OSError:
        return None

    latest_event: dict[str, Any] | None = None
    latest_key: tuple[str, str] | None = None

    for path in files[:max_files]:
        try:
            with path.open(encoding="utf-8", errors="ignore") as handle:
                for raw_line in handle:
                    if '"rate_limits"' not in raw_line or '"token_count"' not in raw_line:
                        continue
                    try:
                        event = json.loads(raw_line)
                    except json.JSONDecodeError:
                        continue
                    payload = event.get("payload") or {}
                    if payload.get("type") != "token_count":
                        continue
                    rate_limits = payload.get("rate_limits") or {}
                    if not rate_limits:
                        continue
                    timestamp = str(event.get("timestamp") or "")
                    key = (timestamp, str(path))
                    if latest_key is None or key > latest_key:
                        latest_key = key
                        latest_event = {
                            "timestamp": event.get("timestamp"),
                            "path": str(path),
                            "rate_limits": rate_limits,
                        }
        except OSError:
            continue

    return latest_event


def _read_local_codex_plan() -> dict[str, Any]:
    event = _latest_codex_rate_limits_event()
    if not event:
        return {
            "available": False,
            "source": "codex_local_sessions",
            "updated_at": None,
            "plan_type": None,
            "weekly_limit": _serialize_codex_limit_window(None),
            "code_review": _serialize_codex_limit_window(None),
            "credits": {
                "available": False,
                "has_credits": False,
                "unlimited": False,
                "balance": None,
                "message": "No local Codex session telemetry was found yet.",
            },
            "error": "No local Codex session telemetry found.",
        }

    rate_limits = event["rate_limits"]
    credits = rate_limits.get("credits") or {}
    balance = credits.get("balance")
    credits_available = balance is not None

    return {
        "available": bool(rate_limits.get("primary")),
        "source": "codex_local_sessions",
        "updated_at": event.get("timestamp"),
        "plan_type": rate_limits.get("plan_type"),
        "weekly_limit": _serialize_codex_limit_window(rate_limits.get("primary")),
        "code_review": _serialize_codex_limit_window(rate_limits.get("secondary")),
        "credits": {
            "available": credits_available,
            "has_credits": bool(credits.get("has_credits")),
            "unlimited": bool(credits.get("unlimited")),
            "balance": balance,
            "message": (
                None
                if credits_available
                else "Codex reports that credits exist, but the local session telemetry"
                " does not expose a balance."
                if credits.get("has_credits")
                else "No credits signal was present in local Codex session telemetry."
            ),
        },
        "error": None,
    }


def _claude_error_payload(message: str) -> dict[str, Any]:
    return {
        "provider": "claude",
        "source": "anthropic_oauth",
        "available": False,
        "plan": None,
        "windows": [],
        "error": message,
    }


async def _fetch_claude_usage_payload() -> tuple[dict[str, Any], int]:
    """Fetch Claude usage data for both /utilization and combined usage views."""
    global _oauth_cache, _oauth_cached_at

    now = time.monotonic()
    if _oauth_cache and (now - _oauth_cached_at) < _OAUTH_CACHE_TTL:
        return _oauth_cache, 200

    token = await asyncio.get_event_loop().run_in_executor(None, _read_keychain_token)
    if not token:
        return _claude_error_payload("Could not read OAuth token from Keychain"), 503

    try:
        import aiohttp as aiohttp_client
        async with aiohttp_client.ClientSession() as session, session.get(
            USAGE_URL,
            headers={
                "Authorization": f"Bearer {token}",
                "anthropic-beta": "oauth-2025-04-20",
            },
            timeout=aiohttp_client.ClientTimeout(total=10),
        ) as resp:
            if resp.status == 429:
                # Cache the rate limit error so we don't keep hammering
                error_result = _claude_error_payload("Rate limited")
                _oauth_cache = error_result
                _oauth_cached_at = now
                return error_result, 429
            if resp.status != 200:
                return _claude_error_payload(f"API returned {resp.status}"), 502
            data = await resp.json()
    except Exception as exc:
        logger.warning("Failed to fetch OAuth usage: %s", exc)
        return _claude_error_payload(str(exc)), 502

    windows = []
    for key in ("five_hour", "seven_day", "seven_day_sonnet", "seven_day_opus", "seven_day_cowork"):
        entry = data.get(key)
        if entry is None:
            continue
        windows.append({
            "name": key,
            "utilization": entry["utilization"] / 100.0,
            "resets_at": entry["resets_at"],
        })

    extra = data.get("extra_usage")
    if extra and extra.get("is_enabled"):
        windows.append({
            "name": "overage",
            "utilization": extra["utilization"] / 100.0,
            "resets_at": None,
        })

    result = {
        "provider": "claude",
        "source": "anthropic_oauth",
        "available": bool(windows),
        "windows": windows,
        "plan": data.get("plan", "unknown"),
        "error": None,
    }
    _oauth_cache = result
    _oauth_cached_at = now
    # Record utilization snapshot for history tracking
    _record_util_snapshot(windows)
    return result, 200


def _empty_codex_usage_payload() -> dict[str, Any]:
    admin_key = _config_value("OPENAI_ADMIN_KEY")
    project_key = _config_value("OPENAI_API_KEY")
    project_ids = _config_csv("OPENAI_USAGE_PROJECT_IDS")
    return {
        "provider": "codex",
        "source": "codex_local_sessions",
        "available": False,
        "window_days": OPENAI_USAGE_WINDOW_DAYS,
        "configuration": {
            "admin_key_configured": bool(admin_key),
            "project_key_configured": bool(project_key),
            "org_id_configured": bool(_config_value("OPENAI_ORG_ID")),
            "project_ids": project_ids,
        },
        "local_plan": {
            "available": False,
            "source": "codex_local_sessions",
            "updated_at": None,
            "plan_type": None,
            "weekly_limit": _serialize_codex_limit_window(None),
            "code_review": _serialize_codex_limit_window(None),
            "credits": {
                "available": False,
                "has_credits": False,
                "unlimited": False,
                "balance": None,
                "message": "No local Codex session telemetry found.",
            },
            "error": "No local Codex session telemetry found.",
        },
        "api_usage": {
            "available": False,
            "requests": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "cost_usd": 0.0,
            "models": [],
            "cost_breakdown": [],
            "scope": "organization" if not project_ids else "project",
            "error": None,
        },
        "plan_limits": {
            "supported": False,
            "items": [
                {"label": "Weekly usage limit", "value": "Read from local Codex session telemetry"},
                {
                    "label": "Code review",
                    "value": "Shown when local Codex telemetry includes a secondary limit",
                },
                {
                    "label": "Credits remaining",
                    "value": "Shown when local Codex telemetry includes a credit balance",
                },
            ],
        },
        "error": None,
        "note": (
            "Weekly Codex plan usage is read from local Codex session telemetry. "
            "OpenAI organization usage and costs remain optional secondary telemetry."
        ),
        "requirements": {
            "required": ["OPENAI_ADMIN_KEY"],
            "optional": ["OPENAI_ORG_ID", "OPENAI_USAGE_PROJECT_IDS"],
        },
    }


def _is_codex_model(model: str | None) -> bool:
    return bool(model) and "codex" in model.lower()


def _is_codex_line_item(line_item: str | None) -> bool:
    return bool(line_item) and "codex" in line_item.lower()


def _aggregate_openai_usage(
    buckets: list[dict[str, Any]],
) -> tuple[dict[str, int], list[dict[str, Any]]]:
    totals = {
        "requests": 0,
        "input_tokens": 0,
        "output_tokens": 0,
    }
    by_model: dict[str, dict[str, Any]] = {}

    for bucket in buckets:
        for result in bucket.get("results", []):
            model = result.get("model")
            if not _is_codex_model(model):
                continue

            requests = int(result.get("num_model_requests") or 0)
            input_tokens = (
                int(result.get("input_tokens") or 0)
                + int(result.get("input_cached_tokens") or 0)
            )
            output_tokens = int(result.get("output_tokens") or 0)

            entry = by_model.setdefault(model, {
                "model": model,
                "requests": 0,
                "input_tokens": 0,
                "output_tokens": 0,
            })
            entry["requests"] += requests
            entry["input_tokens"] += input_tokens
            entry["output_tokens"] += output_tokens

            totals["requests"] += requests
            totals["input_tokens"] += input_tokens
            totals["output_tokens"] += output_tokens

    models = sorted(
        by_model.values(),
        key=lambda item: (item["requests"], item["input_tokens"] + item["output_tokens"]),
        reverse=True,
    )
    return totals, models


def _aggregate_openai_costs(buckets: list[dict[str, Any]]) -> tuple[float, list[dict[str, Any]]]:
    total_cost = 0.0
    by_line_item: dict[str, float] = {}

    for bucket in buckets:
        for result in bucket.get("results", []):
            line_item = result.get("line_item")
            if not _is_codex_line_item(line_item):
                continue
            amount = float(((result.get("amount") or {}).get("value")) or 0.0)
            total_cost += amount
            by_line_item[line_item] = by_line_item.get(line_item, 0.0) + amount

    lines = sorted(
        (
            {"line_item": line_item, "cost_usd": amount}
            for line_item, amount in by_line_item.items()
        ),
        key=lambda item: item["cost_usd"],
        reverse=True,
    )
    return total_cost, lines


async def _fetch_codex_usage_payload() -> dict[str, Any]:
    """Fetch OpenAI organization usage data for Codex models."""
    global _openai_usage_cache, _openai_usage_cached_at

    now = time.monotonic()
    if _openai_usage_cache and (now - _openai_usage_cached_at) < _OPENAI_USAGE_CACHE_TTL:
        return _openai_usage_cache

    payload = _empty_codex_usage_payload()
    payload["local_plan"] = _read_local_codex_plan()
    payload["available"] = bool(payload["local_plan"]["available"])

    admin_key = _config_value("OPENAI_ADMIN_KEY")
    if not admin_key:
        if not payload["available"]:
            if payload["configuration"]["project_key_configured"]:
                payload["error"] = (
                    "OPENAI_ADMIN_KEY is required for organization usage endpoints; "
                    "OPENAI_API_KEY alone is not enough."
                )
            else:
                payload["error"] = payload["local_plan"]["error"] or (
                    "Set OPENAI_ADMIN_KEY to enable OpenAI/Codex API usage reporting."
                )
        _openai_usage_cache = payload
        _openai_usage_cached_at = now
        return payload

    headers = {"Authorization": f"Bearer {admin_key}"}
    org_id = _config_value("OPENAI_ORG_ID")
    if org_id:
        headers["OpenAI-Organization"] = org_id

    project_ids = payload["configuration"]["project_ids"]
    start_time = int(time.time()) - (OPENAI_USAGE_WINDOW_DAYS * 24 * 60 * 60)
    usage_params: dict[str, Any] = {
        "start_time": start_time,
        "bucket_width": "1d",
        "limit": OPENAI_USAGE_WINDOW_DAYS,
        "group_by": ["model"],
    }
    cost_params: dict[str, Any] = {
        "start_time": start_time,
        "bucket_width": "1d",
        "limit": OPENAI_USAGE_WINDOW_DAYS,
        "group_by": ["line_item"],
    }
    if project_ids:
        usage_params["project_ids"] = project_ids
        cost_params["project_ids"] = project_ids

    try:
        import aiohttp as aiohttp_client
        async with aiohttp_client.ClientSession() as session:
            async with session.get(
                OPENAI_USAGE_URL,
                headers=headers,
                params=usage_params,
                timeout=aiohttp_client.ClientTimeout(total=10),
            ) as usage_resp:
                if usage_resp.status != 200:
                    detail = await usage_resp.text()
                    if usage_resp.status in {401, 403}:
                        payload["error"] = (
                            "OpenAI rejected the configured credentials."
                            " OPENAI_ADMIN_KEY must be an"
                            " organization admin key with usage access."
                        )
                    else:
                        payload["error"] = (
                            f"OpenAI usage API returned {usage_resp.status}: {detail[:160]}"
                        )
                    payload["api_usage"]["error"] = payload["error"]
                    if payload["available"]:
                        payload["error"] = None
                    _openai_usage_cache = payload
                    _openai_usage_cached_at = now
                    return payload
                usage_body = await usage_resp.json()

            async with session.get(
                OPENAI_COSTS_URL,
                headers=headers,
                params=cost_params,
                timeout=aiohttp_client.ClientTimeout(total=10),
            ) as cost_resp:
                if cost_resp.status != 200:
                    detail = await cost_resp.text()
                    if cost_resp.status in {401, 403}:
                        payload["error"] = (
                            "OpenAI rejected the configured credentials."
                            " OPENAI_ADMIN_KEY must be an"
                            " organization admin key with usage access."
                        )
                    else:
                        payload["error"] = (
                            f"OpenAI costs API returned {cost_resp.status}: {detail[:160]}"
                        )
                    payload["api_usage"]["error"] = payload["error"]
                    if payload["available"]:
                        payload["error"] = None
                    _openai_usage_cache = payload
                    _openai_usage_cached_at = now
                    return payload
                cost_body = await cost_resp.json()
    except Exception as exc:
        logger.warning("Failed to fetch OpenAI usage: %s", exc)
        payload["error"] = str(exc)
        payload["api_usage"]["error"] = payload["error"]
        if payload["available"]:
            payload["error"] = None
        _openai_usage_cache = payload
        _openai_usage_cached_at = now
        return payload

    usage_buckets = usage_body.get("data", [])
    cost_buckets = cost_body.get("data", [])
    totals, models = _aggregate_openai_usage(usage_buckets)
    cost_total, cost_lines = _aggregate_openai_costs(cost_buckets)

    payload["available"] = True
    payload["api_usage"] = {
        "available": True,
        "requests": totals["requests"],
        "input_tokens": totals["input_tokens"],
        "output_tokens": totals["output_tokens"],
        "cost_usd": cost_total,
        "models": models[:5],
        "cost_breakdown": cost_lines[:5],
        "scope": "organization" if not project_ids else "project",
        "error": None,
    }
    payload["available"] = payload["available"] or payload["api_usage"]["available"]
    if not payload["local_plan"]["available"] and not (totals["requests"] or cost_total or models):
        payload["error"] = "No Codex usage signals found yet."
    else:
        payload["error"] = None

    _openai_usage_cache = payload
    _openai_usage_cached_at = now
    return payload


# ---------------------------------------------------------------------------
# Endpoint handlers
# ---------------------------------------------------------------------------

async def utilization(request: web.Request) -> web.Response:
    """Fetch Claude API rate limit utilization from OAuth endpoint."""
    payload, status = await _fetch_claude_usage_payload()
    if status != 200:
        return web.json_response({"error": payload.get("error", "Unknown error")}, status=status)
    return web.json_response({"windows": payload["windows"], "plan": payload["plan"]})


async def utilization_history(request: web.Request) -> web.Response:
    """Return utilization history from the database."""
    global _pending_util_snapshot
    store = request.app.get("store")
    if not store:
        return web.json_response([])

    # Flush any pending snapshot
    if _pending_util_snapshot:
        with contextlib.suppress(Exception):
            await store.record_utilization(
                _pending_util_snapshot["timestamp"],
                _pending_util_snapshot["five_hour"],
                _pending_util_snapshot["seven_day"],
            )
        _pending_util_snapshot = None

    history = await store.get_utilization_history(limit=2000)
    return web.json_response(history)


async def usage_overview(request: web.Request) -> web.Response:
    """Fetch combined Claude and Codex usage for the dashboard."""
    claude_payload, _ = await _fetch_claude_usage_payload()
    codex_payload = await _fetch_codex_usage_payload()
    return web.json_response({
        "claude": claude_payload,
        "codex": codex_payload,
    })
