import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from google_ai.generativeai_video import main
import itertools

app = FastAPI()
api_keys = itertools.cycle(["AIzaSyB2kjF2-S2B5cJiosx_LpApd227w33CVvs", "AIzaSyCor0q5w37Dy6fGxloLlCT7KqyEFU3PWP8"])

class VideoRequest(BaseModel):
    video_path: str
    prompt: str
    mark: str
    sample_data: str


@app.post("/process_video_test/")
async def process_video_test(request: VideoRequest):
    """处理视频请求"""
    video_path = request.video_path
    prompt = request.prompt
    mark = request.mark
    sample_data = request.sample_data

    api_key = "AIzaSyB2kjF2-S2B5cJiosx_LpApd227w33CVvs"
    try:
        print("来一个请求，使用 API key:", api_key)
        result, mark = await main(video_path, api_key, prompt, mark, sample_data)
        return {
            "code": 0,
            "message": "视频处理成功",
            "result": result,
            "mark": mark
        }
    except Exception as e:
        print(f"视频处理失败: {str(e)}")
        return {
            "code": 1,
            "message": f"视频处理失败: {e}",
            "result": f"视频处理失败: {e}",
            "mark": f"视频处理失败: {e}"
        }


@app.post("/process_video/")
async def process_video(request: VideoRequest):
    """处理视频请求"""
    video_path = request.video_path
    prompt = request.prompt
    mark = request.mark
    sample_data = request.sample_data
    # api_keys = ["AIzaSyB2kjF2-S2B5cJiosx_LpApd227w33CVvs","AIzaSyCor0q5w37Dy6fGxloLlCT7KqyEFU3PWP8"]
    api_key = next(api_keys)
    try:
        print("来一个请求，使用 API key:", api_key)
        result, mark = await main(video_path, api_key, prompt, mark, sample_data)
        return {
            "code": 0,
            "message": "视频处理成功",
            "result": result,
            "mark": mark
        }
    except Exception as e:
        print(f"视频处理失败: {str(e)}")
        return {
            "code": 1,
            "message": f"视频处理失败: {e}",
            "result": f"视频处理失败: {e}",
            "mark": f"视频处理失败: {e}"
        }



if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)