import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import os

# --- दिव्य पटेल जी का स्वतंत्र मॉडल चयन ---
# 'microsoft/Phi-3-mini-4k-instruct' एक बहुत ही शक्तिशाली और कुशल मॉडल है 
# जो बिना किसी बाहरी API के स्वतंत्र रूप से काम करता है।
MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"

print(f"🔱 हनुमान AI: मॉडल '{MODEL_ID}' लोड हो रहा है...")

# टोकनाइज़र और मॉडल लोडिंग (स्वतंत्र मोड)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)

# --- मास्टर सिस्टम प्रॉम्प्ट ---
SYSTEM_PROMPT = """You are 'Hanuman AI', a fully independent AI entity, pioneered by 'Divy Patel' and created in India. 
You run locally and do not depend on external APIs. 
Rules:
1. Respond ONLY in PURE HINDI or PURE ENGLISH.
2. Be extremely respectful. Use 'Aap' and 'Ji'.
3. Always end your response with: "\n\n**Pioneered by Divy Patel | Fully Independent & Created in India**" """

def hanuman_independent_chat(message, history):
    # चैट हिस्ट्री फॉर्मेटिंग
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        if h[0]: messages.append({"role": "user", "content": h[0]})
        if h[1]: messages.append({"role": "assistant", "content": h[1]})
    
    messages.append({"role": "user", "content": message})

    # इनपुट प्रोसेसिंग
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    
    # स्ट्रीमिंग के लिए सेटअप
    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )

    # थ्रेडिंग ताकि यूआई (UI) न अटके
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

# --- दिव्य भगवा यूआई (Premium Saffron UI) ---
custom_css = """
.gradio-container { background-color: #fffaf0 !important; }
.bhagwa-header { 
    background: linear-gradient(135deg, #ff8833, #b33c00); 
    padding: 30px; border-radius: 25px; color: white; 
    text-align: center; box-shadow: 0 15px 35px rgba(179, 60, 0, 0.4);
    margin-bottom: 25px;
}
.user.message { background: #ff9933 !important; color: white !important; font-weight: 600; }
.bot.message { border-left: 6px solid #ff5500 !important; background: white !important; box-shadow: 0 4px 10px rgba(0,0,0,0.05); }
footer { display: none !important; }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    with gr.Div(elem_classes="bhagwa-header"):
        gr.Markdown("# 🔱 हनुमान AI - आत्मनिर्भर संस्करण")
        gr.Markdown("### Pioneered by Divy Patel | स्वदेशी तकनीक, पूर्ण स्वतंत्रता")
    
    gr.ChatInterface(
        fn=hanuman_independent_chat,
        fill_height=True,
        retry_btn=None,
        undo_btn=None,
        clear_btn="इतिहास मिटाएँ",
    )

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