import os
# Thiết lập môi trường trước khi import torch
os.environ.setdefault("device", "cuda:0")

import boogu.utils.import_utils as _import_utils
_import_utils._triton_available = False

import io
import json
import torch
import spaces
from PIL import Image
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse
import gradio as gr

from boogu.pipelines.boogu.pipeline_boogu import BooguImagePipeline

# --- KHỞI TẠO MODEL TURBO DUY NHẤT ---
TURBO_ID = "Boogu/Boogu-Image-0.1-Edit-Turbo"
AOTI_REPO = "multimodalart/Boogu-Image-0.1-Edit-aoti"

print("Đang tải model Boogu Image Turbo...")
# Load trực tiếp pipeline từ Turbo_ID, tái sử dụng các thành phần để tối ưu RAM
turbo_pipe = BooguImagePipeline.from_pretrained(
    TURBO_ID, 
    torch_dtype=torch.bfloat16, 
    trust_remote_code=True
)
turbo_pipe.to("cuda")

# Vá lỗi AoTI để tăng tốc độ nếu môi trường Hugging Face hỗ trợ
try:
    from pathlib import Path
    from huggingface_hub import snapshot_download
    from spaces.zero.torch.aoti import aoti_load_from_module_dir
    _block_dir = Path(snapshot_download(AOTI_REPO)) / "BooguImageTransformerBlock"
    if (_block_dir / "package.pt2").exists():
        aoti_load_from_module_dir(turbo_pipe.transformer.single_stream_layers, _block_dir)
    print("AoTI đã được cấu hình thành công cho Turbo!")
except Exception as exc:
    print(f"Không load được AoTI, chạy mặc định (Eager mode): {exc}")

MAX_SEED = 2**31 - 1
RESOLUTIONS = {
    "1K": {"pixels": 1024 * 1024, "side": 2048},
    "2K": {"pixels": 2048 * 2048, "side": 4096},
}

# --- HÀM XỬ LÝ CHÍNH TRÊN ZEROGPU ---
def _duration(image, instruction, resolution, num_inference_steps, *args, **kwargs):
    # Pinned cho model Turbo, thời gian chạy rất ngắn
    base = int(num_inference_steps) * 4 + 40
    return base * 2 if resolution == "2K" else base

@spaces.GPU(duration=_duration)
def generate_turbo_core(input_image_pil, instruction, resolution, num_inference_steps, seed):
    res = RESOLUTIONS[resolution]
    generator = torch.Generator("cuda").manual_seed(seed)

    # Model Turbo ép cố định guidance scale về 1.0 (CFG off) theo đặc tả của tác giả
    text_guidance_scale = 1.0
    image_guidance_scale = 1.0

    if input_image_pil is None:
        # Text to Image
        size = 1024 if resolution == "1K" else 2048
        result = turbo_pipe(
            instruction=[instruction.strip()],
            negative_instruction="",
            height=size,
            width=size,
            max_input_image_pixels=res["pixels"],
            max_input_image_side_length=res["side"],
            num_inference_steps=int(num_inference_steps),
            text_guidance_scale=float(text_guidance_scale),
            generator=generator,
            device="cuda",
        ).images[0]
    else:
        # Image to Image / Edit
        temp_path = "temp_input_turbo.jpg"
        input_image_pil.save(temp_path)
        
        result = turbo_pipe(
            instruction=[instruction.strip()],
            input_image_paths=[[temp_path]],
            input_images=[[input_image_pil]],
            negative_instruction="",
            height=None,
            width=None,
            max_input_image_pixels=res["pixels"],
            max_input_image_side_length=res["side"],
            align_res=True,
            num_inference_steps=int(num_inference_steps),
            text_guidance_scale=float(text_guidance_scale),
            image_guidance_scale=float(image_guidance_scale),
            generator=generator,
            device="cuda",
        ).images[0]
        
        if os.path.exists(temp_path):
            os.remove(temp_path)

    return result

# --- DỰNG FASTAPI ROUTER ---
app = FastAPI(title="Boogu Image Turbo API")

@app.post("/api/generate")
async def api_generate(
    instruction: str = Form(...),
    image: UploadFile = File(None),
    resolution: str = Form("1K"),       # "1K" hoặc "2K"
    num_inference_steps: int = Form(4),  # Mặc định lý tưởng cho Turbo là 4 steps
    seed: int = Form(0)
):
    if not instruction.strip():
        raise HTTPException(status_code=400, detail="Prompt không được để trống")
    
    input_image_pil = None
    if image:
        try:
            image_bytes = await image.read()
            input_image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        except Exception:
            raise HTTPException(status_code=400, detail="File ảnh gửi lên không hợp lệ")

    if seed == 0:
        seed = int(torch.randint(0, MAX_SEED, (1,)).item())

    try:
        # Gọi hàm core xử lý trên ZeroGPU
        output_pil = generate_turbo_core(
            input_image_pil, instruction, resolution, num_inference_steps, seed
        )
        
        # Đóng gói ảnh thành định dạng WEBP trả về stream binary trực tiếp
        img_io = io.BytesIO()
        output_pil.save(img_io, format="WEBP", quality=95)
        img_io.seek(0)
        
        return StreamingResponse(img_io, media_type="image/webp")
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi xử lý GPU: {str(e)}")

# --- PHẦN GRADIO ĐỂ GIỮ CHỖ CHẠY TRÊN HUGGING FACE SPACES ZEROGPU ---
with gr.Blocks() as demo:
    gr.Markdown("# 🍊 Boogu Image TURBO API is Active!")
    gr.Markdown("Gửi request POST tới endpoint: `https://<your-space-url>/api/generate` để tạo hoặc edit ảnh siêu tốc.")

# Nhúng FastAPI app vào Gradio Server
gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    demo.queue().launch()