import gradio as gr
from gradio_client import Client
import random
import os
import requests

HF_TOKEN = os.environ.get("girlToken")
TARGET_SPACE_URL = "https://prithivmlmods-qwen-image-edit-2511-loras-fast.hf.space"

space_client = Client(
    "prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast",
    hf_token=HF_TOKEN
)

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 upload_file_to_space(local_path):
    """手动上传文件到目标 Space，返回远端路径"""
    upload_url = f"{TARGET_SPACE_URL}/upload"

    headers = {}
    if HF_TOKEN:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"

    mime_type = "image/jpeg"
    if local_path.lower().endswith(".png"):
        mime_type = "image/png"
    elif local_path.lower().endswith(".webp"):
        mime_type = "image/webp"

    with open(local_path, "rb") as f:
        response = requests.post(
            upload_url,
            headers=headers,
            files={"files": (os.path.basename(local_path), f, mime_type)},
        )

    print(f"上传状态码: {response.status_code}")
    print(f"上传响应: {response.text}")

    if response.status_code == 200:
        result = response.json()
        remote_path = result[0] if isinstance(result, list) else result
        return remote_path
    else:
        raise Exception(f"上传失败: {response.status_code} {response.text}")


def infer(
    image,
    prompt,
    lora_adapter,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True),
):
    if image is None:
        print("未上传图片")
        return None, seed

    if not os.path.exists(image):
        print(f"图片路径不存在: {image}")
        return None, seed

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

    try:
        remote_path = upload_file_to_space(image)
        print(f"远端路径: {remote_path}")
    except Exception as e:
        print(f"上传图片失败: {e}")
        return None, seed

    mime_type = "image/jpeg"
    if image.lower().endswith(".png"):
        mime_type = "image/png"
    elif image.lower().endswith(".webp"):
        mime_type = "image/webp"

    images_input = [{
        "image": {
            "path": remote_path,
            "url": f"{TARGET_SPACE_URL}/file={remote_path}",
            "size": os.path.getsize(image),
            "orig_name": os.path.basename(image),
            "mime_type": mime_type,
            "is_stream": False,
            "meta": {}
        },
        "caption": None
    }]

    print("[调用API] 输入参数：")
    print(f"  remote_path: {remote_path}")
    print(f"  prompt: {prompt}")
    print(f"  lora_adapter: {lora_adapter}")
    print(f"  seed: {seed}")
    print(f"  guidance_scale: {guidance_scale}")
    print(f"  steps: {steps}")

    try:
        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] 返回值: {result}")

        image_info, seed_used = result

        if isinstance(image_info, dict):
            img_out = image_info.get("path") or image_info.get("url")
        else:
            img_out = image_info

        return img_out, int(seed_used)

    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"[调用API] 异常: {e}")
        return None, seed


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

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

        image = gr.Image(
            label="上传图片",
            sources=["upload"],
            type="filepath",
        )

        prompt = gr.Text(
            label="编辑描述（Prompt）",
            placeholder="请输入图片编辑描述...",
        )

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

        run_button = gr.Button("执行编辑", 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,
            )

    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__":
    demo.launch(ssr_mode=False, share=True)