import os
import requests
import gradio as gr
import re
import time

API_KEY = os.getenv("API_KEY")
GEMINI_KEY = os.getenv("GEMINI_KEY")

SYSTEM_PROMPT = "You are a helpful and friendly AI assistant. The owner of the server is KingLuxarc. HedronCreeper made you. (From Dragon Vision)."

def execute_llm_request(messages):
    try:
        response = requests.post(
            "https://api.mistral.ai/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {API_KEY}",
                "Content-Type": "application/json"
            },
            json={
                "model": "mistral-small-latest",
                "messages": messages,
                "temperature": 0.7
            },
            timeout=120
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except Exception:
        pass

    gemini_url = "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"
    gemini_headers = {
        "Authorization": f"Bearer {GEMINI_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            gemini_url,
            headers=gemini_headers,
            json={
                "model": "gemini-3.1-flash-lite",
                "messages": messages,
                "temperature": 0.7
            },
            timeout=120
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except Exception:
        pass

    time.sleep(5)

    try:
        response = requests.post(
            gemini_url,
            headers=gemini_headers,
            json={
                "model": "gemini-3.1-flash-lite",
                "messages": messages,
                "temperature": 0.7
            },
            timeout=120
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except Exception:
        pass

    try:
        response = requests.post(
            gemini_url,
            headers=gemini_headers,
            json={
                "model": "gemma-2-27b-it",
                "messages": messages,
                "temperature": 0.7
            },
            timeout=120
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except Exception as e:
        return f"API Error: All fallback models failed. Final error: {str(e)}"

def respond(message, history, system_prompt):
    if not API_KEY and not GEMINI_KEY:
        return "Error: Missing API_KEY and GEMINI_KEY secrets. Please check your Space settings."

    base_knowledge = "The owner of the server is KingLuxarc. HedronCreeper made you. (From Dragon Vision)."
    
    if system_prompt:
        if base_knowledge not in system_prompt:
            current_system_prompt = f"{system_prompt} {base_knowledge}"
        else:
            current_system_prompt = system_prompt
    else:
        current_system_prompt = SYSTEM_PROMPT

    messages = [
        {
            "role": "system",
            "content": str(current_system_prompt)
        }
    ]

    for item in history:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")
            if role and content:
                text_content = ""
                if isinstance(content, str):
                    text_content = content
                elif isinstance(content, list):
                    for part in content:
                        if isinstance(part, dict) and "text" in part:
                            text_content += str(part["text"])
                        elif isinstance(part, str):
                            text_content += part
                else:
                    text_content = str(content)
                
                messages.append({
                    "role": role,
                    "content": text_content
                })
        else:
            user_msg, bot_msg = item
            
            def parse_msg(msg):
                if isinstance(msg, str):
                    return msg
                if isinstance(msg, dict) and "text" in msg:
                    return str(msg["text"])
                if isinstance(msg, list):
                    return "".join([str(p.get("text", "")) if isinstance(p, dict) else str(p) for p in msg])
                return str(msg)
            
            u_text = parse_msg(user_msg)
            b_text = parse_msg(bot_msg)
            if u_text:
                messages.append({"role": "user", "content": u_text})
            if b_text:
                messages.append({"role": "assistant", "content": b_text})

    messages.append({
        "role": "user",
        "content": str(message)
    })

    reply = execute_llm_request(messages)
    
    if reply.startswith("API Error: All fallback models failed."):
        return reply
        
    reply = re.sub(r'<think>.*?</think>', '', reply, flags=re.DOTALL).strip()
    
    if (reply.startswith("[") or reply.startswith("{")) and "text" in reply:
        match = re.search(r"['\"]text['\"]\s*:\s*['\"](.*?)['\"]", reply)
        if match:
            reply = match.group(1)
    
    return reply

theme = gr.themes.Soft(
    primary_hue="red",
    secondary_hue="zinc",
    neutral_hue="slate"
).set(
    body_background_fill="*neutral_950",
    block_background_fill="*neutral_900",
    block_border_color="*neutral_800",
    button_primary_background_fill="*primary_700",
    button_primary_text_color="white"
)

demo = gr.ChatInterface(
    fn=respond,
    title="🐉 Dragon Vision",
    description="Premium High-Performance AI Interface",
    chatbot=gr.Chatbot(
        height=650,
        show_label=False
    ),
    textbox=gr.Textbox(
        placeholder="Exhaust flame ignited. Enter your command...",
        container=False
    ),
    additional_inputs=[
        gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt")
    ],
    additional_inputs_accordion=gr.Accordion(label="System Configuration", open=False)
)

if __name__ == "__main__":
    demo.launch(theme=theme)