from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone

from database import get_db
from remnawave import RemnawaveClient, RemnawaveError

logger = logging.getLogger("traffic_grants")


def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


async def revert_expired_grants() -> int:
    """Снимает временный трафик после expires_at. Возвращает число успешных откатов."""
    now = _now_iso()
    db = await get_db()
    try:
        cursor = await db.execute(
            """
            SELECT * FROM traffic_grants
            WHERE status = 'active' AND expires_at <= ?
            ORDER BY id ASC
            """,
            (now,),
        )
        grants = [dict(row) for row in await cursor.fetchall()]
    finally:
        await db.close()

    if not grants:
        return 0

    client = RemnawaveClient()
    done = 0
    for grant in grants:
        try:
            result = await client.remove_traffic_bytes(
                grant["user_uuid"], int(grant["bytes_added"])
            )
            db = await get_db()
            try:
                await db.execute(
                    """
                    UPDATE traffic_grants
                    SET status = 'reverted', reverted_at = ?, last_error = ''
                    WHERE id = ?
                    """,
                    (_now_iso(), grant["id"]),
                )
                await db.commit()
            finally:
                await db.close()
            done += 1
            logger.info(
                "Reverted traffic grant #%s user=%s -%s bytes (limit %s → %s)",
                grant["id"],
                grant.get("username") or grant["user_uuid"],
                grant["bytes_added"],
                result.get("old_limit_bytes"),
                result.get("new_limit_bytes"),
            )
        except RemnawaveError as exc:
            db = await get_db()
            try:
                await db.execute(
                    """
                    UPDATE traffic_grants
                    SET last_error = ?
                    WHERE id = ?
                    """,
                    (str(exc)[:500], grant["id"]),
                )
                await db.commit()
            finally:
                await db.close()
            logger.warning("Failed to revert grant #%s: %s", grant["id"], exc)

    return done


async def traffic_grant_worker(stop_event: asyncio.Event) -> None:
    while not stop_event.is_set():
        try:
            await revert_expired_grants()
        except Exception:
            logger.exception("traffic grant worker error")
        try:
            await asyncio.wait_for(stop_event.wait(), timeout=60)
        except asyncio.TimeoutError:
            pass