import gradio as gr
import requests
import base64
import time

# --- 核心配置 ---
API_KEY = "zz2026022"
# 根据截图，Base URL 是 http://<IP>:<端口>/flow/v1
BASE_URL = "http://154.40.59.124:8000/flow/v1" 

def encode_image(image_path):
    """将本地图片路径转换为 Base64 字符串"""
    if not image_path:
        return None
    with open(image_path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode('utf-8')
        return f"data:image/png;base64,{encoded}"

def generate_video_task(prompt, start_img_path, end_img_path):
    if not start_img_path:
        return None, "错误：请至少上传起始帧图片。"

    headers = {
        "Authorization": f"Bearer {API_KEY}", #
        "Content-Type": "application/json"
    }

    # 1. 转换图片为 Base64
    start_b64 = encode_image(start_img_path)
    end_b64 = encode_image(end_img_path) if end_img_path else None

    # 2. 构造符合 /v1/videos 接口要求的 payload
    input_refs = [start_b64]
    if end_b64:
        input_refs.append(end_b64) # 传入 [起始帧, 结束帧] 列表

    payload = {
        "prompt": prompt if prompt else "Generate transition video", #
        "model": "veo_3_1", # 使用文档指定的模型名称
        "input_reference": input_refs # 接口参数名为 input_reference
    }

    try:
        # 3. 提交生成任务 (POST /v1/videos)
        response = requests.post(f"{BASE_URL}/videos", json=payload, headers=headers)
        res_data = response.json()

        if response.status_code != 200:
            return None, f"任务提交失败: {res_data}"

        video_id = res_data.get("id")
        if not video_id:
            return None, "未获取到任务 ID，请检查接口返回格式。"

        # 4. 轮询视频生成状态 (GET /v1/videos/{video_id})
        status_url = f"{BASE_URL}/videos/{video_id}"
        for _ in range(60):  # 最多等待 300 秒
            time.sleep(5)
            status_res = requests.get(status_url, headers=headers)
            status_data = status_res.json()
            
            # 这里的状态判断逻辑需根据您的 API 实际返回字段名进行适配（如 status 或 state）
            status = status_data.get("status") 
            if status == "completed" or status == "succeeded":
                return status_data.get("url"), f"生成成功！ID: {video_id}"
            elif status == "failed":
                return None, f"视频生成失败: {status_data.get('error')}"
            
            print(f"正在生成中... 状态: {status}")

        return None, "任务超时，请稍后重试。"

    except Exception as e:
        return None, f"程序异常: {str(e)}"

# --- Gradio 交互界面 ---
with gr.Blocks(title="首尾帧视频生成测试") as demo:
    gr.Markdown("## 🎬 视频接口生成测试 (/v1/videos)")
    gr.Markdown("当前使用模型：`veo_3_1`")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="视频提示词 (Prompt)", value="画小猫") #
            start_img = gr.Image(label="起始帧 (Start Frame)", type="filepath")
            end_img = gr.Image(label="结束帧 (End Frame)", type="filepath")
            btn = gr.Button("提交视频生成任务", variant="primary")
        
        with gr.Column():
            video_out = gr.Video(label="生成结果")
            info_out = gr.Textbox(label="任务状态信息", interactive=False)

    btn.click(
        fn=generate_video_task,
        inputs=[prompt_input, start_img, end_img],
        outputs=[video_out, info_out]
    )

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