#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : AI.  @by PyCharm
# @File         : videos
# @Time         : 2024/10/21 20:17
# @Author       : betterme
# @WeChat       : meutils
# @Software     : PyCharm
# @Description  : https://platform.minimaxi.com/document/video_generation?key=66d1439376e52fcee2853049
# https://useapi.net/docs/start-here/setup-minimax
# token 过期时间一个月: 看下free hailuo
# https://jwt.io/
import asyncio

# todo: check token


from meutils.pipe import *
from meutils.hash_utils import md5
from meutils.io.files_utils import to_bytes
from meutils.schemas.hailuo_types import BASE_URL_ABROAD as BASE_URL

from meutils.decorators.retry import retrying

from meutils.apis.hailuoai.yy import get_yy
from meutils.apis.hailuoai.utils import PARAMS as params, get_access_token, upload

from meutils.schemas.image_types import ImageRequest, ImagesResponse


async def generate(request: ImageRequest, api_key: Optional[str] = None):
    refresh_token = api_key
    token = await get_access_token(refresh_token)

    payload = {
        "quantity": request.n,
        "parameter": {
            "modelID": request.model,
            "desc": request.prompt,
            "fileList": [],
            "useOriginPrompt": True if "--enhance" not in request.prompt else False,  # enhance-prompt
            "aspectRatio": request.aspect_ratio or "Auto",
            "resolution": (request.resolution or "1K").upper()
        }
    }

    if request.image:
        payload['parameter']['fileList'] = [
            {
                "url": url,

                # "id": "469899357624119304",
                # "name": "kling_watermark.png",
                # "type": "png",
                # "frameType": 3
            }
            for url in request.image_urls
        ]

    path = "/v2/api/multimodal/generate/image"
    headers = {
        'Content-Type': 'application/json',
        'token': token,
        'yy': get_yy(payload, params, path),
    }

    async with httpx.AsyncClient(base_url=BASE_URL, headers=headers, timeout=60) as client:
        response = await client.post(path, params=params, content=json.dumps(payload))
        response.raise_for_status()

        data = response.json()

        logger.debug(bjson(data))

        image_response = None
        if task_id := (data.get('data') or {}).get('task', {}).get('batchID'):
            logger.debug(f"TASK_ID: {task_id}")
            for _ in range(60):
                await asyncio.sleep(5) if _ else await asyncio.sleep(16)

                if image_response := await get_task(task_id, token):
                    if image_response.data and len(image_response.data) == request.n:
                        return image_response

        if image_response and image_response.data:
            return image_response
        else:
            raise Exception("invalid image")


async def get_task(task_id: str, token: str):
    batch_id = task_id
    if '_' in task_id:
        batch_id, task_id = task_id.split('_', 1)

    payload = {
        "batchInfoList": [
            {
                "batchID": batch_id,
                "batchType": 1
            }
        ],
        "type": 1
    }

    path = "/v4/api/multimodal/video/processing"
    headers = {
        'Content-Type': 'application/json',
        'token': token,
        'yy': get_yy(payload, params, url=path),
    }

    async with httpx.AsyncClient(base_url=BASE_URL, headers=headers, timeout=60) as client:
        response = await client.post(path, params=params, content=json.dumps(payload)
                                     )
        response.raise_for_status()
        data = response.json()
        logger.debug(bjson(data))

        if batchs := data['data']['batchVideos']:
            if (items := batchs) and (assets := items[0]['assets']):
                # if (items := batchs | xfilter_(lambda x: x['batchID'] == batch_id)) and (assets := items[0]['assets']):

                logger.debug(bjson(assets))

                if "失败" in str(assets): raise Exception(f"invalid image: {task_id}")

                data = [
                    {"url": image.get("downloadURL")}
                    for image in assets if image.get("downloadURL")
                ]

                return ImagesResponse(data=data, metadata=payload)


if __name__ == '__main__':
    token = None
    token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3NzIyODE1NzcsInVzZXIiOnsiaWQiOiI0NDQyMjk2MDAzMzA0OTgwNTUiLCJuYW1lIjoibWZ1aiBiamhuIiwiYXZhdGFyIjoiIiwiZGV2aWNlSUQiOiIzMzkxMTQ5Mjg4NjU1Mjk4NjQiLCJpc0Fub255bW91cyI6ZmFsc2V9fQ.__NDyZQQqyYb7TLrumo944EfuCmrbzYngQloNBK4CmM"
    token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3NzIzNTA5NjUsInVzZXIiOnsiaWQiOiI0Njk4ODIxOTY3NDM1Mjg0NDkiLCJuYW1lIjoiYWZzbCBkcnF2IiwiYXZhdGFyIjoiIiwiZGV2aWNlSUQiOiIzMTE2NzAxODUwMDg0ODg0NTIiLCJpc0Fub255bW91cyI6ZmFsc2V9fQ.3pO0O36-um2fQs0ML0eHwpi0D7rV5yjmnjcpuiZcNKw"
    model = "nano-banana2"
    # model = "nano-banana2_2k"
    # model = "nano-banana2_4k"

    # seedream-4.5
    # gpt-image-1.5 {"quantity":1,"parameter":{"modelID":"","desc":"a dog","fileList":[],"useOriginPrompt":true,"aspectRatio":"Auto","resolution":"Low"}}
    # Low
    # Medium
    # High

    request = ImageRequest(
        model=model,
        n=1,
        # prompt="笑起来",
        prompt="哭起来",

        image="https://s3.ffire.cc/files/jimeng.jpg",

        size="16:9"
    )

    r = arun(generate(request, api_key=token))

    # model = "veo3.1-t2v-fast"
    task_id = "469896308304285702_469873163803475974"
    # "message": "内容生成失败，请重试",

    # task_id = "469872667202084865_469872667202084866"
    # 469872667202084865
    # task_id = "469901294006366214_469901294006366212"  # n=1
    task_id = "469901776552665096_469901776552665097"  # n=2

    task_id = "469902511826763781"  # n=3

    task_id = "469911153359634438"  # 4k

    task_id = "469916812486545413"

    # arun(get_task(task_id=task_id, token=token))
