"""In-process token bucket per principal (single API process). Quotas (rows/day) live in usage_daily."""
from __future__ import annotations

import threading
import time

from findata_core.db import transaction


class TokenBucket:
    def __init__(self):
        self._b: dict[str, tuple[float, float]] = {}
        self._lock = threading.Lock()

    def allow(self, key: str, per_minute: int) -> tuple[bool, int]:
        now = time.monotonic()
        with self._lock:
            tokens, last = self._b.get(key, (float(per_minute), now))
            tokens = min(per_minute, tokens + (now - last) * per_minute / 60.0)
            if tokens >= 1:
                self._b[key] = (tokens - 1, now)
                return True, int(tokens - 1)
            self._b[key] = (tokens, now)
            return False, 0


bucket = TokenBucket()


def record_usage(principal: str, rows: int = 0, bytes_exported: int = 0) -> None:
    with transaction() as conn:
        conn.execute("INSERT INTO usage_daily(principal, day, requests, rows_returned, bytes_exported) VALUES (%s, current_date, 1, %s, %s) "
                     "ON CONFLICT (principal, day) DO UPDATE SET requests=usage_daily.requests+1, rows_returned=usage_daily.rows_returned+EXCLUDED.rows_returned, "
                     "bytes_exported=usage_daily.bytes_exported+EXCLUDED.bytes_exported", (principal, rows, bytes_exported))