import gradio as gr
import requests
import json
import time

# --- 配置資訊 ---
API_KEY = "d72ef10876b9403a9479c6e31f27a4be".strip()
WORKFLOW_ID = "1958842411961233410"
UPLOAD_URL = "https://www.runninghub.cn/task/openapi/upload"
CREATE_TASK_URL = "https://www.runninghub.cn/task/openapi/create"
GET_TASK_STATUS_URL = "https://www.runninghub.cn/task/openapi/status"
GET_TASK_OUTPUTS_URL = "https://www.runninghub.cn/task/openapi/outputs"
HEADERS = {"Host": "www.runninghub.cn"}

# 工作流節點 ID
LOAD_IMAGE_NODE_ID = "141"
POSITIVE_PROMPT_NODE_ID = "144"
SEED_NODE_ID = "132"

# --- 核心函式 ---
def process_image_workflow(image_path, seed=-1):
    positive_prompt = "a charming young woman with bright eyes and a subtle smile"

    if image_path is None:
        return "請先上傳圖片！", None

    # --- 步驟 1: 上傳圖片 ---
    yield "正在上傳圖片...", None
    try:
        with open(image_path, "rb") as f:
            files = {"file": f}
            data = {"apiKey": API_KEY, "fileType": "image"}
            upload_response = requests.post(UPLOAD_URL, headers=HEADERS, data=data, files=files)
            upload_response.raise_for_status()
            upload_data = upload_response.json()
            print("上傳回傳:", upload_data)

            if upload_data.get("code") != 0:
                return f"圖片上傳失敗！訊息: {upload_data.get('msg')}", None

            rh_image_path = upload_data["data"]["fileName"]
            yield f"圖片上傳成功！fileName: {rh_image_path}", None

    except requests.exceptions.RequestException as e:
        return f"上傳連線錯誤: {e}", None

    # --- 步驟 2: 建立工作流任務 ---
    yield "正在建立任務...", None
    node_info_list = [
        {"nodeId": LOAD_IMAGE_NODE_ID, "fieldName": "image", "fieldValue": rh_image_path},
        {"nodeId": POSITIVE_PROMPT_NODE_ID, "fieldName": "ShowText_0", "fieldValue": positive_prompt},
        {"nodeId": SEED_NODE_ID, "fieldName": "seed", "fieldValue": seed}
    ]
    payload = {"apiKey": API_KEY, "workflowId": WORKFLOW_ID, "nodeInfoList": node_info_list}

    try:
        create_response = requests.post(CREATE_TASK_URL, headers={"Content-Type": "application/json"}, data=json.dumps(payload))
        create_response.raise_for_status()
        create_data = create_response.json()
        print("建立任務回傳:", create_data)

        if create_data.get("code") != 0:
            return f"任務建立失敗！訊息: {create_data.get('msg')}", None

        task_id = create_data["data"]["taskId"]
        yield f"任務建立成功！TaskID: {task_id}", None

    except requests.exceptions.RequestException as e:
        return f"任務建立連線錯誤: {e}", None

    # --- 步驟 3: 輪詢任務狀態 ---
    max_retries = 40
    retry_interval = 5

    for i in range(max_retries):
        yield f"查詢任務狀態... (第 {i+1} 次)", None
        time.sleep(retry_interval)
        try:
            status_payload = {"apiKey": API_KEY, "taskId": task_id}
            status_response = requests.post(GET_TASK_STATUS_URL, headers={"Content-Type": "application/json"}, data=json.dumps(status_payload))
            status_response.raise_for_status()
            status_data = status_response.json()
            print("查詢狀態回傳:", status_data)

            if status_data.get("code") != 0:
                yield f"查詢任務失敗！訊息: {status_data.get('msg')}", None
                return f"查詢任務失敗！訊息: {status_data.get('msg')}", None

            current_status = status_data.get("data")
            if current_status == "SUCCESS":
                yield "任務狀態為成功，正在取得圖片結果...", None
                break  # 跳出輪詢迴圈
            elif current_status == "FAILED":
                yield f"任務失敗！", None
                return f"任務失敗！", None
        
        except requests.exceptions.RequestException as e:
            yield f"查詢連線錯誤: {e}", None
            return f"查詢連線錯誤: {e}", None
    else:
        yield "任務查詢超時，請使用 TaskID 手動查詢。", None
        return "任務查詢超時，請使用 TaskID 手動查詢。", None

    # --- 步驟 4: 取得任務結果 ---
    try:
        outputs_payload = {"apiKey": API_KEY, "taskId": task_id}
        outputs_response = requests.post(GET_TASK_OUTPUTS_URL, headers={"Content-Type": "application/json"}, data=json.dumps(outputs_payload))
        outputs_response.raise_for_status()
        outputs_data = outputs_response.json()
        print("取得結果回傳:", outputs_data)

        if outputs_data.get("code") != 0:
            return f"取得結果失敗！訊息: {outputs_data.get('msg')}", None

        image_url = outputs_data["data"][0]["fileUrl"]
        yield "任務完成！", image_url
        return "任務完成！", image_url

    except requests.exceptions.RequestException as e:
        return f"取得結果連線錯誤: {e}", None
    except Exception as e:
        return f"處理結果時發生錯誤: {e}", None

# --- Gradio 介面 ---
with gr.Blocks() as demo:
    gr.Markdown("# RunningHub 工作流圖生圖")
    gr.Markdown("上傳圖片，輸入種子值，啟動遠端工作流任務。")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="上傳圖片")
            seed_input = gr.Number(label="種子值 (Seed)", value=-1, precision=0)
            process_button = gr.Button("開始處理圖像")
        with gr.Column():
            status_output = gr.Textbox(label="任務狀態", interactive=False)
            image_output = gr.Image(label="生成圖片")

    process_button.click(
        fn=process_image_workflow,
        inputs=[image_input, seed_input],
        outputs=[status_output, image_output]
    )

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