import gradio as gr
import numpy as np
import random
import torch
import spaces
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
import math
import os

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Scheduler configuration for Lightning
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}

# Initialize scheduler
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

# Load model
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    scheduler=scheduler,
    torch_dtype=dtype
).to(device)

pipe.load_lora_weights(
    "2vXpSwA7/iroiro-lora",
    weight_name="qwen_lora/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16_dim1.safetensors"
)
pipe.fuse_lora(lora_scale=0.7)
pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")

# --- Constants ---
MAX_SEED = np.iinfo(np.int32).max
PROMPTS = {
    "front": "Move the camera to a front-facing position so the full body of the character is visible. The character stands with both arms extended slightly downward and close to the thighs, keeping the body evenly balanced on both sides. The legs are positioned symmetrically with a narrow stance. The background is plain white.",
    "back": "Move the camera to a back-facing position so the full body of the character is visible. Background is plain white.",
    "left": "Move the camera to a side view (profile) from the left so the full body of the character is visible. Background is plain white.",
    "right": "Move the camera to a side view (profile) from the right so the full body of the character is visible. Background is plain white."
}

# NEW: 出力解像度プリセット
RESOLUTIONS = {
    "1:4": (512, 2048),
    "1:3": (576, 1728),
    "nealy 9:16": (768, 1344),
    "nealy 2:3": (832, 1216),
    "3:4": (896, 1152),
}

def _append_prompt(base: str, extra: str) -> str:
    extra = (extra or "").strip()
    return (base if not extra else f"{base} {extra}").strip()

def generate_single_view(input_images, prompt, seed, num_inference_steps, true_guidance_scale):
    generator = torch.Generator(device=device).manual_seed(seed)
    result = pipe(
        image=input_images if input_images else None,
        prompt=prompt,
        negative_prompt=" ",
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=1,
    ).images
    return result[0]

def concat_images_horizontally(images, bg_color=(255, 255, 255)):
    images = [img.convert("RGB") for img in images if img is not None]
    if not images:
        return None
    h = max(img.height for img in images)
    resized = []
    for img in images:
        if img.height != h:
            w = int(img.width * (h / img.height))
            img = img.resize((w, h), Image.LANCZOS)
        resized.append(img)
    w_total = sum(img.width for img in resized)
    canvas = Image.new("RGB", (w_total, h), bg_color)
    x = 0
    for img in resized:
        canvas.paste(img, (x, 0))
        x += img.width
    return canvas

# NEW: リサイズユーティリティ
def resize_to_preset(img: Image.Image, preset_key: str) -> Image.Image:
    w, h = RESOLUTIONS[preset_key]
    return img.resize((w, h), Image.LANCZOS)

@spaces.GPU()
def generate_turnaround(
    image,
    extra_prompt="",
    preset_key="nealy 9:16",  # NEW: デフォルト
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    progress=gr.Progress(track_tqdm=True),
):
    """4視点＋横連結PNG生成（ユーザー追記プロンプト対応 & 出力解像度プリセット対応）"""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    if image is None:
        return None, None, None, None, None, seed, "エラー: 入力画像をアップロードしてください"

    if isinstance(image, Image.Image):
        input_image = image.convert("RGB")
    else:
        input_image = Image.open(image).convert("RGB")

    pil_images = [input_image]

    # 各プロンプト末尾に追記
    p_front = _append_prompt(PROMPTS["front"], extra_prompt)
    p_back  = _append_prompt(PROMPTS["back"],  extra_prompt)
    p_left  = _append_prompt(PROMPTS["left"],  extra_prompt)
    p_right = _append_prompt(PROMPTS["right"], extra_prompt)

    progress(0.25, desc="正面生成中...")
    front = generate_single_view(pil_images, p_front, seed, num_inference_steps, true_guidance_scale)

    progress(0.5, desc="背面生成中...")
    back = generate_single_view([front], p_back, seed+1, num_inference_steps, true_guidance_scale)

    progress(0.75, desc="左側面生成中...")
    left = generate_single_view([front], p_left, seed+2, num_inference_steps, true_guidance_scale)

    progress(1.0, desc="右側面生成中...")
    right = generate_single_view([front], p_right, seed+3, num_inference_steps, true_guidance_scale)

    # NEW: ここで指定プリセットにリサイズ
    front_r = resize_to_preset(front, preset_key)
    back_r  = resize_to_preset(back,  preset_key)
    left_r  = resize_to_preset(left,  preset_key)
    right_r = resize_to_preset(right, preset_key)

    # NEW: リサイズ後を連結（横：正面→右→背面→左）
    concat = concat_images_horizontally([front_r, right_r, back_r, left_r])

    return front_r, back_r, left_r, right_r, concat, seed, f"✅ {preset_key} にリサイズして4視点＋連結画像を生成しました"

# --- UI ---
css = """
#col-container {margin: 0 auto; max-width: 1400px;}
.image-container img {object-fit: contain !important; max-width: 100%; max-height: 100%;}
/* 追加: 注意ボックスのスタイル */
.notice {
  background: #fff5f5;
  border: 1px solid #fca5a5;
  color: #7f1d1d;
  padding: 12px 14px;
  border-radius: 10px;
  font-weight: 600;
  line-height: 1.5;
  margin-bottom: 10px;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# キャラクター4視点立ち絵自動生成")
    gr.Markdown("アップロードしたキャラクター画像から正面・背面・左右側面、さらに4枚連結のPNG画像を出力します。")

    with gr.Column(elem_id="col-container"):
        gr.HTML(
            "<div class='notice'>"
            "注意：他者が作成した画像のアップロードはご遠慮ください。"
            "他人の著作物・肖像権を侵害する恐れがあります。"
            "当アプリ作成者は、アップロード内容による権利侵害について一切の責任を負いません。"
            "</div>"
        )

        input_image = gr.Image(label="入力画像", type="pil", height=500)

        # 追記プロンプト欄
        extra_prompt = gr.Textbox(
            label="追加プロンプト（各視点プロンプトの末尾に追記）",
            placeholder="例: high detail, anime style, soft lighting, 4k, pastel colors",
            lines=2
        )

        # NEW: 出力解像度プリセットのプルダウン
        preset_dropdown = gr.Dropdown(
            label="出力解像度プリセット",
            choices=list(RESOLUTIONS.keys()),
            value="nealy 9:16"
        )

        run_button = gr.Button("🎨 生成開始", variant="primary")
        status_text = gr.Textbox(label="ステータス", interactive=False)

        with gr.Row():
            result_front = gr.Image(label="正面", type="pil", format="png", height=400, show_download_button=True)
            result_back = gr.Image(label="背面", type="pil", format="png", height=400, show_download_button=True)
        with gr.Row():
            result_left = gr.Image(label="左側面", type="pil", format="png", height=400, show_download_button=True)
            result_right = gr.Image(label="右側面", type="pil", format="png", height=400, show_download_button=True)

        # PNG連結出力
        result_concat = gr.Image(label="連結画像（正面→右→背面→左）", type="pil", format="png", height=400, show_download_button=True)

        with gr.Accordion("⚙️ 詳細設定", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="ランダムシード", value=True)
            true_guidance_scale = gr.Slider(label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
            num_inference_steps = gr.Slider(label="生成ステップ数", minimum=1, maximum=40, step=1, value=4)

    # NEW: クリック時に preset_dropdown を引数として渡す
    run_button.click(
        fn=generate_turnaround,
        inputs=[input_image, extra_prompt, preset_dropdown, seed, randomize_seed, true_guidance_scale, num_inference_steps],
        outputs=[result_front, result_back, result_left, result_right, result_concat, seed, status_text],
    )

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