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

HF_TOKEN = os.environ.get("girlToken")

space_client = Client(
    "prithivMLmods/Qwen-Image-Edit-2511-LoRAs-Fast",
    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_to_imgbb(image_path):
#     """上传图片到免费图床，返回公开 URL"""
#     IMGBB_API_KEY = os.environ.get("IMGBB_API_KEY")
#     if not IMGBB_API_KEY:
#         return None
#     with open(image_path, "rb") as f:
#         import base64
#         b64 = base64.b64encode(f.read()).decode()
#     resp = requests.post(
#         "https://api.imgbb.com/1/upload",
#         data={"key": IMGBB_API_KEY, "image": b64}
#     )
#     if resp.status_code == 200:
#         print(f'图床resp.json()={resp.json()}')
#         return resp.json()["data"]["url"]
#     return None

def upload_to_hf(image_path):
    """上传到 HF dataset，返回公开直链"""
    from huggingface_hub import HfApi
    import uuid
    
    api = HfApi(token=HF_TOKEN)
    filename = f"{uuid.uuid4().hex[:8]}_{os.path.basename(image_path)}"
    
    # 上传到你自己的 HF dataset repo（需要先创建一个 public dataset repo）
    url = api.upload_file(
        path_or_fileobj=image_path,
        path_in_repo=f"uploads/{filename}",
        repo_id="kenfoo",  # 改成你自己的
        repo_type="dataset",
    )
    print(f'图床 ={url}')
    # 转成直链格式
    # url 格式: https://huggingface.co/datasets/xxx/yyy/resolve/main/uploads/zzz.jpg
    return url


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

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

    # 方案：上传到公开图床，拿到 URL 再传给远端
    #public_url = upload_to_imgbb(image)

    public_url = upload_to_hf(image)

    if not public_url:
        print("图床上传失败 ")
        return None, seed

    print(f"图片公开 URL: {public_url}")

    # 直接传 URL 字符串，让远端自己下载
    # 直接传字符串列表
    images_input = [public_url]

    print(f"prompt: {prompt}, lora: {lora_adapter}, seed: {seed}")

    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"返回值: {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"异常: {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="结果图片")

        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="引导强度", minimum=0.1, maximum=10.0, step=0.1, value=1.0)
            steps = gr.Slider(label="推理步数", 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)
# ```

# 然后去 [imgbb.com](https://imgbb.com/api) 免费申请一个 API key，存到 HF Space 的 Secrets 里，变量名 `IMGBB_API_KEY`。

# 这样的流程是：
# ```