"""pip install httpx; python a7_video_api.py --image reference.jpg --prompt '镜头缓缓推近'

Set A7_API_KEY, or enter the key at the hidden prompt. To resume polling use --job-id.
If submission times out, retry with the printed --request-id; do not create a new ID.
"""
import argparse
import getpass
import mimetypes
import os
import time
import uuid
from contextlib import ExitStack
from pathlib import Path

import httpx


def run(args, key):
    base = args.base_url.rstrip('/')
    with httpx.Client(headers={'Authorization': f'Bearer {key}'}, timeout=180) as client:
        if args.job_id:
            job_id = args.job_id
        else:
            request_id = args.request_id or str(uuid.uuid4())
            print(f'请求编号: {request_id}', flush=True)
            with ExitStack() as stack:
                files = [('images', (path.name, stack.enter_context(path.open('rb')),
                                     mimetypes.guess_type(path.name)[0] or 'application/octet-stream'))
                         for path in args.image]
                try:
                    response = client.post(base + '/api/jobs', headers={'Idempotency-Key': request_id},
                                           data={'prompt': args.prompt, 'duration': str(args.seconds),
                                                 'resolution': args.resolution, 'workflow': args.workflow},
                                           files=files)
                    response.raise_for_status()
                    job_id = response.json()['id']
                except (httpx.HTTPError, ValueError, KeyError) as exc:
                    status = exc.response.status_code if isinstance(exc, httpx.HTTPStatusError) else '未收到明确响应'
                    raise RuntimeError(f'提交未确认（{status}）。如需重试，请复用 --request-id {request_id}') from None
            print(f'任务编号: {job_id}', flush=True)
        deadline = time.monotonic() + 1200
        while time.monotonic() < deadline:
            response = client.get(f'{base}/api/jobs/{job_id}')
            response.raise_for_status()
            job = response.json()
            status = job['status']
            if status == 'succeeded':
                # Resolve the authenticated redirect first; a separate client downloads
                # from R2 without forwarding the API key to another host.
                media = client.get(f'{base}/api/jobs/{job_id}/video', params={'download': 'true'})
                if media.status_code != 307:
                    raise RuntimeError(f'视频下载地址未返回（{media.status_code}）')
                with httpx.stream('GET', media.headers['location'], timeout=180, follow_redirects=True) as stream:
                    stream.raise_for_status()
                    with args.output.open('wb') as output:
                        for chunk in stream.iter_bytes():
                            output.write(chunk)
                print(f'已保存: {args.output.resolve()}', flush=True)
                return job
            if status in ('failed', 'submission_unknown', 'storage_failed'):
                raise RuntimeError(f'{status}: {job.get("error") or "任务未完成"}；任务编号: {job_id}')
            print('生成中', flush=True)
            time.sleep(5)
        raise RuntimeError(f'等待超时，请用 --job-id {job_id} 继续查询，不要重新提交')


def main():
    parser = argparse.ArgumentParser(description='A7 Motion API 图生视频测试')
    parser.add_argument('--base-url', default='https://h3-test.kpass.top')
    parser.add_argument('--image', type=Path, action='append', help='可重复传入，最多 2 张，依次为首帧、尾帧')
    parser.add_argument('--prompt', default='镜头缓缓推近，保持主体细节，画面稳定自然。')
    parser.add_argument('--resolution', default='768p横')
    parser.add_argument('--workflow', default='h3_fastest')
    parser.add_argument('--request-id', help='同一提交的重试必须复用此 UUID')
    parser.add_argument('--seconds', type=int, choices=range(3, 16), default=3, help='视频时长，默认 3 秒')
    parser.add_argument('--job-id', help='查询已有任务，不提交新生成')
    parser.add_argument('--output', type=Path, default=Path('a7-video.mp4'))
    args = parser.parse_args()
    if not args.job_id and not args.image:
        parser.error('请提供 --image，或使用 --job-id 查询已有任务')
    key = os.environ.get('A7_API_KEY') or getpass.getpass('API Key（不显示）: ')
    if not key.strip():
        parser.error('API Key 不能为空')
    try:
        run(args, key.strip())
    except httpx.HTTPStatusError as exc:
        parser.exit(1, f'HTTP {exc.response.status_code}，请检查密钥、任务编号及服务状态。\n')
    except httpx.HTTPError:
        parser.exit(1, '请求中断，请使用已输出的任务编号继续查询。\n')
    except (OSError, RuntimeError) as exc:
        parser.exit(1, f'{exc}\n')


if __name__ == '__main__':
    main()