"""LLM guardrails: concurrency limits, circuit breaker, and routing snapshots."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional

import redis

from .. import settings

log = logging.getLogger(__name__)

_redis_client: Optional[redis.Redis] = None


def _get_redis() -> Optional[redis.Redis]:
    global _redis_client
    if _redis_client is not None:
        return _redis_client
    redis_url = settings.REDIS_URL or "redis://127.0.0.1:6379/0"
    try:
        _redis_client = redis.from_url(redis_url, decode_responses=True)
        return _redis_client
    except Exception as exc:
        log.warning("LLM guard Redis init failed: %s", exc)
        return None


class LLMConcurrencyError(RuntimeError):
    """Raised when LLM concurrency limit is reached."""

    def __init__(
        self,
        retry_after_s: int = 30,
        user_message: str = "Too many requests right now. Please try again in 30 seconds.",
    ) -> None:
        super().__init__(user_message)
        self.retry_after_s = retry_after_s
        self.user_message = user_message


class LLMCircuitOpenError(RuntimeError):
    """Raised when the circuit breaker is open for a model."""

    def __init__(
        self,
        retry_after_s: int,
        user_message: str = "Service temporarily unavailable. Please try again shortly.",
    ) -> None:
        super().__init__(user_message)
        self.retry_after_s = retry_after_s
        self.user_message = user_message


@dataclass
class RoutingSnapshot:
    """Pinned routing state for a job/session."""

    primary_model: str
    fallback_model: Optional[str]
    using_fallback: bool
    locked_at: str
    tier: Optional[str] = None

    def to_json(self) -> str:
        return json.dumps(
            {
                "primary_model": self.primary_model,
                "fallback_model": self.fallback_model,
                "using_fallback": self.using_fallback,
                "locked_at": self.locked_at,
                "tier": self.tier,
            }
        )

    @classmethod
    def from_json(cls, raw: str) -> Optional["RoutingSnapshot"]:
        try:
            payload = json.loads(raw)
            return cls(
                primary_model=str(payload.get("primary_model") or ""),
                fallback_model=payload.get("fallback_model"),
                using_fallback=bool(payload.get("using_fallback")),
                locked_at=str(payload.get("locked_at") or ""),
                tier=payload.get("tier"),
            )
        except Exception:
            return None


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _concurrency_ttl_s(timeout_s: Optional[float]) -> int:
    base = int(
        max(
            float(timeout_s or 0),
            float(settings.LLM_HEAVY_REASON_BUDGET_S),
        )
    )
    return max(60, base + 30)


def _circuit_key(model: str) -> str:
    return f"llm:circuit:{model}"


def _circuit_fail_key(model: str) -> str:
    return f"llm:circuit:{model}:fail"


def is_circuit_open(model: str) -> bool:
    redis_conn = _get_redis()
    if not redis_conn:
        return False
    try:
        return bool(redis_conn.get(_circuit_key(model)))
    except Exception as exc:
        log.debug("LLM circuit check failed: %s", exc)
        return False


def record_failure(model: str) -> bool:
    """Return True if circuit opens."""
    redis_conn = _get_redis()
    if not redis_conn:
        return False
    try:
        count = redis_conn.incr(_circuit_fail_key(model))
        redis_conn.expire(_circuit_fail_key(model), settings.LLM_CIRCUIT_RESET_S)
        if count >= settings.LLM_CIRCUIT_THRESHOLD:
            redis_conn.set(
                _circuit_key(model),
                "open",
                ex=settings.LLM_CIRCUIT_RESET_S,
            )
            return True
    except Exception as exc:
        log.debug("LLM circuit record failed: %s", exc)
    return False


def record_success(model: str) -> None:
    redis_conn = _get_redis()
    if not redis_conn:
        return
    try:
        redis_conn.delete(_circuit_fail_key(model))
    except Exception as exc:
        log.debug("LLM circuit reset failed: %s", exc)


class LLMGuard:
    """Context manager for circuit breaker + concurrency limits."""

    def __init__(
        self,
        *,
        model: str,
        account_id: Optional[str],
        timeout_s: Optional[float],
    ) -> None:
        self.model = model
        self.account_id = account_id
        self.timeout_s = timeout_s
        self.redis = _get_redis()
        self._acquired_global = False
        self._acquired_org = False

    def __enter__(self) -> "LLMGuard":
        if is_circuit_open(self.model):
            raise LLMCircuitOpenError(settings.LLM_CIRCUIT_RESET_S)

        if not self.redis:
            return self

        ttl_s = _concurrency_ttl_s(self.timeout_s)
        global_key = "llm:global:semaphore"
        org_key = (
            f"llm:org:{self.account_id}" if self.account_id else None
        )

        try:
            if settings.LLM_GLOBAL_CONCURRENCY > 0:
                current = int(self.redis.incr(global_key))
                self.redis.expire(global_key, ttl_s)
                if current > settings.LLM_GLOBAL_CONCURRENCY:
                    self.redis.decr(global_key)
                    raise LLMConcurrencyError()
                self._acquired_global = True

            if org_key and settings.LLM_ORG_CONCURRENCY > 0:
                current = int(self.redis.incr(org_key))
                self.redis.expire(org_key, ttl_s)
                if current > settings.LLM_ORG_CONCURRENCY:
                    if self._acquired_global:
                        self.redis.decr(global_key)
                        self._acquired_global = False
                    self.redis.decr(org_key)
                    raise LLMConcurrencyError()
                self._acquired_org = True
        except LLMConcurrencyError:
            raise
        except Exception as exc:
            log.debug("LLM guard Redis error: %s", exc)
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if not self.redis:
            return
        try:
            if self._acquired_global:
                self.redis.decr("llm:global:semaphore")
        except Exception as exc:
            log.debug("LLM guard global release failed: %s", exc)
        try:
            if self._acquired_org and self.account_id:
                self.redis.decr(f"llm:org:{self.account_id}")
        except Exception as exc:
            log.debug("LLM guard org release failed: %s", exc)


def get_routing_snapshot(key: str) -> Optional[RoutingSnapshot]:
    redis_conn = _get_redis()
    if not redis_conn:
        return None
    try:
        raw = redis_conn.get(key)
        if not raw:
            return None
        return RoutingSnapshot.from_json(raw)
    except Exception as exc:
        log.debug("Routing snapshot read failed: %s", exc)
        return None


def save_routing_snapshot(key: str, snapshot: RoutingSnapshot, ttl_s: int) -> None:
    redis_conn = _get_redis()
    if not redis_conn:
        return
    try:
        redis_conn.set(key, snapshot.to_json(), ex=ttl_s)
    except Exception as exc:
        log.debug("Routing snapshot write failed: %s", exc)


def ensure_routing_snapshot(
    key: str,
    *,
    primary_model: str,
    fallback_model: Optional[str],
    ttl_s: int,
    tier: Optional[str] = None,
) -> RoutingSnapshot:
    snapshot = get_routing_snapshot(key)
    if snapshot:
        return snapshot
    snapshot = RoutingSnapshot(
        primary_model=primary_model,
        fallback_model=fallback_model,
        using_fallback=False,
        locked_at=_now_iso(),
        tier=tier,
    )
    save_routing_snapshot(key, snapshot, ttl_s)
    return snapshot
