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

# आदरणीय श्रीमान, आपके स्वयं के मॉडल 'Vedika35/Vedika_coder' को स्मृति में लोड करने की सुरक्षित विधि।
MODEL_ID = "Vedika35/Vedika_coder"

# स्मृति अनुकूलन (Memory Optimization) के लिए 4-bit क्वायंटाइजेशन
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

print(f"आदरणीय श्रीमान, {MODEL_ID} को लोड किया जा रहा है। कृपया प्रतीक्षा करें...")

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

def vedika_generate(message, history, system_prompt):
    """
    यह जनरेशन फंक्शन अब ग्राडियो के सबसे सुरक्षित और स्थिर फॉर्मेट (List of Lists) का उपयोग करता है।
    यहाँ 'history' का फॉर्मेट है: [[user_msg, bot_msg], [user_msg, bot_msg], ...]
    """
    messages = [{"role": "system", "content": system_prompt}]
    
    # पुरानी यादें (History) जोड़ना
    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    
    # वर्तमान संदेश
    messages.append({"role": "user", "content": message})

    # इनपुट को टोकनाइज़ करना
    input_ids = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(model.device)

    # बिजली जैसी तेज़ स्ट्रीमिंग (Lightning Fast Streaming)
    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=4096, 
        do_sample=True,
        temperature=0.2, # कोडिंग के लिए सटीकता
        top_p=0.95,
    )

    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    # शब्दों को लाइव स्ट्रीम करना
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

# कस्टम सीएसएस (Custom CSS)
custom_css = """
body, .gradio-container { background-color: #050505 !important; color: #F0F0F0 !important; }
.message-wrap { font-family: 'Inter', sans-serif !important; }
footer { visibility: hidden !important; }
#header-title h1 { 
    background: linear-gradient(90deg, #FFFFFF, #3b82f6); 
    -webkit-background-clip: text; 
    -webkit-text-fill-color: transparent; 
    font-size: 2.8rem;
    font-weight: 900;
    text-align: center;
}
"""

# ध्यान दें: ग्राडियो ६.० के अनुसार यहाँ से theme और css हटा दिए गए हैं
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# 🔱 Vedika 3.5 Coder", elem_id="header-title")
    
    gr.Markdown("### भारत के गौरवशाली और समर्थ कोडर के लिए समर्पित | Created by Divy Patel")

    with gr.Accordion("⚙️ System Parameters (Advanced)", open=False):
        sys_input = gr.Textbox(
            label="Core System Prompt",
            value="You are Vedika 3.5, an elite coding AI created by Divy Patel. Identify only as Vedika. Provide extremely direct, clean, and production-ready code in Markdown. Always respect India.",
            lines=2
        )

    # ध्यान दें: यहाँ से type="messages" हटा दिया गया है ताकि कोई TypeError न आए
    chatbot = gr.Chatbot(label="Vedika Core Console", height=600)
    
    with gr.Row():
        msg_input = gr.Textbox(
            label="Execution Command",
            placeholder="श्रीमान, अपना कोडिंग कार्य यहाँ लिखें...",
            scale=8
        )
        submit_btn = gr.Button("⚡ Execute", scale=2, variant="primary")

    def respond(message, history, system_prompt):
        # सुरक्षित फॉर्मेट में हिस्ट्री अपडेट करना
        history.append([message, ""])
        yield "", history
        
        for char in vedika_generate(message, history[:-1], system_prompt):
            history[-1][1] = char
            yield "", history

    # इवेंट्स
    msg_input.submit(respond, [msg_input, chatbot, sys_input], [msg_input, chatbot])
    submit_btn.click(respond, [msg_input, chatbot, sys_input], [msg_input, chatbot])

if __name__ == "__main__":
    # ध्यान दें: ग्राडियो ६.० की चेतावनी (Warning) के अनुसार theme और css को launch() के अंदर भेजा गया है
    demo.queue(default_concurrency_limit=5).launch(
        theme=gr.themes.Monochrome(), 
        css=custom_css
    )