import os
import gradio as gr
from huggingface_hub import InferenceClient

# ---------------------------------------------------------
# 1. API CLIENT SETUP (Using HF Serverless Infrastructure)
# ---------------------------------------------------------
hf_token = os.getenv("HF_TOKEN")

# Setting up clients for both required models
gemma_client = InferenceClient("google/gemma-2b-it", token=hf_token)
param_client = InferenceClient("bharatgenai/Param-1-2.9B-Instruct", token=hf_token)

# ---------------------------------------------------------
# 2. PURE PYTHON ADVANCED MOE ROUTER
# ---------------------------------------------------------
def advanced_router(prompt):
    prompt_lower = prompt.lower()
    
    # Target keywords for Indian Context, Agriculture, Medical, and Studies
    param_keywords = [
        "india", "indian", "hindi", "kheti", "farmer", "agriculture", "crop", 
        "medical", "doctor", "health", "bukhar", "fever", "medicine", "ayurved",
        "history", "study", "exam", "syllabus", "board", "pm", "constitution",
        "history of india", "kise kehte hain"
    ]
    
    for word in param_keywords:
        if word in prompt_lower:
            return "param"
    return "gemma"

# ---------------------------------------------------------
# 3. CORE LOGIC & IDENTITY OVERRIDE
# ---------------------------------------------------------
def dibakar_1_response(user_query, history):
    query_lower = user_query.lower()
    
    # --- STRICT PERSONAL IDENTITY CHECK ---
    if "kis na banaya" in query_lower or "kisne banaya" in query_lower or "who created you" in query_lower:
        return "Mujha Dibakar munshi na banaya hai yo akala kala banaya hai."
    
    if "dibakar ka bara ma" in query_lower or "about dibakar" in query_lower:
        return "Dibakar munshi ka garh indian ka west bangal ma purba barwaman ma samudragrarh ka natunpara ma hai yon ka pita ka name kartick munshi hai."

    # --- MOE ROUTING WORKFLOW ---
    selected_expert = advanced_router(user_query)
    
    try:
        if selected_expert == "gemma":
            system_prompt = f"System: You are Dibakar 1, an advanced AI system.\nUser: {user_query}\nAI:"
            response = gemma_client.text_generation(system_prompt, max_new_tokens=250, temperature=0.7)
            return f"[🟢 Routed to Gemma Expert]\n\n{response.strip()}"
            
        elif selected_expert == "param":
            system_prompt = f"System: You are Dibakar 1, an expert in Indian history, farming, and medical knowledge.\nUser: {user_query}\nAI:"
            response = param_client.text_generation(system_prompt, max_new_tokens=250, temperature=0.7)
            return f"[🟠 Routed to Param Expert]\n\n{response.strip()}"
            
    except Exception as e:
        return f"⚠️ Engine Routing Error, processing fallback... Details: {str(e)}"

# ---------------------------------------------------------
# 4. PREMIUM GRADIO WEB CHAT INTERFACE
# ---------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🤖 Dibakar 1 - MoE Engine v1.0
        ### Powered by Google Gemma & BharatGenAI Param | Created by Dibakar Munshi
        *This MoE intelligently routes technical/general queries to Gemma and Indian context/farming/medical queries to Param.*
        """
    )
    
    # Fixed for Gradio 6.0+: Removed incompatible button arguments and added modern placeholder
    chatbot = gr.ChatInterface(
        fn=dibakar_1_response,
        textbox=gr.Textbox(placeholder="Ask Dibakar 1 something...", container=False, scale=7),
        type="messages" # Enables new clean UI with default control buttons safely
    )

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