# multi_inference.py

import requests
import os

# DeepSeek API keys
DEEPSEEK_KEYS = [
    os.getenv("DEEPSEEK_KEY", "").strip(),
    os.getenv("DEEPSEEK_KEY_1", "").strip(),
    os.getenv("DEEPSEEK_KEY_2", "").strip(),
    os.getenv("DEEPSEEK_KEY_3", "").strip(),
    os.getenv("DEEPSEEK_KEY_4", "").strip()
]

# OpenAI keys (named 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)}"

def multi_query(prompt):
    for key in DEEPSEEK_KEYS:
        if key:
            response = try_deepseek(prompt, key)
            if not response.startswith("[ERROR]"):
                return response

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

    return "[ALL APIs FAILED: Out of quota or invalid keys]"