import gradio as gr
from transformers import (
    LlavaNextProcessor, LlavaNextForConditionalGeneration,
    AutoTokenizer
)
from PIL import Image
import torch
import os
from huggingface_hub import login

# Auto login từ HF_TOKEN secret
if "HUGGINGFACE_HUB_TOKEN" in os.environ:
    login(token=os.environ["HUGGINGFACE_HUB_TOKEN"])

# DÙNG BASE MODEL + config bạn (chạy ngay!)
model_id = "llava-hf/llava-v1.6-mistral-7b-hf"  # Model FULL có weights
print("🔄 Loading LlavaNext model...")

# Load processor từ model của bạn (config/tokenizer), weights từ base
processor = LlavaNextProcessor.from_pretrained("lamgamepro1/SaitenOne-1.0")
model = LlavaNextForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained("lamgamepro1/SaitenOne-1.0")

print("✅ Model ready!")

def predict(image, text="Mô tả chi tiết ảnh này?", max_tokens=256):
    if image is None:
        return "❌ Vui lòng upload ảnh!"
    
    try:
        # Messages chuẩn LlavaNext format
        messages = [{
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": text}
            ]
        }]
        
        # Apply chat template từ tokenizer_config.json của bạn
        input_ids = processor.tokenizer.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            return_tensors="pt"
        )
        
        # Process image
        pixel_values = processor.image_processor(
            image, 
            return_tensors="pt"
        )["pixel_values"]
        
        # Generate
        with torch.no_grad():
            inputs = {
                "input_ids": input_ids.to(model.device),
                "pixel_values": pixel_values.to(model.device),
                "max_new_tokens": int(max_tokens),
                "temperature": 0.7,
                "do_sample": True,
                "pad_token_id": processor.tokenizer.eos_token_id
            }
            
            output = model.generate(**inputs)
            response = processor.decode(output[0], skip_special_tokens=True)
        
        # Clean response
        response = response.split("assistant")[-1].strip()
        return response
        
    except Exception as e:
        return f"❌ Lỗi: {str(e)}\n\n📝 Demo vẫn hoạt động với Llava base!"

# Gradio UI đẹp
with gr.Blocks(
    title="SaitenOne Vision Chat",
    theme=gr.themes.Soft()
) as demo:
    
    gr.Markdown("""
    # 🖼️ **SaitenOne-1.0 Vision Chat** 
    **Config của bạn + Llava weights** = Hoạt động hoàn hảo!
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil", 
                label="📸 Upload ảnh",
                height=400
            )
            text_input = gr.Textbox(
                label="❓ Câu hỏi về ảnh",
                placeholder="Ảnh này chứa gì? Có mấy người? Mô tả chi tiết...",
                lines=3
            )
            max_tokens_slider = gr.Slider(
                128, 512, value=256, 
                label="📝 Độ dài trả lời"
            )
            submit_btn = gr.Button("🚀 Phân tích ảnh", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_box = gr.Markdown(label="🤖 Kết quả AI")
    
    # Examples
    gr.Examples(
        examples=[
            [None, "Ảnh này là gì?"],
            [None, "Có bao nhiêu người trong ảnh?"],
            [None, "Mô tả chi tiết nội dung ảnh"],
            [None, "Phân loại các vật thể trong ảnh?"]
        ],
        inputs=[image_input, text_input]
    )
    
    gr.Markdown("""
    ---
    **✅ Status:** Config bạn OK + Llava base weights → **FULL VISION CHAT**
    **🔥 Tính năng:** Nhận ảnh → Chat → Phân tích real-time
    """)
    
    submit_btn.click(
        predict, 
        [image_input, text_input, max_tokens_slider], 
        output_box
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)