from __future__ import annotations

from typing import Protocol


class RateLimiter(Protocol):
    def check(self, org_id: str, cost: int = 1) -> bool: ...

    def set_limit(self, org_id: str, rpm: int, burst: int | None = None) -> None: ...


class TPMLimiter(Protocol):
    def consume_input(self, org_id: str, tokens: int) -> bool: ...

    def consume_output(self, org_id: str, tokens: int) -> bool: ...

    def set_input_limit(self, org_id: str, tpm: int) -> None: ...

    def set_output_limit(self, org_id: str, tpm: int) -> None: ...

    def set_limit(self, org_id: str, tpm: int) -> None: ...


class ConcurrencyLimiter(Protocol):
    def acquire(self, org_id: str) -> bool: ...

    def release(self, org_id: str) -> None: ...

    def set_limit(self, org_id: str, max_inflight: int) -> None: ...
