import os
import io
import base64
import torch
from fastapi import FastAPI, Form, File, UploadFile, HTTPException
from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
from PIL import Image

app = FastAPI()

MODEL_ID = "stabilityai/sdxl-turbo"

print("🚀 Memuat Mamboro Engine - Dual Mode (Text & Image)...")

try:
    # 1. Load untuk Text-to-Image (Generate)
    pipe_text = AutoPipelineForText2Image.from_pretrained(
        MODEL_ID, torch_dtype=torch.float32, variant="fp16", low_cpu_mem_usage=True
    )
    pipe_text.to("cpu")
    pipe_text.enable_attention_slicing()

    # 2. Load untuk Image-to-Image (Edit) - Menggunakan model yang sama dari cache
    pipe_edit = AutoPipelineForImage2Image.from_pipe(pipe_text)
    
    print("✅ Server Ready!")
except Exception as e:
    print(f"❌ Gagal Load Model: {e}")

@app.get("/")
def home():
    return {"status": "Online", "mode": "Dual (Generate & Edit)"}

@app.post("/generate")
async def generate(prompt: str = Form(...)):
    try:
        # Gunakan pipe_text karena hanya ada input teks
        image = pipe_text(
            prompt=prompt, 
            num_inference_steps=2, 
            guidance_scale=0.0,
            width=512,
            height=512
        ).images[0]

        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return {"image": f"data:image/jpeg;base64,{img_str}"}
    except Exception as e:
        print(f"❌ Error Generate: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/edit")
async def edit(
    prompt: str = Form(...), 
    image_file: UploadFile = File(...), 
    strength: float = Form(0.5)
):
    try:
        # Baca semua jenis file gambar (PNG, JPG, WebP, dll)
        contents = await image_file.read()
        if not contents:
            raise HTTPException(status_code=400, detail="File gambar kosong")

        # Buka gambar dengan PIL
        raw_image = Image.open(io.BytesIO(contents))
        
        # KONVERSI KRITIS: Ubah ke RGB
        # Ini penting karena PNG sering punya 4 channel (RGBA), AI cuma mau 3 (RGB)
        init_image = raw_image.convert("RGB")
        
        # Resize tetap diperlukan agar RAM tidak meledak
        init_image = init_image.resize((512, 512))

        print(f"🎨 Editing berbagai format dengan Strength: {strength}")

        image = pipe_edit(
            prompt=prompt, 
            image=init_image, 
            strength=strength, 
            num_inference_steps=2, 
            guidance_scale=0.0
        ).images[0]

        # Kembalikan sebagai JPEG agar ringan saat dikirim balik ke HP
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG", quality=80)
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        return {"image": f"data:image/jpeg;base64,{img_str}"}
        
    except Exception as e:
        print(f"❌ Error Detail: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Gagal memproses gambar: {str(e)}")

@app.post("/edit_ref")
async def edit_with_ref(
    prompt: str = Form(...), 
    main_image: UploadFile = File(...), 
    ref_image: UploadFile = File(...),
    strength: float = Form(0.6)
):
    try:
        # Baca Gambar Utama
        main_bytes = await main_image.read()
        main_img = Image.open(io.BytesIO(main_bytes)).convert("RGB").resize((512, 512))

        # Baca Gambar Referensi
        ref_bytes = await ref_image.read()
        ref_img = Image.open(io.BytesIO(ref_bytes)).convert("RGB").resize((512, 512))

        # Gabungkan instruksi: Prompt + gaya dari Ref Image
        # Di SDXL Turbo sederhana, kita gunakan main_img sebagai init, 
        # dan ref_img bisa kita olah untuk memperkuat prompt (atau dikirim sebagai controlnet jika ada)
        
        image = pipe_edit(
            prompt=f"{prompt}, follow style of reference image", 
            image=main_img, 
            strength=strength, 
            num_inference_steps=2
        ).images[0]

        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return {"image": f"data:image/jpeg;base64,{img_str}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))