#!/usr/bin/env python3
"""Query an existing H3 V2 task and download its result. Never creates a task."""
import json
import os
from pathlib import Path
import time
import urllib.error
import urllib.parse
import urllib.request


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def main():
    origin = 'https://api.jusuanhub.com'
    key = os.environ['JUSP_API_KEY']
    task_id = json.loads(Path('create.json').read_text())['task_id']
    query = origin + '/v2/query/video_generation/' + urllib.parse.quote(task_id, safe='')
    opener = urllib.request.build_opener(NoRedirect)
    deadline = time.monotonic() + 1800
    failures = 0
    while time.monotonic() < deadline:
        request = urllib.request.Request(query, headers={'Authorization': 'Bearer ' + key})
        try:
            with opener.open(request, timeout=min(60, max(1, deadline - time.monotonic()))) as response:
                task = json.load(response)['task']
            failures = 0
        except (urllib.error.URLError, TimeoutError) as error:
            if isinstance(error, urllib.error.HTTPError) and error.code not in {429, 502, 503, 504}:
                # Do not print credentials, signed URLs or arbitrary remote response bodies.
                raise RuntimeError(f'查询 HTTP {error.code}；保留原 task_id 联系支持') from None
            failures += 1
            if failures > 5:
                raise RuntimeError('连续查询失败；稍后运行本脚本继续原任务') from None
            time.sleep(min(2 ** failures, 30, max(0, deadline - time.monotonic())))
            continue
        if task['status'] == 'succeeded':
            url = task['content']['url']
            parsed = urllib.parse.urlsplit(url)
            expected = urllib.parse.urlsplit(origin)
            if (parsed.scheme, parsed.netloc) != (expected.scheme, expected.netloc) or parsed.username:
                raise RuntimeError('结果来源异常；不要转发 API Key')
            # Ticket URL already authorizes the result. No Bearer header.
            partial = Path('result.mp4.part')
            try:
                with opener.open(url, timeout=60) as response, partial.open('wb') as output:
                    if response.status != 200 or response.headers.get_content_type() != 'video/mp4':
                        raise RuntimeError('未返回完整 MP4；保留任务 ID 并重新查询')
                    total = 0
                    download_deadline = time.monotonic() + 300
                    while True:
                        if time.monotonic() >= download_deadline:
                            raise TimeoutError('下载超时，重新运行可查询原任务并重下')
                        chunk = response.read(1024 * 1024)
                        if not chunk:
                            break
                        total += len(chunk)
                        if total > 1024 * 1024 * 1024:
                            raise RuntimeError('超过示例的 1 GiB 下载预算')
                        output.write(chunk)
                partial.replace('result.mp4')
            except urllib.error.URLError:
                raise RuntimeError('下载失败；保留原任务 ID 后重查，勿新建任务') from None
            print('已保存 result.mp4；实际画布', task.get('width'), task.get('height'))
            return
        if task['status'] in {'failed', 'cancelled'}:
            raise RuntimeError('任务终止：' + task.get('error', {}).get('code', task['status']))
        if task['status'] not in {'queued', 'running'}:
            raise RuntimeError('未知任务状态；保留 task_id 联系支持')
        time.sleep(min(5, max(0, deadline - time.monotonic())))
    raise TimeoutError('本地等待结束；服务端任务未取消，稍后重跑本脚本查询原任务')


if __name__ == '__main__':
    main()