"""レート制限。

軸は 4 つ。全キー横断の同時実行、キー単位の同時実行、毎分・毎日のリクエスト数、
毎分トークン数（ベスト エフォート）。取得順は必ず

    全キー横断の同時実行 -> キー単位の同時実行 -> 毎分 / 毎日 / トークン

とする。先に横断上限を取ることで、1 本のキーが上流のスロットを占有して
他の利用者を停止させる事態を防ぐ。解放は呼び出し側の ``try/finally`` で必ず行う。
"""

from __future__ import annotations

import asyncio
import logging
import math
import sqlite3
import time
from collections import deque
from dataclasses import dataclass, field
from datetime import timedelta

from . import errors
from .config import Settings
from .db import Database, iso_jst, now_jst, parse_iso
from .keys import KeyRecord

logger = logging.getLogger(__name__)

MINUTE = 60.0
_SNAPSHOT_MINUTE = "minute"
_SNAPSHOT_DAY_PREFIX = "day:"


def _seconds_until_jst_midnight() -> int:
    """次の JST 日付境界までの秒数。"""

    now = now_jst()
    tomorrow = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
    return max(1, int((tomorrow - now).total_seconds()))


@dataclass(slots=True)
class KeyState:
    """1 キー分のカウンタ。"""

    minute_hits: deque[float] = field(default_factory=deque)
    day_date: str = ""
    day_count: int = 0
    token_hits: deque[tuple[float, int]] = field(default_factory=deque)
    semaphore: asyncio.Semaphore | None = None
    semaphore_capacity: int | None = None

    def prune(self, now: float) -> None:
        cutoff = now - MINUTE
        while self.minute_hits and self.minute_hits[0] <= cutoff:
            self.minute_hits.popleft()
        while self.token_hits and self.token_hits[0][0] <= cutoff:
            self.token_hits.popleft()

    def tokens_in_window(self) -> int:
        return sum(count for _, count in self.token_hits)


@dataclass(slots=True)
class Reservation:
    """取得済みのスロット。``release`` は何度呼んでも安全。"""

    limiter: "RateLimiter"
    key_semaphore: asyncio.Semaphore | None = None
    _released: bool = False

    def release(self) -> None:
        if self._released:
            return
        self._released = True
        if self.key_semaphore is not None:
            self.key_semaphore.release()
        self.limiter.release_global()


class RateLimiter:
    """プロセス内のカウンタで制限を判定する。"""

    def __init__(self, settings: Settings) -> None:
        self._settings = settings
        self._global = asyncio.Semaphore(settings.global_concurrency)
        self._queue_timeout = settings.global_queue_timeout
        self._states: dict[int, KeyState] = {}

    # ------------------------------------------------------------------ 状態

    def state(self, key_id: int) -> KeyState:
        state = self._states.get(key_id)
        if state is None:
            state = KeyState()
            self._states[key_id] = state
        return state

    def forget(self, key_id: int) -> None:
        """削除されたキーの状態を捨てる。"""

        self._states.pop(key_id, None)

    # ------------------------------------------------------ 全キー横断の上限

    async def acquire_global(self) -> None:
        """横断上限を取得する。待ち切れなければ 429。"""

        if self._queue_timeout <= 0:
            if self._global.locked():
                raise errors.rate_limited(
                    "The gateway is at its global concurrency limit. Please retry shortly.",
                    retry_after=1,
                    code="gateway_busy",
                )
            await self._global.acquire()
            return
        try:
            await asyncio.wait_for(self._global.acquire(), timeout=self._queue_timeout)
        except (asyncio.TimeoutError, TimeoutError) as exc:
            raise errors.rate_limited(
                "The gateway is at its global concurrency limit. Please retry shortly.",
                retry_after=math.ceil(self._queue_timeout),
                code="gateway_busy",
            ) from exc

    def release_global(self) -> None:
        self._global.release()

    # ------------------------------------------------------------------ 取得

    async def acquire(self, key: KeyRecord) -> Reservation:
        """横断上限とキー単位の同時実行を取得し、毎分・毎日の上限を検査する。

        いずれかで弾かれた場合、取得済みのスロットは本メソッド内で解放する。
        """

        await self.acquire_global()
        reservation = Reservation(self)
        try:
            reservation.key_semaphore = await self._acquire_key_slot(key)
            self.check_rates(key)
        except BaseException:
            reservation.release()
            raise
        return reservation

    async def _acquire_key_slot(self, key: KeyRecord) -> asyncio.Semaphore | None:
        """キー単位の同時実行を待たずに取得する。空きが無ければ 429。"""

        capacity = key.limits.concurrency
        if capacity is None or capacity <= 0:
            return None
        state = self.state(key.id)
        if state.semaphore is None or state.semaphore_capacity != capacity:
            # 上限が変更された場合は新しい Semaphore へ切り替える。
            # 実行中のリクエストは取得時の Semaphore を持ち続けるため解放は破綻しない。
            state.semaphore = asyncio.Semaphore(capacity)
            state.semaphore_capacity = capacity
        semaphore = state.semaphore
        if semaphore.locked():
            raise errors.rate_limited(
                f"Concurrent request limit reached for this API key (limit: {capacity}).",
                retry_after=1,
                code="concurrency_limit_exceeded",
            )
        # locked() が偽の間に await は挟まれないため、この acquire は待たずに完了する。
        await semaphore.acquire()
        return semaphore

    def check_rates(self, key: KeyRecord, *, record: bool = True) -> None:
        """毎分・毎日・毎分トークンの上限を検査し、通過した場合のみ計上する。"""

        now = time.monotonic()
        state = self.state(key.id)
        state.prune(now)

        rpm = key.limits.rpm
        if rpm is not None and rpm > 0 and len(state.minute_hits) >= rpm:
            oldest = state.minute_hits[0]
            retry_after = max(1, math.ceil(MINUTE - (now - oldest)))
            raise errors.rate_limited(
                f"Rate limit reached for this API key (limit: {rpm} requests per minute).",
                retry_after=retry_after,
            )

        today = now_jst().strftime("%Y-%m-%d")
        if state.day_date != today:
            state.day_date = today
            state.day_count = 0

        rpd = key.limits.rpd
        if rpd is not None and rpd > 0 and state.day_count >= rpd:
            raise errors.rate_limited(
                f"Daily limit reached for this API key (limit: {rpd} requests per day).",
                retry_after=_seconds_until_jst_midnight(),
                code="daily_limit_exceeded",
            )

        tpm = key.limits.tpm
        if tpm is not None and tpm > 0 and state.token_hits:
            used = state.tokens_in_window()
            if used >= tpm:
                oldest = state.token_hits[0][0]
                retry_after = max(1, math.ceil(MINUTE - (now - oldest)))
                raise errors.rate_limited(
                    f"Token rate limit reached for this API key (limit: {tpm} tokens per minute).",
                    retry_after=retry_after,
                    code="token_rate_limit_exceeded",
                )

        if record:
            state.minute_hits.append(now)
            state.day_count += 1

    def record_tokens(self, key_id: int, tokens: int) -> None:
        """応答から読み取れたトークン数を計上する。"""

        if tokens <= 0:
            return
        state = self.state(key_id)
        now = time.monotonic()
        state.prune(now)
        state.token_hits.append((now, tokens))

    # ---------------------------------------------------------- スナップショット

    async def restore(self, db: Database) -> None:
        """起動時にカウンタを復元する。

        プロセスの異常終了で制限がリセットされることを悪用させないため。
        """

        rows = await db.fetchall('SELECT key_id, "window", count, updated_at FROM rate_snapshots')
        today = now_jst().strftime("%Y-%m-%d")
        wall_now = now_jst()
        mono_now = time.monotonic()
        stale: list[tuple[int, str]] = []
        for row in rows:
            window = row["window"]
            key_id = int(row["key_id"])
            count = int(row["count"])
            updated = parse_iso(row["updated_at"])
            if window == _SNAPSHOT_MINUTE:
                if updated is None or (wall_now - updated).total_seconds() >= MINUTE or count <= 0:
                    stale.append((key_id, window))
                    continue
                age = (wall_now - updated).total_seconds()
                stamp = mono_now - age
                state = self.state(key_id)
                state.minute_hits.extend([stamp] * count)
            elif window.startswith(_SNAPSHOT_DAY_PREFIX):
                date = window[len(_SNAPSHOT_DAY_PREFIX) :]
                if date != today:
                    stale.append((key_id, window))
                    continue
                state = self.state(key_id)
                state.day_date = date
                state.day_count = count
        if stale:
            await db.run(lambda conn: _delete_snapshots(conn, stale))

    async def snapshot(self, db: Database) -> None:
        """現在のカウンタを SQLite へ書き出す。"""

        now = time.monotonic()
        stamp = iso_jst()
        rows: list[tuple[int, str, int, str]] = []
        for key_id, state in self._states.items():
            state.prune(now)
            if state.minute_hits:
                rows.append((key_id, _SNAPSHOT_MINUTE, len(state.minute_hits), stamp))
            if state.day_date and state.day_count:
                rows.append((key_id, _SNAPSHOT_DAY_PREFIX + state.day_date, state.day_count, stamp))
        await db.run(lambda conn: _upsert_snapshots(conn, rows, stamp))

    async def snapshot_loop(self, db: Database, interval: float) -> None:
        """周期スナップショット。呼び出し側がタスクとして起動する。"""

        while True:
            try:
                await asyncio.sleep(interval)
                await self.snapshot(db)
            except asyncio.CancelledError:
                raise
            except Exception:  # pragma: no cover
                logger.exception("レート制限カウンタのスナップショットに失敗しました")


def _delete_snapshots(conn: sqlite3.Connection, rows: list[tuple[int, str]]) -> None:
    conn.executemany('DELETE FROM rate_snapshots WHERE key_id = ? AND "window" = ?', rows)


def _upsert_snapshots(conn: sqlite3.Connection, rows: list[tuple[int, str, int, str]], stamp: str) -> None:
    if rows:
        conn.executemany(
            'INSERT INTO rate_snapshots (key_id, "window", count, updated_at) VALUES (?, ?, ?, ?)'
            ' ON CONFLICT(key_id, "window") DO UPDATE SET count = excluded.count,'
            " updated_at = excluded.updated_at",
            rows,
        )
    # 昨日以前の日次スナップショットは復元対象にならないため掃除する。
    today_window = _SNAPSHOT_DAY_PREFIX + now_jst().strftime("%Y-%m-%d")
    conn.execute(
        'DELETE FROM rate_snapshots WHERE "window" LIKE ? AND "window" <> ?',
        (_SNAPSHOT_DAY_PREFIX + "%", today_window),
    )