import os
import time
import threading
from dotenv import load_dotenv

load_dotenv()

# Module-level singleton cache
_llm_cache = {}
_llm_cache_lock = threading.Lock()

# ─── Circuit Breaker State ──────────────────────────────────────────────────
# Tracks consecutive LLM failures per provider to avoid hammering a dead endpoint
_circuit_state = {
    "groq":   {"failures": 0, "open_until": 0.0},
    "gemini": {"failures": 0, "open_until": 0.0},
}
_circuit_lock = threading.Lock()
CIRCUIT_FAILURE_THRESHOLD = 3   # open circuit after 3 consecutive failures
CIRCUIT_RECOVERY_TIMEOUT  = 60  # seconds before retrying a broken provider

def _is_circuit_open(provider: str) -> bool:
    """Return True if this provider's circuit is open (broken), False if usable."""
    with _circuit_lock:
        state = _circuit_state[provider]
        if state["failures"] >= CIRCUIT_FAILURE_THRESHOLD:
            if time.time() < state["open_until"]:
                return True          # still in cooldown
            else:
                # Recovery window passed — reset and try again (half-open)
                state["failures"] = 0
                state["open_until"] = 0.0
        return False

def _record_failure(provider: str):
    with _circuit_lock:
        _circuit_state[provider]["failures"] += 1
        if _circuit_state[provider]["failures"] >= CIRCUIT_FAILURE_THRESHOLD:
            _circuit_state[provider]["open_until"] = time.time() + CIRCUIT_RECOVERY_TIMEOUT
            print(f"[circuit_breaker] {provider} circuit OPEN — pausing for {CIRCUIT_RECOVERY_TIMEOUT}s")

def _record_success(provider: str):
    with _circuit_lock:
        _circuit_state[provider]["failures"] = 0
        _circuit_state[provider]["open_until"] = 0.0

# ─── Exponential Backoff Retry ──────────────────────────────────────────────
def _invoke_with_backoff(llm, messages, provider: str, max_retries: int = 3):
    """
    Invoke LLM with exponential backoff retry.
    Retries: 1s → 2s → 4s on failure.
    Records circuit breaker state after each attempt.
    """
    last_error = None
    for attempt in range(max_retries):
        try:
            result = llm.invoke(messages)
            _record_success(provider)
            return result
        except Exception as e:
            last_error = e
            _record_failure(provider)
            wait = 2 ** attempt  # 1s, 2s, 4s
            print(f"[llm_factory] {provider} attempt {attempt + 1}/{max_retries} failed: {str(e)[:80]}. Retrying in {wait}s...")
            if attempt < max_retries - 1:
                time.sleep(wait)
    raise last_error

# ─── LLM Factory ────────────────────────────────────────────────────────────
def get_llm(temperature: float = 0.0):
    """
    Returns a cached LLM client.
    Primary:  Groq  llama-3.1-8b-instant
    Fallback: Gemini 2.0 Flash
    Thread-safe singleton with circuit breaker awareness.
    """
    cache_key = f"temp_{temperature}"

    with _llm_cache_lock:
        if cache_key in _llm_cache:
            return _llm_cache[cache_key]

    # Try Groq
    if not _is_circuit_open("groq"):
        try:
            from langchain_groq import ChatGroq
            api_key = os.getenv("GROQ_API_KEY")
            if not api_key or api_key == "your_groq_key":
                raise ValueError("GROQ_API_KEY not configured")
            llm = ChatGroq(
                model="llama-3.1-8b-instant",
                temperature=temperature,
                groq_api_key=api_key,
                request_timeout=20,
                max_retries=1,      # SDK-level retry; backoff is handled by _invoke_with_backoff
            )
            with _llm_cache_lock:
                _llm_cache[cache_key] = llm
            return llm
        except Exception as e:
            _record_failure("groq")
            print(f"[llm_factory] Groq init failed: {e}. Falling back to Gemini.")
    else:
        print("[llm_factory] Groq circuit is OPEN — skipping directly to Gemini.")

    # Fallback: Gemini
    if not _is_circuit_open("gemini"):
        try:
            from langchain_google_genai import ChatGoogleGenerativeAI
            api_key = os.getenv("GOOGLE_API_KEY")
            if not api_key or api_key == "your_gemini_key":
                raise ValueError("GOOGLE_API_KEY not configured")
            llm = ChatGoogleGenerativeAI(
                model="gemini-2.0-flash",
                temperature=temperature,
                google_api_key=api_key,
                request_timeout=20,
                max_retries=1,
            )
            with _llm_cache_lock:
                _llm_cache[cache_key] = llm
            return llm
        except Exception as fallback_error:
            _record_failure("gemini")
            raise RuntimeError(
                f"Both providers failed. Groq circuit open: {_is_circuit_open('groq')}. "
                f"Gemini error: {fallback_error}"
            )
    else:
        raise RuntimeError(
            "[llm_factory] Both Groq and Gemini circuits are OPEN. "
            "System is in cooldown. Try again in 60 seconds."
        )