#!/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/

# todo: check token

import oss2

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

from meutils.schemas.hailuo_types import VideoRequest, VideoResponse
from meutils.decorators.retry import retrying
from meutils.notice.feishu import send_message as _send_message, VIDEOS

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

from openai import AsyncOpenAI
from meutils.schemas.video_types import SoraVideoRequest, Video

send_message = partial(
    _send_message,
    title=__name__,
    url=VIDEOS
)

APP_ID = '3001'
VERSION_CODE = '22203'

MODEL_MAPPING = {
    # video-01 video-01 video-01-live2d S2V-01

    "t2v-01": "23000",  # 23010
    "t2v-01-director": "23010",

    "i2v-01": "23001",
    "i2v-01-live": "23011",
    "video-01-live2d": "23011",
    "s2v-01": "23021",

    # "23210" # 要积分
}


class Tasks(object):

    def __init__(self, base_url: Optional[str] = None, api_key: str = None):
        base_url = BASE_URL
        self.client = AsyncOpenAI(base_url=base_url, api_key=api_key)

    async def create(self, request: SoraVideoRequest):
        payload = {
            "model": request.model,
            "prompt": request.prompt,

            **(request.metadata or {})
        }

        if request.seconds:
            payload["duration"] = int(request.seconds)

        if request.resolution:
            payload['resolution'] = request.resolution

        if request.ratio:
            payload["aspect_ratio"] = request.ratio

        if image_urls := request.input_reference:
            payload['model'] = request.model.replace("text-to-video", "image-to-video").replace("t2v", "i2v")

            # if len(image_urls) > 1:
            #     payload["image_list"] = image_urls
            # else:
            #     payload["image_url"] = image_urls[0]

            payload["image_url"] = image_urls[0]
            payload["image_list"] = image_urls

            # https://aimlapi.com/models/veo-3-1-reference-to-video
            # 模型区分 参考图 首尾帧
            # if any( i in request.model for i in {"veo"}):

            # klingai/video-o1-reference-to-video image_list
            # klingai/video-o1-image-to-video image_url last_image_url

        else:
            payload['model'] = request.model.replace("image-to-video", "text-to-video").replace("i2v", "t2v")

        # 首尾帧
        if request.first_frame_image:
            if request.model.startswith(("klingai/video-o1",)):
                payload['model'] = "klingai/video-o1-image-to-video"

            payload["image_url"] = request.first_frame_image
            payload["first_frame_image"] = request.first_frame_image

        if request.last_frame_image:
            payload["last_image_url"] = request.last_frame_image

        # 数字人 对口型
        if request.audio and any(i in request.model for i in {"omnihuman", 'avatar'}):
            payload['image_url'] = request.image
            payload['audio_url'] = request.audio

        if request.video:
            if request.model.startswith(("alibaba/wan-2-6",)):
                payload["video_urls"] = [request.video]
                payload['model'] = "alibaba/wan-2-6-r2v"

            else:
                payload["video_url"] = request.video
                payload['keep_audio'] = True

                if request.model.endswith('edit'):  # kling
                    payload['model'] = "klingai/video-o1-video-to-video-edit"
                else:
                    payload['model'] = "klingai/video-o1-video-to-video-reference"

        logany(bjson(payload))

        response = await self.client.post(
            path="/video/generations",
            body=payload,
            cast_to=object
        )
        """
        {
    "id": "9913239d-4fa8-47ea-b51d-d313e29caba5:alibaba/wan2.5-i2v-preview",
    "status": "queued",
    "meta": {
        "usage": {
            "tokens_used": 105000
        }
    }

    {
    "generation_id": "339995387916622:minimax/hailuo-02",
    "status": "queued",
    "meta": {
        "usage": {
            "tokens_used": 588000
        }
    }
}
}
        """

        logger.debug(bjson(response))

        return response

    async def get(self, task_id: str):
        logger.debug(task_id)

        response = await self.client.get(
            path=f"/video/generations?generation_id={task_id}",
            cast_to=object
        )
        """
        {
    "id": "9913239d-4fa8-47ea-b51d-d313e29caba5:alibaba/wan2.5-i2v-preview",
    "status": "completed",
    "video": {
        "url": "https://cdn.aimlapi.com/alpaca/1d/dd/20251107/30b07d9c/42740107-9913239d-4fa8-47ea-b51d-d313e29caba5.mp4?Expires=1762593280&OSSAccessKeyId=LTAI5tBLUzt9WaK89DU8aECd&Signature=Guk6apyEnKeuniLv0mcBJhkHO%2FI%3D"
    }
}

{
    "id": "",
    "status": "unknown",
    "error": {
        "name": "Error",
        "message": "invalid params, task_id cannot by empty"
    }
}

        """
        logger.debug(bjson(response))

        status = response
        if response.get('error'):
            status = 'failed'

        video = Video(
            id=task_id,
            status=status,
            video_url=(response.get("video") or {}).get("url"),

            error=response.get("error")
        )

        # logger.debug(bjson(video))

        return video


# @retrying(predicate=lambda r: r.base_resp.status_code in {1000061, 1500009})  # 限流
async def create_task(request: SoraVideoRequest, token: Optional[str] = None):
    refresh_token = token
    token = await get_access_token(refresh_token)

    payload = {
        "desc": request.prompt,
        "useOriginPrompt": not request.prompt_optimizer,
        "fileList": [],
        "modelID": "23000",  # 文生视频
        "quantity": f"{request.n}",

        # "extra": {
        #     "promptStruct": "{"value":[{"type":"paragraph","children":[{"text":"动起来"}]}],"length":3,"plainLength":3,"
        #     rawLength":3}",
        #     "templateID": ""
        # },

        # "durationType": 1,
        # "resolutionType": 3
    }

    payload = {
        "quantity": 1,
        "parameter": {
            "modelID": request.model,
            "desc": request.prompt,
            "fileList": [],
            "useOriginPrompt": request.enhance_prompt or True,
            "duration": int(request.seconds),
            "resolution": (request.resolution or "720").rstrip('p').rstrip('P'),
            "aspectRatio": ""
        },
        "videoExtra": {
            "promptStruct": "{\"value\":[{\"type\":\"paragraph\",\"children\":[{\"text\":\"a cat\"}]}],\"length\":5,\"plainLength\":5,\"rawLength\":5}"
        }
    }
    """
    {"quantity":1,"parameter":{"modelID":"veo3.1-t2v-fast","desc":"a cat","fileList":[],"useOriginPrompt":true,"resolution":"720","duration":8,"aspectRatio":"16:9"},"videoExtra":{"promptStruct":"{\"value\":[{\"type\":\"paragraph\",\"children\":[{\"text\":\"a cat\"}]}],\"length\":5,\"plainLength\":5,\"rawLength\":5}"}}
    """
    if request.ratio:
        payload['parameter']['aspectRatio'] = request.ratio

    if request.first_frame_image:
        file = await to_bytes(request.first_frame_image)
        if request.first_frame_image.startswith("http") and BASE_URL.endswith(".video") and 0:  # 必须上传
            file_data = {
                # "id": data['data']['fileID'],
                # "name": "_.png",
                # "type": "png",
                "url": request.first_frame_image,  #######
            }

            # {"desc": "跳起来", "useOriginPrompt": false, "fileList": [{"id": "338311163211288581",
            #                                                            "url": "https://cdn.hailuoai.video/moss/prod/2025-01-22-11/user/multi_chat_file/de5a4cec-eb26-4380-94e4-13b268bf5c0d.jpg",
            #                                                            "name": "duikoux.jpg", "type": "jpg"}],
            #  "modelID": "23021", "quantity": "1"}
        else:
            data = await upload(file, token=refresh_token)
            file_data = {
                "id": data['data']['fileID'],
                "name": "_.png",
                "type": "png",
                "url": data['data']['ossPath'],
            }

        payload["fileList"].append(file_data)
        payload["modelID"] = MODEL_MAPPING.get(request.model.lower(), "23001")

    logger.debug(bjson(payload))

    # payload = {
    #     "desc": "笑起来",
    #     "useOriginPrompt": True,
    #     "fileList": [
    #         {
    #             # "id": "335822260344578050",
    #             ############# 不需要走上传了
    #             # "url": "https://cdn.hailuoai.video/moss/prod/2025-01-15-13/user/multi_chat_file/6ca7b141-383c-4481-8bba-b9148c1339c0.jpg?x-oss-process=image/resize,p_50/format,webp",
    #             "url": "https://oss.ffire.cc/files/kling_watermark.png",
    #             "name": "baobao.jpg",
    #             "type": "jpg"
    #         }
    #     ],
    #     "modelID": "23001",
    #     "quantity": "1"
    # }

    path = "/v2/api/multimodal/generate/video"

    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))

        task_id = (data.get('data') or {}).get('id', '')
        response = VideoResponse(
            task_id=f"hailuoai-{task_id}",
            base_resp=data.get('statusInfo', {}),
            system_fingerprint=refresh_token
        )
        if response.base_resp.status_code != 0:  # 451
            from fastapi import HTTPException, status
            raise HTTPException(status_code=status.HTTP_451_UNAVAILABLE_FOR_LEGAL_REASONS,
                                detail=response.base_resp)

        return response
    # {
    #     "statusInfo": {
    #         "code": 1000061,
    #         "httpCode": 0,
    #         "message": "上一个视频任务未完成，请稍后再试",
    #         "serviceTime": 1729512914,
    #         "requestID": "82bc8c60-4dc3-4ad0-b5b6-b1836e0c88ab",
    #         "debugInfo": "",
    #         "serverAlert": 0
    #     }
    # }

    # {
    #     "data": {
    #         "id": "304746220940677121"
    #     },
    #     "statusInfo": {
    #         "code": 0,
    #         "httpCode": 0,
    #         "message": "成功",
    #         "serviceTime": 1729513305,
    #         "requestID": "caaf2364-d2ed-45df-b79a-827810a5d58c",
    #         "debugInfo": "",
    #         "serverAlert": 0
    #     }
    # }


# 307134660730421250
async def get_task(task_id: str, token: str):
    task_id = task_id.rsplit('-', 1)[-1]

    # params = {
    #     'device_platform': 'web',
    #     'app_id': APP_ID,
    #     'version_code': VERSION_CODE,
    #     'biz_id': 0,
    #     'uuid': '8c059369-00bf-4777-a426-d9c9b7984ee6',
    #     'device_id': '243713252545986562',
    #     'os_name': 'Mac',
    #     'browser_name': 'chrome',
    #     'device_memory': '8',
    #     'cpu_core_num': '10',
    #     'browser_language': 'zh-CN',
    #     'browser_platform': 'MacIntel',
    #     'screen_width': '1920',
    #     'screen_height': '1080',
    #     'unix': f'{int(time.time())}000',
    #     'idList': task_id
    # }

    # _params = {
    #     **params,
    #     'unix': f'{int(time.time())}000',
    # }

    payload = {
        "batchInfoList": [
            {
                "batchID": task_id,
                "batchType": 0
            }
        ],
        "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))

        response = VideoResponse(
            task_id=task_id,
            base_resp=data.get('statusInfo', {}),
            videos=(data.get('data') or {}).get("videos", [])[:1]
        )
        return response


if __name__ == '__main__':
    token = None
    token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3NzIyODE1NzcsInVzZXIiOnsiaWQiOiI0NDQyMjk2MDAzMzA0OTgwNTUiLCJuYW1lIjoibWZ1aiBiamhuIiwiYXZhdGFyIjoiIiwiZGV2aWNlSUQiOiIzMzkxMTQ5Mjg4NjU1Mjk4NjQiLCJpc0Fub255bW91cyI6ZmFsc2V9fQ.__NDyZQQqyYb7TLrumo944EfuCmrbzYngQloNBK4CmM"

    request = VideoRequest(
        # model="t2v-01",
        model="I2V-01-live",
        # model="S2V-01-live",


        # prompt="smile",  # 307145017365086216
        prompt="动起来",  # 307145017365086216
        # first_frame_image="https://oss.ffire.cc/files/kling_watermark.png"  # 307173162217783304
    )

    # r = arun(create_task(request, token=token))

    model = "veo3.1-t2v-fast"
    task_id = "hailuoai-469852272096808964"

    # {"quantity": 1, "parameter": {"modelID": "sora2-i2v", "desc": "笑起来啊", "fileList": [{"id": "469914687591321600",
    #                                                                                         "url": "https://cdn.hailuoai.video/moss/prod/2026-01-20-15/user/multi_chat_file/997398be-fe17-4af1-ac9f-7057a54e5587.jpeg?x-oss-process=image/resize,p_50/format,webp",
    #                                                                                         "name": "cropped_1768892531810.jpeg",
    #                                                                                         "type": "jpeg",
    #                                                                                         "frameType": 0}],
    #                               "useOriginPrompt": true, "resolution": "720", "duration": 4, "aspectRatio": "16:9"},
    #  "videoExtra": {
    #      "promptStruct": "{\"value\":[{\"type\":\"paragraph\",\"children\":[{\"text\":\"笑起来啊\"}]}],\"length\":4,\"plainLength\":4,\"rawLength\":4}"}}

    # data = {
    #     "model": "video-01",
    #     "prompt": "画面中两个人非常缓慢地拥抱在一起",
    #     "prompt_optimizer": True,
    #     # "first_frame_image": "https://hg-face-domestic-hz.oss-cn-hangzhou.aliyuncs.com/avatarapp/ai-cache/54883340-954c-11ef-8920-db8e7bfa3fdf.jpeg"
    # }
    # request = VideoRequest(**data)
    arun(get_task(task_id=task_id, token=token))
