#!/usr/bin/env python3
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

# Model configuration
MODEL_ID = "mx-llms/Lychee-GPT-9B"

# Get HF token from environment
HF_TOKEN = os.getenv("HF_TOKEN")

print("🚀 Loading Lychee-GPT-9B model...")
print("⏱️ This may take a few minutes...")

try:
    # Load tokenizer with token
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID, 
        trust_remote_code=True,
        token=HF_TOKEN
    )
    
    # Load model with optimizations
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        device_map="cpu",
        trust_remote_code=True,
        token=HF_TOKEN,
    )
    
    model.eval()
    print("✅ Model loaded successfully!")
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    print("\n⚠️ Make sure:")
    print("1. Model is not gated OR")
    print("2. HF_TOKEN environment variable is set")
    model = None
    tokenizer = None

def generate_text(prompt, max_length=256, temperature=0.7, top_p=0.9):
    """Generate text using Lychee-GPT-9B"""
    if model is None or tokenizer is None:
        return "❌ Model failed to load. Check that you have access to mx-llms/Lychee-GPT-9B model."
    
    try:
        # Tokenize input
        inputs = tokenizer(prompt, return_tensors="pt")
        
        # Generate with torch.no_grad for memory efficiency
        with torch.no_grad():
            output = model.generate(
                inputs["input_ids"],
                max_new_tokens=int(max_length),
                temperature=float(temperature),
                top_p=float(top_p),
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        # Decode output
        response = tokenizer.decode(output[0], skip_special_tokens=True)
        
        # Remove input prompt from response
        if prompt in response:
            response = response.replace(prompt, "", 1).strip()
        
        return response if response else "No response generated"
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Lychee-GPT-9B", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🎉 Lychee-GPT-9B Demo
        
        আপনার নিজস্ব LLM Model!
        
        ⏱️ **নোট:** 
        - প্রথম load: 2-3 মিনিট
        - Response time: 30-90 সেকেন্ড (CPU তে)
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="প্রশ্ন/Prompt",
                    placeholder="কিছু লিখুন...",
                    lines=4,
                    info="আপনার প্রশ্ন বা prompt দিন"
                )
                
                with gr.Row():
                    max_len = gr.Slider(
                        label="Max Length",
                        minimum=10,
                        maximum=512,
                        value=256,
                        step=10,
                        info="Response এর সর্বোচ্চ length"
                    )
                
                with gr.Row():
                    temp = gr.Slider(
                        label="Temperature",
                        minimum=0.0,
                        maximum=1.0,
                        value=0.7,
                        step=0.1,
                        info="কম = সুসংগত, বেশি = সৃজনশীল"
                    )
                    
                    top_p = gr.Slider(
                        label="Top P",
                        minimum=0.0,
                        maximum=1.0,
                        value=0.9,
                        step=0.05,
                        info="শব্দ নির্বাচন নিয়ন্ত্রণ"
                    )
                
                submit_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Response",
                    lines=10,
                    interactive=False,
                    info="Model এর response এখানে দেখা যাবে"
                )
        
        # Examples
        gr.Examples(
            examples=[
                ["বাংলা ভাষা সম্পর্কে বলুন"],
                ["পাইথন প্রোগ্রামিং কি?"],
                ["একটি সংক্ষিপ্ত গল্প বলুন"],
                ["কৃত্রিম বুদ্ধিমত্তা কি?"],
            ],
            inputs=prompt,
            label="উদাহরণ প্রশ্ন"
        )
        
        # Connect button click
        submit_btn.click(
            fn=generate_text,
            inputs=[prompt, max_len, temp, top_p],
            outputs=output,
            api_name="generate"
        )
        
        # Allow Enter key
        prompt.submit(
            fn=generate_text,
            inputs=[prompt, max_len, temp, top_p],
            outputs=output,
            api_name="generate"
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    # Launch with share=True for HF Spaces (they handle public .  URL)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,  # Important for Docker!
        show_error=True,
    )