import requests
import os
import threading
from dotenv import load_dotenv
from src.chatbot.prompts import AI_MODES, DEFAULT_PROMPT
import logging
import time

# ----------------------------
# Load environment variables
# ----------------------------
load_dotenv()

API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("OPENAI_BASE_URL")

if not API_KEY or not BASE_URL:
    raise ValueError("API_KEY or BASE_URL not set in environment variables.")

# ----------------------------
# Logging setup
# ----------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ----------------------------
# Session storage (thread-safe)
# ----------------------------
sessions_db = {}
db_lock = threading.Lock()
CONTEXT_WINDOW = 15  # Number of previous messages to keep for context

# ----------------------------
# Main LLM call function
# ----------------------------
def call_llm(user_query: str, mode: str = "socratic", session_id: str = None, reasoning_enabled=True):
    """
    Call the best available model with smart routing, exponential backoff retries, and auto-fallback.
    Models: Llama-3.3-70b (Arabic/General) & GPT-OSS-120b (Logic/Math).
    """
    session_id = session_id or "default_user"
    mode = (mode or "socratic").lower().strip()
    system_instruction = AI_MODES.get(mode, DEFAULT_PROMPT).strip()

    # Thread-safe history retrieval
    with db_lock:
        if session_id not in sessions_db:
            sessions_db[session_id] = []
        history = sessions_db[session_id][-CONTEXT_WINDOW:]

    # Build the message payload (shared for all attempts)
    messages = [{"role": "system", "content": system_instruction}]
    messages.extend(history)
    messages.append({"role": "user", "content": user_query})

    # --- SMART ROUTING LOGIC ---
    # Keywords to trigger high-reasoning models
    logic_keywords = ["solve", "math", "code",
                      "physics", "calculate", "احسب", "معادلة", "برمج"]

    if any(word in user_query.lower() for word in logic_keywords):
        # Priority 1: GPT-OSS-120B (Best for logic)
        model_priority = ["openai/gpt-oss-120b",
                          "llama-3.3-70b-versatile"]
    else:
        # Priority 1: Llama-3.3 (Best for Arabic/General chat)
        model_priority = [
            "llama-3.3-70b-versatile", "openai/gpt-oss-120b"]

    # --- API URL SANITIZATION ---
    # Prevent double /v1 suffix in the base URL
    base_url_cleaned = BASE_URL.rstrip("/")
    if not base_url_cleaned.endswith("/v1") and "/v1" not in base_url_cleaned:
        url = f"{base_url_cleaned}/v1/chat/completions"
    else:
        url = f"{base_url_cleaned}/chat/completions"
    headers = {
        "Authorization": f"Bearer {API_KEY.strip()}",
        "Content-Type": "application/json"
    }

    last_error = ""
    MAX_RETRIES = 2       # Number of retries per model
    INITIAL_DELAY = 2     # Initial sleep time in seconds

    # --- FALLBACK LOOP (Iterate through models) ---
    for current_model in model_priority:
        payload = {
            "model": current_model,
            "messages": messages,
            "max_tokens": 2048,
            "temperature": 0.7,
            "stream": False
        }

        # --- RETRY LOOP (Exponential Backoff per model) ---
        for attempt in range(MAX_RETRIES + 1):
            try:
                logger.info(
                    f"Attempt {attempt + 1} with {current_model} (Session: {session_id})")

                # Request timeout set to 40s to allow heavy reasoning models to finish
                response = requests.post(
                    url, headers=headers, json=payload, timeout=40)

                # Case 1: Success
                if response.status_code == 200:
                    result = response.json()
                    choice = result.get("choices", [{}])[0].get("message", {})
                    answer = choice.get("content", "No content returned.")
                    reasoning_details = choice.get("reasoning_details")

                    # Thread-safe history update
                    with db_lock:
                        sessions_db[session_id].append(
                            {"role": "user", "content": user_query})
                        sessions_db[session_id].append({
                            "role": "assistant",
                            "content": answer,
                            "reasoning_details": reasoning_details,
                            "model_used": current_model
                        })
                    return answer

                # Case 2: Unauthorized (Do not retry, check .env)
                elif response.status_code == 401:
                    return "Error: Unauthorized. Check API Key in .env"

                # Case 3: Retriable errors (Rate limits 429 or Server errors 5xx)
                elif response.status_code in [429, 500, 502, 503, 504]:
                    last_error = f"Model {current_model} returned {response.status_code}"
                    if attempt < MAX_RETRIES:
                        # Exponential backoff: 2s, 4s, etc.
                        delay = INITIAL_DELAY * (2 ** attempt)
                        logger.warning(
                            f"{last_error}. Retrying in {delay}s...")
                        time.sleep(delay)
                        continue  # Retry the same model
                    else:
                        logger.error(
                            f"{current_model} exhausted all retries. Falling back...")
                        break  # Move to next model in priority list

                # Case 4: Other non-retriable errors
                else:
                    last_error = f"Status {response.status_code}: {response.text}"
                    logger.error(
                        f"Unrecoverable error for {current_model}: {last_error}")
                    break  # Move to next model

            except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
                last_error = f"Network Error: {str(e)}"
                if attempt < MAX_RETRIES:
                    delay = INITIAL_DELAY * (2 ** attempt)
                    logger.warning(
                        f"Connection issue. Retrying in {delay}s...")
                    time.sleep(delay)
                    continue
                break  # Move to next model
            except Exception as e:
                last_error = f"Unexpected error: {str(e)}"
                logger.error(last_error)
                break  # Move to next model

    # Final response if both models and all retries fail
    return f"All models failed. Last error: {last_error}"