# multi_inference.py

import requests
import os

# Load DeepSeek keys from Hugging Face Secrets
DEEPSEEK_KEYS = [
    os.getenv("DEEPSEEK_KEY_1", "").strip(),
    os.getenv("DEEPSEEK_KEY_2", "").strip(),
    os.getenv("DEEPSEEK_KEY_3", "").strip(),
    os.getenv("DEEPSEEK_KEY_4", "").strip(),
    os.getenv("DEEPSEEK_KEY_5", "").strip()
]

# Load OpenAI keys named as open1, open2, etc.
OPENAI_KEYS = [
    os.getenv("open1", "").strip(),
    os.getenv("open2", "").strip(),
    os.getenv("open3", "").strip(),
    os.getenv("open4", "").strip(),
    os.getenv("open5", "").strip()
]


def try_deepseek(prompt, key):
    try:
        url = "https://api.deepseek.com/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json"
        }
        data = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ]
        }
        res = requests.post(url, headers=headers, json=data)
        result = res.json()
        if "choices" in result:
            return result["choices"][0]["message"]["content"]
        return f"[ERROR] DeepSeek: {result.get('error', {}).get('message', 'Unknown Error')}"
    except Exception as e:
        return f"[ERROR] DeepSeek Exception: {str(e)}"


def try_openai(prompt, key):
    try:
        url = "https://api.openai.com/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json"
        }
        data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ]
        }
        res = requests.post(url, headers=headers, json=data)
        result = res.json()
        if "choices" in result:
            return result["choices"][0]["message"]["content"]
        return f"[ERROR] OpenAI: {result.get('error', {}).get('message', 'Unknown Error')}"
    except Exception as e:
        return f"[ERROR] OpenAI Exception: {str(e)}"


# Main smart fallback wrapper
def multi_query(prompt):
    for key in DEEPSEEK_KEYS:
        output = try_deepseek(prompt, key)
        if not output.startswith("[ERROR]"):
            return output

    for key in OPENAI_KEYS:
        output = try_openai(prompt, key)
        if not output.startswith("[ERROR]"):
            return output

    return "[ALL APIs FAILED: Out of quota or network error]"