import gradio as gr
from gradio_client import Client, handle_file
import random
import os

# API client for the external Space
space_client = Client("prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast")

LORA_STYLES = [
    'Multiple-Angles',
    'Photo-to-Anime',
    'Anime-V2',
    'Light-Migration',
    'Upscaler',
    'Style-Transfer',
    'Manga-Tone',
    'Anything2Real',
    'Fal-Multiple-Angles',
    'Polaroid-Photo',
    'Unblur-Anything',
    'Midnight-Noir-Eyes-Spotlight',
    'Hyper-Realistic-Portrait',
    'Ultra-Realistic-Portrait',
    'Pixar-Inspired-3D',
    'Noir-Comic-Book',
    'Any-light',
    'Studio-DeLight',
    'Cinematic-FlatLog',
]
MAX_SEED = 2**31 - 1

def encode_image_to_gallery_dict(image_path):
    if not image_path or not isinstance(image_path, str):
        return None
    try:
        return handle_file(image_path)
    except Exception as e:
        print(f"无法读取图片: {image_path}: {e}")
        return None

def infer(
    image,
    prompt,
    lora_adapter,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True),
):
    # Process input image(s) as a list of dict, each of form {"image": {"data": ...}}
    images_input = []
    if image is not None:
        if isinstance(image, list):  # Gradio Gallery
            for im in image:
                img_obj = encode_image_to_gallery_dict(im)
                if img_obj:
                    images_input.append(img_obj)
                else:
                    print(f"警告: 路径无效或无法读取: {im}")
        else:
            img_obj = encode_image_to_gallery_dict(image)
            if img_obj:
                images_input.append(img_obj)
            else:
                print(f"警告: 路径无效或无法读取: {image}")

    if len(images_input) == 0:
        print("未检测到有效图片，未能上传图片。")
        return None, seed

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    print("[调用API] space_client.predict 输入参数：")
    print(f"  images count: {len(images_input)}")
    print(f"  prompt: {prompt}")
    print(f"  lora_adapter: {lora_adapter}")
    print(f"  seed: {seed}")
    print(f"  randomize_seed: {randomize_seed}")
    print(f"  guidance_scale: {guidance_scale}")
    print(f"  steps: {steps}")

    try:
        # Gradio Space expects list of {"image": {"data": <b64 str>}} elements (see AppError in prompt)
        result = space_client.predict(
            images=images_input,
            prompt=prompt,
            lora_adapter=lora_adapter,
            seed=float(seed),
            randomize_seed=bool(randomize_seed),
            guidance_scale=float(guidance_scale),
            steps=float(steps),
            api_name="/infer",
        )
        print(f"[调用API] space_client.predict 返回值: {result}")
        # result可能不是元组形式，需增加健壮性
        if isinstance(result, dict):
            # 新API只返回dict
            image_info = result
            seed_used = result.get("seed", seed)
        elif isinstance(result, (tuple, list)) and len(result) == 2:
            image_info, seed_used = result
        else:
            print(f"[错误] space_client.predict 返回类型未知: {type(result)}，内容: {result}")
            return None, seed

        # 检查image_info是否包含url/path
        if isinstance(image_info, dict):
            img_url = image_info.get("url") or image_info.get("path") or None
            # 获取base64图片（如果有）
            if img_url is None and "data" in image_info:
                # 返回base64图片（data格式）
                return f"data:image/png;base64,{image_info['data']}", seed_used
            return img_url, seed_used
        else:
            print(f"[错误] image_info 格式异常: {image_info}")
            return None, seed
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"[调用API] 调用接口异常: {e}")
        if "Could not process uploaded images" in str(e):
            print(
                "\n[错误] 上传图片处理失败。请确保图片文件有效且未损坏。\n"
                "可以尝试重新选择图片或修改图片格式。"
            )
        else:
            print("\n[错误] 调用推理服务失败，请稍后重试或联系开发者。")
        return None, seed

examples = [
    [None, "Astronaut in jungle, anime style", "Photo-to-Anime", 0, True, 1.0, 4],
    [None, "A delicious ceviche cheesecake slice", "Style-Transfer", 0, True, 1.0, 4],
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # 图像编辑 API Demo (基于 prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast)")

        with gr.Row():
            image = gr.Image(
                label="上传图片",
                sources=["upload"],
                type="filepath",
                elem_id="input-image"
            )
        with gr.Row():
            prompt = gr.Text(
                label="编辑描述（Prompt）",
                placeholder="请输入图片编辑描述...",
            )

        with gr.Row():
            lora_adapter = gr.Dropdown(
                label="编辑风格（Style）",
                choices=LORA_STYLES,
                value="Photo-to-Anime"
            )

        run_button = gr.Button("执行编辑", scale=1, variant="primary")

        result = gr.Image(label="结果图片", show_label=True)

        with gr.Accordion("高级设置", open=False):
            seed = gr.Slider(
                label="随机种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="随机种子", value=True)

            guidance_scale = gr.Slider(
                label="引导强度(Guidance Scale)",
                minimum=0.1,
                maximum=10.0,
                step=0.1,
                value=1.0,
            )

            steps = gr.Slider(
                label="推理步数(Steps)",
                minimum=1,
                maximum=50,
                step=1,
                value=4,
            )

        # Only show example text/inputs, but avoid file path errors (set image to None)
        gr.Examples(
            examples=examples,
            inputs=[image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
            label="示例",
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            image,
            prompt,
            lora_adapter,
            seed,
            randomize_seed,
            guidance_scale,
            steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    # See: https://prithivmlmods-qwen-image-edit-2511-loras-fast.hf.space
    # To use SSR, remove 'ssr_mode=False' (can help with speed in some settings)
    # To create a public link, set share=True
    demo.launch(css=css, ssr_mode=True, share=True)