import json, time, threading
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

CONFIG_DIR = Path("config")
POLICY_PATH = CONFIG_DIR / "crawl_policy.json"
_lock = threading.RLock()

def _read_json(path: Path) -> Dict[str, Any]:
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        return {}
    return json.loads(path.read_text())

def _write_json(path: Path, data: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2))

def get_policy() -> Dict[str, Any]:
    with _lock:
        data = _read_json(POLICY_PATH)
        if not data:
            data = {
                "global": {
                    "crawl_mode": "all",
                    "per_hour_quota": 100,
                    "per_run_limit": 500,
                    "burst_concurrency": 2,
                    "token_bucket": {"capacity": 100, "refill_rate_per_hour": 100},
                },
                "brokers": {"*": {"crawl_mode": "all"}},
            }
            _write_json(POLICY_PATH, data)
        return data

def set_policy(payload: Dict[str, Any]) -> None:
    assert "global" in payload and "brokers" in payload, "invalid policy"
    _write_json(POLICY_PATH, payload)
    TokenBucketManager.reload_from_policy(payload)

def effective_policy_for(broker: str, policy: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    p = policy or get_policy()
    eff = dict(p.get("global", {}))
    b = p.get("brokers", {})
    if "*" in b:
        eff.update(b["*"])
    if broker in b:
        eff.update(b[broker])
    tb = eff.get("token_bucket", {})
    refill = tb.get("refill_rate_per_hour", eff.get("per_hour_quota", 0))
    eff.setdefault("token_bucket", {})["refill_rate_per_hour"] = refill
    tb.setdefault("capacity", max(1, int(refill) if refill else 1))
    return eff

class TokenBucket:
    def __init__(self, capacity: int, refill_rate_per_sec: float):
        self.capacity = float(capacity)
        self.tokens = float(capacity)
        self.refill_rate = float(refill_rate_per_sec)
        self.timestamp = time.monotonic()
        self.lock = threading.Lock()

    def take(self, n: int = 1) -> float:
        with self.lock:
            now = time.monotonic()
            elapsed = now - self.timestamp
            self.timestamp = now
            self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
            if self.tokens >= n:
                self.tokens -= n
                return 0.0
            deficit = n - self.tokens
            wait = deficit / max(self.refill_rate, 1e-9)
            return wait

class TokenBucketManager:
    _global_bucket: Optional[TokenBucket] = None
    _broker_buckets: Dict[str, TokenBucket] = {}
    _lock = threading.RLock()

    @classmethod
    def reload_from_policy(cls, policy: Optional[Dict[str, Any]] = None) -> None:
        p = policy or get_policy()
        with cls._lock:
            g = p["global"]["token_bucket"]
            cls._global_bucket = TokenBucket(
                capacity=int(g["capacity"]),
                refill_rate_per_sec=float(g["refill_rate_per_hour"]) / 3600.0,
            )
            cls._broker_buckets = {}

    @classmethod
    def get_buckets_for(cls, broker: str) -> Tuple[TokenBucket, TokenBucket]:
        with cls._lock:
            if cls._global_bucket is None:
                cls.reload_from_policy()
            if broker not in cls._broker_buckets:
                eff = effective_policy_for(broker)
                tb = eff["token_bucket"]
                cls._broker_buckets[broker] = TokenBucket(
                    capacity=int(tb["capacity"]),
                    refill_rate_per_sec=float(tb["refill_rate_per_hour"]) / 3600.0,
                )
            return cls._global_bucket, cls._broker_buckets[broker]
