from __future__ import annotations
import asyncio, json, time
from app.services.storage import async_session_maker
from app.models.request_log import RequestLog

_buffer: list[dict] = []
_lock: asyncio.Lock = asyncio.Lock()
_flush_task: asyncio.Task | None = None
FLUSH_INTERVAL = 2.0
FLUSH_BATCH_SIZE = 100


async def log(entry: dict) -> None:
    async with _lock:
        _buffer.append(entry)
        if len(_buffer) >= FLUSH_BATCH_SIZE:
            asyncio.create_task(flush())


async def flush() -> None:
    async with _lock:
        entries = _buffer[:]
        _buffer.clear()
    if not entries:
        return
    try:
        async with async_session_maker() as session:
            for e in entries:
                session.add(RequestLog(
                    request_id=e.get("request_id", ""),
                    provider=e.get("provider", ""),
                    model=e.get("model", ""),
                    endpoint=e.get("endpoint", ""),
                    key_id=e.get("key_id"),
                    hub_key_id=e.get("hub_key_id"),
                    tokens_in=e.get("tokens_in", 0),
                    tokens_out=e.get("tokens_out", 0),
                    cost=e.get("cost", 0),
                    latency_ms=e.get("latency_ms", 0),
                    status_code=e.get("status_code", 200),
                    error=e.get("error"),
                    cached=e.get("cached", 0),
                    cache_type=e.get("cache_type"),
                    streaming=e.get("streaming", 0),
                    fallback_count=e.get("fallback_count", 0),
                    ip=e.get("ip"),
                    session_id=e.get("session_id"),
                    injection_flagged=e.get("injection_flagged", 0),
                    has_images=e.get("has_images", 0),
                    ab_variant=e.get("ab_variant"),
                    cache_read_tokens=e.get("cache_read_tokens", 0),
                    cache_write_tokens=e.get("cache_write_tokens", 0),
                ))
            await session.commit()
    except Exception as exc:
        print(f"[batch_logger] Flush error: {exc}")


async def _periodic_flush():
    while True:
        await asyncio.sleep(FLUSH_INTERVAL)
        async with _lock:
            if _buffer:
                await flush()


def start() -> None:
    global _flush_task
    if _flush_task is None or _flush_task.done():
        _flush_task = asyncio.create_task(_periodic_flush())