"""
Phase 82: Distributed Rate Limiting — Global token bucket across the mesh.

Shared rate-limit buckets that can be synced across nodes. Each bucket uses
the token-bucket algorithm with configurable capacity and refill rate.

Storage:
  state/broker/rate_limits/
    {bucket_name}.json     — bucket definition + current state
    consumption.jsonl      — consumption event log
"""

import json
import logging
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

LIMITS_DIR = BROKER_DIR / "rate_limits"
CONSUMPTION_LOG = LIMITS_DIR / "consumption.jsonl"

for _d in (LIMITS_DIR,):
    _d.mkdir(parents=True, exist_ok=True)


def reset() -> None:
    """Reset. Used by tests."""
    pass  # All state is file-based


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _bucket_file(name: str) -> Path:
    safe = name.replace("/", "_").replace("\\", "_")
    return LIMITS_DIR / f"{safe}.json"


def _load_bucket(name: str) -> Optional[dict]:
    f = _bucket_file(name)
    if not f.exists():
        return None
    try:
        return json.loads(f.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_bucket(bucket: dict) -> None:
    _write_json_sync(_bucket_file(bucket["name"]), bucket)


def _refill(bucket: dict) -> dict:
    """Refill tokens based on elapsed time since last refill."""
    now = time.time()
    elapsed = now - bucket.get("last_refill", now)
    if elapsed > 0:
        added = elapsed * bucket["refill_rate_per_sec"]
        bucket["tokens"] = min(bucket["capacity"], bucket["tokens"] + added)
        bucket["last_refill"] = now
    return bucket


# ── Public API ─────────────────────────────────────────────────────────────

def create_bucket(
    name: str,
    capacity: float,
    refill_rate_per_sec: float,
    scope: str = "global",
) -> dict:
    """Create a rate-limit bucket."""
    if capacity <= 0 or refill_rate_per_sec < 0:
        raise ValueError("capacity must be > 0 and refill_rate_per_sec must be >= 0")

    bucket = {
        "name": name,
        "capacity": capacity,
        "refill_rate_per_sec": refill_rate_per_sec,
        "scope": scope,
        "tokens": capacity,  # Start full
        "last_refill": time.time(),
        "created_at": _now(),
        "total_consumed": 0,
        "total_rejected": 0,
    }
    _save_bucket(bucket)
    logger.info("Created rate-limit bucket: %s (capacity=%s, rate=%s/s)", name, capacity, refill_rate_per_sec)
    return bucket


def consume(name: str, tokens: int = 1, node_id: Optional[str] = None) -> Tuple[bool, float, Optional[float]]:
    """
    Consume tokens from a bucket.
    Returns (allowed, remaining, retry_after_seconds).
    retry_after is None if allowed, otherwise seconds until enough tokens.
    """
    bucket = _load_bucket(name)
    if not bucket:
        raise KeyError(f"Bucket '{name}' not found")

    bucket = _refill(bucket)

    if bucket["tokens"] >= tokens:
        bucket["tokens"] -= tokens
        bucket["total_consumed"] = bucket.get("total_consumed", 0) + tokens
        _save_bucket(bucket)
        _append_jsonl_sync(CONSUMPTION_LOG, {
            "bucket": name, "tokens": tokens, "allowed": True,
            "node_id": node_id, "ts": _now(), "remaining": bucket["tokens"],
        })
        return (True, bucket["tokens"], None)
    else:
        # Calculate retry_after
        deficit = tokens - bucket["tokens"]
        rate = bucket["refill_rate_per_sec"]
        retry_after = deficit / rate if rate > 0 else float("inf")
        bucket["total_rejected"] = bucket.get("total_rejected", 0) + 1
        _save_bucket(bucket)
        _append_jsonl_sync(CONSUMPTION_LOG, {
            "bucket": name, "tokens": tokens, "allowed": False,
            "node_id": node_id, "ts": _now(), "remaining": bucket["tokens"],
            "retry_after": retry_after,
        })
        return (False, bucket["tokens"], retry_after)


def get_bucket(name: str) -> Optional[dict]:
    """Get current bucket state with tokens refilled."""
    bucket = _load_bucket(name)
    if bucket:
        bucket = _refill(bucket)
        _save_bucket(bucket)
    return bucket


def list_buckets() -> List[dict]:
    """List all rate-limit buckets."""
    buckets = []
    if not LIMITS_DIR.exists():
        return buckets
    for f in sorted(LIMITS_DIR.glob("*.json")):
        try:
            b = json.loads(f.read_text("utf-8"))
            if "name" in b and "capacity" in b:
                b = _refill(b)
                buckets.append(b)
        except (json.JSONDecodeError, OSError):
            continue
    return buckets


def delete_bucket(name: str) -> bool:
    """Delete a rate-limit bucket."""
    f = _bucket_file(name)
    if not f.exists():
        return False
    f.unlink()
    return True


def sync_bucket(name: str, remote_state: dict) -> dict:
    """
    Merge bucket state from a peer node.
    Takes the minimum token count (most conservative) and merges counters.
    """
    local = _load_bucket(name)
    if not local:
        # Accept remote state as-is
        _save_bucket(remote_state)
        return remote_state

    local = _refill(local)
    merged = dict(local)
    merged["tokens"] = min(local["tokens"], remote_state.get("tokens", local["tokens"]))
    merged["total_consumed"] = max(
        local.get("total_consumed", 0),
        remote_state.get("total_consumed", 0),
    )
    merged["total_rejected"] = max(
        local.get("total_rejected", 0),
        remote_state.get("total_rejected", 0),
    )
    merged["last_sync"] = _now()
    _save_bucket(merged)
    return merged


def get_consumption_stats(name: str, hours: int = 24) -> dict:
    """Usage statistics over the given time window."""
    bucket = _load_bucket(name)
    if not bucket:
        return {"error": f"Bucket '{name}' not found"}

    entries = _read_jsonl_sync(CONSUMPTION_LOG)
    relevant = [e for e in entries if e.get("bucket") == name]

    allowed = sum(1 for e in relevant if e.get("allowed"))
    rejected = sum(1 for e in relevant if not e.get("allowed"))
    total_tokens = sum(e.get("tokens", 0) for e in relevant if e.get("allowed"))

    return {
        "bucket": name,
        "window_hours": hours,
        "requests_allowed": allowed,
        "requests_rejected": rejected,
        "tokens_consumed": total_tokens,
        "current_tokens": bucket.get("tokens", 0),
        "capacity": bucket.get("capacity", 0),
    }
