import gradio as gr
from openai import OpenAI
import os

# Configuration from HF Secrets
BASE_URL = "https://integrate.api.nvidia.com/v1"
keys_string = os.getenv("API_KEYS", "")
API_KEYS = [k.strip() for k in keys_string.split(",") if k.strip()]

class KeyRotator:
    def __init__(self, keys):
        self.keys = keys
        self.counter = 0
    def get_next_client(self):
        if not self.keys:
            return None
        key = self.keys[self.counter % len(self.keys)]
        self.counter += 1
        return OpenAI(base_url=BASE_URL, api_key=key)

rotator = KeyRotator(API_KEYS)

def predict(message, history):
    client = rotator.get_next_client()
    
    if not client:
        yield "❌ පද්ධති දෝෂයකි: API Keys (Secrets) හමු නොවීය."
        return

    # 🔱 SYSTEM PROMPT: Defining Inachi Identity
    # මෙහිදී AI එකට තමන් කවුද සහ තමන්ව හැදුවේ කවුද යන්න පැහැදිලිව උපදෙස් දී ඇත.
    system_prompt = {
        "role": "system", 
        "content": "You are 'Inachi AI', developed by the 'Inachi Team'. "
        "Always be concise, technical, and maintain the Inachi brand identity. "
            }

    # OpenAI format එකට history සහ System Prompt එක එක් කිරීම
    messages = [system_prompt]
    
    for user_msg, assistant_msg in history:
        if user_msg: messages.append({"role": "user", "content": user_msg})
        if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})

    full_response = ""
    
    try:
        completion = client.chat.completions.create(
            model="z-ai/glm-5.1",
            messages=messages,
            temperature= 1,
            top_p= 1,
            max_tokens=16384,
            stream=True
        )
        
        for chunk in completion:
            if not getattr(chunk, "choices", None):
                continue
            
            # NVIDIA Kimi ගේ 'Thinking' සහ Content කොටස් හසුකර ගැනීම
            reasoning = getattr(chunk.choices[0].delta, "reasoning_content", None)
            content = getattr(chunk.choices[0].delta, "content", None)
            
            if reasoning:
                # සිතන ක්‍රියාවලිය පෙන්වීමට අවශ්‍ය නම් මෙතැනින් ලබාගත හැක
                pass 
                
            if content:
                full_response += content
                yield full_response
                
    except Exception as e:
        yield f"⚠️ වැරැද්දක් සිදු විය: {str(e)}"

# 🔱 INACHI UI Design
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", neutral_hue="zinc")) as demo:
    gr.Markdown("# 🔱 INACHI CORE V1.1")
    gr.Markdown(f"**Identity:** Inachi AI | **Developed by:** Inachi Team | **Architect:** MINZO-PRIME")
    
    chat = gr.ChatInterface(
        fn=predict,
        title="INACHI Intelligence Terminal",
        description="Official Inachi AI Core. Beyond Intelligence. Into the Void."
    )

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