import gradio as gr
import torch
from PIL import Image
import numpy as np
import random
import time

# บังคับใช้ CPU
device = torch.device('cpu')
print(f"Using device: {device}")

MAX_SEED = np.iinfo(np.int32).max

def load_edit_model():
    """โหลดโมเดลสำหรับ CPU"""
    try:
        from transformers import AutoProcessor
        from longcat_image.models import LongCatImageTransformer2DModel
        from longcat_image.pipelines import LongCatImageEditPipeline
        
        edit_model_id = 'meituan-longcat/LongCat-Image-Edit'
        print(f"🔄 Loading Image Edit model from {edit_model_id}...")
        print("⚠️ Using CPU mode - this will be slower but no GPU memory issues")
        
        # โหลดโมเดลแบบ float32 สำหรับ CPU
        edit_text_processor = AutoProcessor.from_pretrained(
            edit_model_id,
            subfolder='tokenizer'
        )

        edit_transformer = LongCatImageTransformer2DModel.from_pretrained(
            edit_model_id,
            subfolder='transformer',
            torch_dtype=torch.float32,  # ใช้ float32 สำหรับ CPU
            use_safetensors=True
        ).to(device)

        edit_pipe = LongCatImageEditPipeline.from_pretrained(
            edit_model_id,
            transformer=edit_transformer,
            text_processor=edit_text_processor,
            torch_dtype=torch.float32
        )
        edit_pipe.to(device)
        
        # ปิด offload เพื่อป้องกันปัญหา
        edit_pipe.unet = edit_pipe.unet.to(device)
        if hasattr(edit_pipe, 'vae'):
            edit_pipe.vae = edit_pipe.vae.to(device)
        
        print(f"✅ Image Edit model loaded successfully on {device}")
        return edit_pipe
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# ตัวแปร global สำหรับโมเดล
edit_pipe = None

def get_edit_pipe():
    """โหลดโมเดลเมื่อต้องการใช้จริง"""
    global edit_pipe
    if edit_pipe is None:
        edit_pipe = load_edit_model()
    return edit_pipe

def edit_image(
    input_image: Image.Image,
    prompt: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float = 4.0,
    num_inference_steps: int = 20,  # ลด steps สำหรับ CPU
    progress=gr.Progress()
):
    """Edit image based on text prompt - CPU version"""
    if input_image is None:
        raise gr.Error("กรุณาอัพโหลดภาพก่อน")
    if not prompt or prompt.strip() == "":
        raise gr.Error("กรุณาระบุคำสั่งการแก้ไขภาพ")
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    try:
        progress(0.1, desc="กำลังโหลดโมเดล...")
        pipe = get_edit_pipe()
        if pipe is None:
            raise gr.Error("ไม่สามารถโหลดโมเดลได้ กรุณาตรวจสอบการติดตั้ง")
        
        progress(0.3, desc="กำลังเตรียมภาพ...")
        if input_image.mode != 'RGB':
            input_image = input_image.convert('RGB')
        
        # Resize image เพื่อเพิ่มความเร็ว (CPU ต้องการภาพขนาดเล็ก)
        max_size = 512  # ขนาดเล็กสำหรับ CPU
        if max(input_image.size) > max_size:
            ratio = max_size / max(input_image.size)
            new_width = int(input_image.width * ratio)
            new_height = int(input_image.height * ratio)
            input_image = input_image.resize((new_width, new_height), Image.LANCZOS)
            print(f"ปรับขนาดภาพเป็น {new_width}x{new_height} สำหรับ CPU")
        
        progress(0.5, desc="กำลังสร้างภาพ... (โปรดรอ 2-5 นาที)")
        generator = torch.Generator(device).manual_seed(seed)
        
        # เริ่มจับเวลา
        start_time = time.time()
        
        with torch.inference_mode():
            output = pipe(
                input_image,
                prompt,
                negative_prompt="",
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                num_images_per_prompt=1,
                generator=generator,
                output_type="pil"
            )
        
        # คำนวณเวลา
        elapsed_time = time.time() - start_time
        minutes = int(elapsed_time // 60)
        seconds = int(elapsed_time % 60)
        
        progress(1.0, desc=f"เสร็จสิ้น! ใช้เวลา {minutes} นาที {seconds} วินาที")
        
        return output.images[0], seed, f"ใช้เวลา: {minutes} นาที {seconds} วินาที"
    
    except Exception as e:
        print(f"Error: {str(e)}")
        raise gr.Error(f"เกิดข้อผิดพลาด: {str(e)}")

css = """
#col-container {
    margin: 0 auto;
    max-width: 900px;
}
.generate-btn {
    background: linear-gradient(90deg, #4CAF50 0%, #2E7D32 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2) !important;
}
.cpu-warning {
    background-color: #e3f2fd;
    border: 2px solid #2196F3;
    border-radius: 10px;
    padding: 15px;
    margin: 15px 0;
    color: #1565c0;
}
.time-info {
    background-color: #e8f5e9;
    border: 1px solid #4CAF50;
    border-radius: 8px;
    padding: 10px;
    margin: 10px 0;
    color: #2E7D32;
}
.image-box {
    border: 2px dashed #ddd;
    border-radius: 10px;
    padding: 10px;
    background: #fafafa;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as app:
    gr.HTML("""
    <center>
        <h1 style="background: linear-gradient(90deg, #4CAF50, #2E7D32); -webkit-background-clip: text; -webkit-text-fill-color: transparent; margin-bottom: 10px;">
            🖼️ LongCat-Image Editor (CPU Version)
        </h1>
        <p style="color: #666; margin-bottom: 20px;">
            แก้ไขภาพด้วยปัญญาประดิษฐ์ - รุ่นประหยัดหน่วยความจำ (ใช้ CPU เท่านั้น)
        </p>
    </center>
    """)
    
    with gr.Column(elem_id="col-container"):
        # Warning about CPU mode
        gr.HTML("""
        <div class="cpu-warning">
            ⚠️ <strong>โหมด CPU:</strong> กำลังใช้ CPU เท่านั้นในการประมวลผล<br>
            ⏱️ <strong>เวลาในการประมวลผล:</strong> ประมาณ 2-5 นาทีต่อภาพ<br>
            📏 <strong>ขนาดภาพแนะนำ:</strong> ไม่เกิน 512x512 พิกเซล<br>
            💡 <strong>ข้อดี:</strong> ไม่ต้องใช้การ์ดจอ ไม่มีปัญหาเรื่องหน่วยความจำ
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input Image
                input_image = gr.Image(
                    label="📤 ภาพต้นฉบับ",
                    type="pil",
                    elem_classes=["image-box"],
                    height=300
                )
                
                # Edit Prompt
                prompt = gr.Textbox(
                    label="✏️ คำสั่งแก้ไขภาพ",
                    placeholder="ตัวอย่าง: 'ทำให้แมวใส่หมวก', 'เปลี่ยนพื้นหลังเป็นทะเล', 'เพิ่มแว่นตาดำ'...",
                    lines=2
                )
                
                # Quick edit suggestions (Thai version)
                with gr.Row():
                    quick_edits = [
                        ("เพิ่มแว่นตา", "เพิ่มแว่นตากันแดดให้บุคคล"),
                        ("เปลี่ยนฤดู", "เปลี่ยนเป็นหน้าหนาวที่มีหิมะ"),
                        ("กลางคืน", "ทำให้เป็นเวลากลางคืนมีดาว"),
                        ("เพิ่มสัตว์", "เพิ่มสุนัขน่ารักในภาพ"),
                    ]
                    
                    for btn_text, edit_text in quick_edits:
                        btn = gr.Button(
                            btn_text,
                            size="sm",
                            variant="secondary",
                            scale=0
                        )
                        
                        def make_click_fn(text):
                            return lambda: text
                        
                        btn.click(
                            fn=make_click_fn(edit_text),
                            outputs=prompt
                        )
                
                with gr.Accordion("⚙️ การตั้งค่าขั้นสูง", open=False):
                    with gr.Row():
                        seed = gr.Slider(
                            label="Seed",
                            value=-1,
                            minimum=-1,
                            maximum=MAX_SEED,
                            step=1
                        )
                        randomize_seed = gr.Checkbox(
                            label="สุ่ม Seed",
                            value=True
                        )
                    
                    with gr.Row():
                        steps = gr.Slider(
                            label="จำนวนขั้นตอน",
                            value=20,
                            minimum=10,
                            maximum=30,
                            step=1,
                            info="น้อย = เร็ว, มาก = คุณภาพดีขึ้น"
                        )
                        guidance_scale = gr.Slider(
                            label="Guidance Scale",
                            value=4.0,
                            minimum=1.0,
                            maximum=7.0,
                            step=0.5,
                            info="ความเข้มงวดในการทำตามคำสั่ง"
                        )
                
                # Generate Button
                generate_btn = gr.Button(
                    "🎨 แก้ไขภาพ",
                    variant='primary',
                    elem_classes=["generate-btn"],
                    scale=1
                )
            
            with gr.Column(scale=1):
                # Output Image
                output_image = gr.Image(
                    label="🖼️ ภาพที่แก้ไขแล้ว",
                    type="pil",
                    elem_classes=["image-box"],
                    height=300
                )
                
                # Time and seed info
                time_info = gr.Markdown(
                    "",
                    elem_classes=["time-info"]
                )
                
                # Seed output
                with gr.Row():
                    seed_output = gr.Number(
                        label="Seed ที่ใช้",
                        visible=True
                    )
                    copy_seed_btn = gr.Button(
                        "📋 คัดลอก Seed",
                        size="sm",
                        variant="secondary"
                    )
                    
                    def copy_seed(seed_val):
                        return seed_val
                    
                    copy_seed_btn.click(
                        fn=copy_seed,
                        inputs=seed_output,
                        outputs=seed
                    )
        
        # Tips section (Thai)
        gr.HTML("""
        <div style="margin-top: 20px; padding: 15px; background: #f8f9fa; border-radius: 8px;">
            <h4 style="margin-top: 0;">💡 คำแนะนำสำหรับผลลัพธ์ที่ดีที่สุด:</h4>
            <ul style="margin-bottom: 0;">
                <li>ใช้คำสั่งที่ชัดเจนและเป็นรูปธรรม</li>
                <li>เริ่มจาก 20 ขั้นตอน (steps) และปรับเพิ่มถ้าจำเป็น</li>
                <li>ใช้ Guidance Scale ระหว่าง 3.5-5.0</li>
                <li>อัพโหลดภาพไม่เกิน 512x512 พิกเซลสำหรับความเร็วสูงสุด</li>
                <li>ถ้าต้องการภาพใหญ่ ให้ใช้การปรับขนาดภายหลัง</li>
            </ul>
        </div>
        """)
        
        # Example images
        gr.Examples(
            examples=[
                ["ตัวอย่างภาพแมว", "ทำให้แมวใส่หมวกใบใหญ่สีแดง"],
                ["ตัวอย่างภาพทิวทัศน์", "เพิ่มรุ้งกินน้ำบนฟ้า"],
                ["ตัวอย่างภาพบุคคล", "เปลี่ยนทรงผมให้เป็นลอน"],
            ],
            inputs=[input_image, prompt],
            label="ตัวอย่างคำสั่ง"
        )
    
    # Connect the button
    generate_btn.click(
        fn=edit_image,
        inputs=[
            input_image,
            prompt,
            seed,
            randomize_seed,
            guidance_scale,
            steps
        ],
        outputs=[output_image, seed_output, time_info]
    )
    
    # Also allow prompt submission with Enter
    prompt.submit(
        fn=edit_image,
        inputs=[
            input_image,
            prompt,
            seed,
            randomize_seed,
            guidance_scale,
            steps
        ],
        outputs=[output_image, seed_output, time_info]
    )

# ล้างหน่วยความจำ
import gc
gc.collect()

if __name__ == "__main__":
    # โหลดโมเดลล่วงหน้า
    print("กำลังโหลดโมเดลล่วงหน้า...")
    _ = get_edit_pipe()
    
    app.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        debug=False,
        enable_queue=True,
        max_threads=1
    )