#!/usr/bin/env python3
"""智叶 API Seedance 示例：模型列表、单次提交、可恢复查询和同站下载。"""
import argparse
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

OPENER = urllib.request.build_opener(NoRedirect())
BASE = os.environ.get('ZHIYE_API_BASE', 'https://api.ziyelian.site/v1').rstrip('/')
ORIGIN = urllib.parse.urlsplit(BASE)

def request(path):
    key = os.environ.get('ZHIYE_API_KEY')
    if not key:
        raise ValueError('请设置 ZHIYE_API_KEY 环境变量')
    url = urllib.parse.urljoin(BASE + '/', path)
    target = urllib.parse.urlsplit(url)
    if target.scheme != 'https' or (target.scheme, target.netloc) != (ORIGIN.scheme, ORIGIN.netloc):
        raise ValueError('仅向配置的 HTTPS 同站地址发送令牌')
    if target.username or target.password:
        raise ValueError('URL 不得包含凭据')
    return urllib.request.Request(url, headers={'Authorization': 'Bearer ' + key, 'Accept': 'application/json'})

def api(path, payload=None):
    req = request(path)
    if payload is not None:
        req.data = json.dumps(payload, ensure_ascii=False).encode()
        req.add_header('Content-Type', 'application/json')
    with OPENER.open(req, timeout=120) as r:
        return json.load(r)

def save(path, value):
    temp = path.with_suffix(path.suffix + '.part')
    temp.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding='utf-8')
    temp.replace(path)

def main():
    os.umask(0o077)
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='cmd', required=True)
    sub.add_parser('models')
    create = sub.add_parser('create')
    create.add_argument('request_json', type=Path)
    create.add_argument('order_dir', type=Path)
    for name in ('status', 'watch', 'download'):
        item = sub.add_parser(name)
        item.add_argument('order_dir', type=Path)
        if name == 'watch':
            item.add_argument('--wait-seconds', type=int, default=900)
    args = parser.parse_args()
    if args.cmd == 'models':
        print(json.dumps(api('models'), ensure_ascii=False, indent=2))
        return
    directory = args.order_dir
    if args.cmd == 'create':
        payload = json.loads(args.request_json.read_text(encoding='utf-8'))
        if not isinstance(payload, dict) or not isinstance(payload.get('content'), list) or not payload['content']:
            raise ValueError('请求须为对象，且包含非空 content 数组')
        request('contents/generations/tasks')  # 提交前检查环境变量与地址。
        directory.mkdir(parents=True, exist_ok=False)
        save(directory / 'request.json', payload)
        save(directory / 'submission.json', {'state': 'submitting', 'created_at': int(time.time())})
        try:
            result = api('contents/generations/tasks', payload)
            save(directory / 'create.json', result)
            task_id = result.get('id')
            if not isinstance(task_id, str) or not task_id:
                raise ValueError('创建响应缺少 id，请保留目录并核对提交结果')
            save(directory / 'task.json', {'id': task_id})
            save(directory / 'submission.json', {'state': 'accepted'})
            print(task_id)
        except Exception:
            print('提交结果待核对。请保留订单目录，不要删除目录重新提交。')
            raise
        return
    task_id = json.loads((directory / 'task.json').read_text())['id']
    path = 'contents/generations/tasks/' + urllib.parse.quote(task_id, safe='')
    deadline = time.monotonic() + (args.wait_seconds if args.cmd == 'watch' else 0)
    while True:
        try:
            result = api(path)
        except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError) as error:
            transient = not isinstance(error, urllib.error.HTTPError) or error.code == 429 or error.code >= 500
            if args.cmd == 'watch' and transient and time.monotonic() < deadline:
                time.sleep(10)
                continue
            raise
        save(directory / 'result.json', result)
        state = result.get('status', 'unknown')
        print(state)
        if args.cmd != 'watch' or state in ('succeeded', 'failed', 'expired', 'cancelled', 'canceled'):
            break
        if time.monotonic() >= deadline:
            print('本次查询窗口结束；稍后可对同一目录继续 watch。')
            return
        time.sleep(5)
    if args.cmd == 'download':
        if state != 'succeeded':
            raise ValueError('任务尚未成功，请先继续查询')
        url = '/v1/tasks/' + urllib.parse.quote(task_id, safe='') + '/artifacts/video/content'
        req = request(url)
        req.add_header('Accept', 'video/mp4')
        partial = directory / 'result.mp4.part'
        with OPENER.open(req, timeout=180) as response, partial.open('wb') as output:
            first = response.read(32)
            if len(first) < 12 or first[4:8] != b'ftyp':
                raise ValueError('下载结果不是预期的 MP4 文件')
            output.write(first)
            while True:
                chunk = response.read(1024 * 1024)
                if not chunk:
                    break
                output.write(chunk)
        partial.replace(directory / 'result.mp4')
        print(directory / 'result.mp4')

if __name__ == '__main__':
    main()