"""Redis cache + in-flight coalescing for identical LiteLLM chat calls."""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
from collections.abc import Awaitable, Callable
from typing import Any

log = logging.getLogger("llm_cache")

_redis = None
_redis_tried = False
_inflight: dict[str, asyncio.Future[str]] = {}
_inflight_lock = asyncio.Lock()


def _get_redis():
    global _redis, _redis_tried
    if _redis_tried:
        return _redis
    _redis_tried = True
    url = os.environ.get("REDIS_URL", "").strip()
    if not url:
        return None
    try:
        import redis

        client = redis.from_url(url, decode_responses=True)
        client.ping()
        _redis = client
    except Exception as exc:  # noqa: BLE001
        log.warning("llm cache redis unavailable: %s", exc)
        _redis = None
    return _redis


def cache_key(
    *,
    model: str,
    messages: list[dict[str, Any]],
    max_tokens: int,
    temperature: float,
) -> str:
    blob = json.dumps(
        {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
        },
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )
    digest = hashlib.sha256(blob.encode("utf-8")).hexdigest()
    return f"ai-router:llm-cache:{digest}"


def get_cached(key: str) -> str | None:
    r = _get_redis()
    if not r:
        return None
    try:
        val = r.get(key)
        return str(val) if val is not None else None
    except Exception as exc:  # noqa: BLE001
        log.warning("llm cache get failed: %s", exc)
        return None


def set_cached(key: str, value: str, ttl_sec: int) -> None:
    r = _get_redis()
    if not r or ttl_sec <= 0:
        return
    try:
        r.setex(key, ttl_sec, value)
    except Exception as exc:  # noqa: BLE001
        log.warning("llm cache set failed: %s", exc)


async def run_cached(
    key: str,
    factory: Callable[[], Awaitable[str]],
    *,
    ttl_sec: int,
    on_hit: Callable[[str], Awaitable[None] | None] | None = None,
) -> tuple[str, str]:
    """
    Return (text, source) where source is cache|coalesce|live.
    Coalesce concurrent identical calls onto one upstream request.
    """
    cached = get_cached(key)
    if cached is not None:
        if on_hit:
            maybe = on_hit("cache")
            if asyncio.iscoroutine(maybe):
                await maybe
        return cached, "cache"

    async with _inflight_lock:
        existing = _inflight.get(key)
        if existing is not None:
            fut: asyncio.Future[str] = existing
            mine = False
        else:
            fut = asyncio.get_running_loop().create_future()
            _inflight[key] = fut
            mine = True

    if not mine:
        if on_hit:
            maybe = on_hit("coalesce")
            if asyncio.iscoroutine(maybe):
                await maybe
        return await fut, "coalesce"

    try:
        text = await factory()
        set_cached(key, text, ttl_sec)
        if not fut.done():
            fut.set_result(text)
        return text, "live"
    except Exception as exc:
        if not fut.done():
            fut.set_exception(exc)
        raise
    finally:
        async with _inflight_lock:
            if _inflight.get(key) is fut:
                _inflight.pop(key, None)