import os
import requests
import concurrent.futures
import gradio as gr
import spaces

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
HF_TOKEN = os.environ.get("HF_TOKEN", "")

# Список моделей за провайдерами
OPENROUTER_MODELS = {
    "Gemma 4 26B": "google/gemma-4-26b-a4b-it:free",
    "Nemotron 3 Ultra 550B": "nvidia/nemotron-3-ultra-550b-a55b:free"
}

HF_MODELS = {
    "Qwen 3.8 2.4T": "Qwen/Qwen3.8-2.4T-A95B",
    "GPT OSS 120B": "openai/gpt-oss-120b",
    "Llama 3.1 8B Instruct": "meta-llama/Llama-3.1-8B-Instruct",
    "Xiaomi MiMo V2.6 Flash": "XiaomiMiMo/MiMo-V2.6-Flash-RL",
    "DeepSeek V4 Flash": "deepseek-ai/DeepSeek-V4-Flash-0731",
    "Microsoft Phi-4": "microsoft/phi-4",
    "Hermes 3 Llama 3.1 8B": "NousResearch/Hermes-3-Llama-3.1-8B"
}

ALL_MODELS = {**OPENROUTER_MODELS, **HF_MODELS}

def query_openrouter(model_name, model_id, messages):
    if not OPENROUTER_API_KEY:
        return f"### 🤖 **{model_name}**\n\n❌ *Помилка:* `OPENROUTER_API_KEY` відсутній у Secrets."
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {"model": model_id, "messages": messages}
    try:
        res = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload, timeout=40)
        if res.status_code == 200:
            return f"### 🤖 **{model_name}**\n\n{res.json()['choices'][0]['message']['content']}"
        return f"### 🤖 **{model_name}**\n\n⚠️ *Помилка OpenRouter ({res.status_code}):* {res.text}"
    except Exception as e:
        return f"### 🤖 **{model_name}**\n\n❌ *Помилка з'єднання:* {str(e)}"

def query_huggingface(model_name, model_id, messages):
    if not HF_TOKEN:
        return f"### 🤖 **{model_name}**\n\n❌ *Помилка:* `HF_TOKEN` відсутній у Secrets."
    
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }
    url = f"https://router.huggingface.co/hf-inference/v1/chat/completions"
    payload = {"model": model_id, "messages": messages, "max_tokens": 1024}
    try:
        res = requests.post(url, headers=headers, json=payload, timeout=40)
        if res.status_code == 200:
            return f"### 🤖 **{model_name}**\n\n{res.json()['choices'][0]['message']['content']}"
        return f"### 🤖 **{model_name}**\n\n⚠️ *Помилка HF API ({res.status_code}):* {res.text}"
    except Exception as e:
        return f"### 🤖 **{model_name}**\n\n❌ *Помилка з'єднання:* {str(e)}"

def query_single_model(model_name, messages):
    if model_name in OPENROUTER_MODELS:
        return query_openrouter(model_name, OPENROUTER_MODELS[model_name], messages)
    elif model_name in HF_MODELS:
        return query_huggingface(model_name, HF_MODELS[model_name], messages)
    return f"### 🤖 **{model_name}**\n\n❌ Невідома модель."

@spaces.GPU
def group_chat_fn(user_message, history, selected_models):
    if history is None:
        history = []

    if not selected_models:
        history.append({"role": "user", "content": user_message})
        history.append({"role": "assistant", "content": "⚠️ Обери хоча б одну модель зі списку."})
        return "", history

    messages = [{"role": msg.get("role", "user"), "content": msg.get("content", "")} for msg in history if isinstance(msg, dict)]
    messages.append({"role": "user", "content": user_message})

    history.append({"role": "user", "content": user_message})

    responses = []
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = [
            executor.submit(query_single_model, m_name, messages)
            for m_name in selected_models
        ]
        for future in concurrent.futures.as_completed(futures):
            responses.append(future.result())

    combined_response = "\n\n---\n\n".join(responses)
    history.append({"role": "assistant", "content": combined_response})

    return "", history

custom_css = """
body, .gradio-container {
    background-color: #0f172a !important;
    color: #f8fafc !important;
}
#sidebar {
    background: #1e293b !important;
    border-radius: 12px;
    padding: 16px;
}
#chatbot {
    background: #0f172a !important;
    border-radius: 12px;
}
"""

with gr.Blocks() as demo:
    gr.Markdown("# ⚡ **Multi-AI Studio**")

    with gr.Row():
        with gr.Column(scale=1, elem_id="sidebar"):
            gr.Markdown("### 🎛️ **Моделі OpenRouter & HF**")
            model_selector = gr.CheckboxGroup(
                choices=list(ALL_MODELS.keys()),
                value=["Gemma 4 26B", "Llama 3.1 8B Instruct"],
                label="Обери нейромережі:"
            )
            clear_btn = gr.Button("🗑️ Очистити чат", variant="secondary")

        with gr.Column(scale=3):
            chatbot = gr.Chatbot(elem_id="chatbot", height=600)
            
            with gr.Row():
                msg_input = gr.Textbox(
                    placeholder="Введіть запит...",
                    show_label=False,
                    scale=4,
                    container=False
                )
                send_btn = gr.Button("Надіслати 🚀", variant="primary", scale=1)

    msg_input.submit(group_chat_fn, inputs=[msg_input, chatbot, model_selector], outputs=[msg_input, chatbot])
    send_btn.click(group_chat_fn, inputs=[msg_input, chatbot, model_selector], outputs=[msg_input, chatbot])
    clear_btn.click(lambda: [], outputs=[chatbot])

demo.launch(css=custom_css)