"""即梦 AI 视频生成工具

支持：
- 文生视频 (t2v): jimeng_t2v_v30 (720P) / jimeng_t2v_v30_1080p (1080P)
- 图生视频-首帧 (i2v): jimeng_i2v_first_v30 (720P)
"""

import json
import os
import sys
import time
import base64
import requests

from jimeng_gen import jimeng_request, _ensure_dir


# req_key 映射
REQ_KEYS = {
    't2v_720p': 'jimeng_t2v_v30',
    't2v_1080p': 'jimeng_t2v_v30_1080p',
    'i2v_720p': 'jimeng_i2v_first_v30',
}


def generate_video_t2v(prompt, output_path, resolution='720p', aspect_ratio='16:9', frames=121, seed=-1):
    """文生视频：提交任务并轮询获取结果"""
    req_key = REQ_KEYS.get(f't2v_{resolution}', REQ_KEYS['t2v_720p'])
    print(f"[即梦视频] 文生视频 ({resolution}): {prompt[:50]}...")

    body = {
        'req_key': req_key,
        'prompt': prompt,
        'frames': frames,
        'aspect_ratio': aspect_ratio,
        'seed': seed,
    }

    return _submit_and_poll(req_key, body, output_path)


def generate_video_i2v(prompt, image_path_or_url, output_path, aspect_ratio='16:9', frames=121, seed=-1):
    """图生视频-首帧：输入首帧图片和提示词生成视频"""
    req_key = REQ_KEYS['i2v_720p']
    print(f"[即梦视频] 图生视频-首帧: {prompt[:50]}...")

    body = {
        'req_key': req_key,
        'prompt': prompt,
        'frames': frames,
        'seed': seed,
    }

    # 判断是 URL 还是本地文件
    if image_path_or_url.startswith('http'):
        body['image_urls'] = [image_path_or_url]
    else:
        with open(image_path_or_url, 'rb') as f:
            img_b64 = base64.b64encode(f.read()).decode('utf-8')
        body['binary_data_base64'] = [img_b64]

    return _submit_and_poll(req_key, body, output_path)


def _submit_and_poll(req_key, body, output_path, max_wait=180, poll_interval=10):
    """提交任务并轮询结果，下载视频到本地"""
    # 1. 提交任务
    submit_resp = jimeng_request('CVSync2AsyncSubmitTask', body)

    if submit_resp.get('code') != 10000:
        print(f"[即梦视频] 提交失败: {submit_resp}")
        return None

    task_id = submit_resp['data']['task_id']
    print(f"[即梦视频] 任务已提交, task_id: {task_id}")

    # 2. 轮询查询结果
    max_polls = max_wait // poll_interval
    for i in range(max_polls):
        time.sleep(poll_interval)
        result = jimeng_request('CVSync2AsyncGetResult', {
            'req_key': req_key,
            'task_id': task_id,
        })

        resp_code = result.get('code')
        resp_data = result.get('data', {})

        if resp_code == 10000 and resp_data.get('status') == 'done':
            video_url = resp_data.get('video_url')
            if video_url:
                # 下载视频
                print(f"[即梦视频] 生成完成，正在下载...")
                video_data = requests.get(video_url).content
                _ensure_dir(output_path)
                with open(output_path, 'wb') as f:
                    f.write(video_data)
                print(f"[即梦视频] 视频已保存: {output_path}")
                return output_path
            else:
                print(f"[即梦视频] 任务完成但未找到视频URL: {json.dumps(result, ensure_ascii=False)[:300]}")
                return None

        if resp_code != 10000 and resp_data.get('status') not in (None, 'running', 'pending'):
            print(f"[即梦视频] 任务失败: {json.dumps(result, ensure_ascii=False)[:300]}")
            return None

        print(f"[即梦视频] 等待中... ({i+1}/{max_polls})")

    print("[即梦视频] 超时")
    return None


def main():
    import argparse
    parser = argparse.ArgumentParser(description='即梦 AI 视频生成')
    parser.add_argument('prompt', help='视频描述提示词')
    parser.add_argument('output', help='输出视频路径 (.mp4)')
    parser.add_argument('--mode', choices=['t2v', 'i2v'], default='t2v', help='生成模式: t2v=文生视频, i2v=图生视频')
    parser.add_argument('--resolution', choices=['720p', '1080p'], default='720p', help='分辨率 (仅 t2v)')
    parser.add_argument('--aspect-ratio', default='16:9', help='宽高比: 16:9, 4:3, 1:1, 3:4, 9:16, 21:9')
    parser.add_argument('--frames', type=int, default=121, help='视频帧数 (默认121)')
    parser.add_argument('--image', help='首帧图片路径或URL (仅 i2v 模式)')
    parser.add_argument('--seed', type=int, default=-1, help='随机种子 (-1=随机)')

    args = parser.parse_args()

    if args.mode == 'i2v':
        if not args.image:
            parser.error('i2v 模式需要 --image 参数')
        result = generate_video_i2v(args.prompt, args.image, args.output,
                                     aspect_ratio=args.aspect_ratio, frames=args.frames, seed=args.seed)
    else:
        result = generate_video_t2v(args.prompt, args.output,
                                     resolution=args.resolution, aspect_ratio=args.aspect_ratio,
                                     frames=args.frames, seed=args.seed)

    if result:
        print(f"\n视频生成成功: {result}")
    else:
        print("\n视频生成失败")
        sys.exit(1)


if __name__ == '__main__':
    main()