#!/usr/bin/env python3
"""Minimal server-side Seedance V1 example. Only explicit 'submit' creates a paid task."""
import argparse
import json
import os
from pathlib import Path
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlparse
from urllib.request import Request, HTTPRedirectHandler, build_opener

BASE = 'https://api.wxksai.com'

class SafeRedirect(HTTPRedirectHandler):
    def redirect_request(self, request, response, code, message, headers, url):
        if urlparse(url).scheme != 'https':
            raise ValueError('Refusing a redirect outside HTTPS')
        redirected = super().redirect_request(request, response, code, message, headers, url)
        if redirected is not None and urlparse(url).netloc != urlparse(request.full_url).netloc:
            redirected.remove_header('Authorization')
        return redirected

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='command', required=True)
    submit = sub.add_parser('submit', help='Create ONE paid task; never automatically retries')
    submit.add_argument('--prompt', required=True)
    submit.add_argument('--seconds', type=int, choices=range(5, 31), default=10)
    submit.add_argument('--ratio', choices=['16:9','9:16','1:1','4:3','3:4','21:9'], default='16:9')
    for name in ['query', 'wait', 'download']:
        command = sub.add_parser(name)
        command.add_argument('task_id')
        if name == 'wait':command.add_argument('--minutes', type=int, choices=range(1, 121), default=20)
        if name == 'download':command.add_argument('--output', default='result.mp4')
    args = parser.parse_args()
    key = os.environ.get('WXKS_API_KEY', '').strip()
    if not key:
        parser.error('Set WXKS_API_KEY in your server environment first')
    opener = build_opener(SafeRedirect())
    path = '/v1/videos' if args.command == 'submit' else '/v1/videos/' + quote(args.task_id, safe='')
    headers = {'Authorization': 'Bearer ' + key}
    if args.command == 'submit':
        body = json.dumps({'model':'seedance-2.5-special2','prompt':args.prompt,'seconds':str(args.seconds),'resolution':'1080p','aspect_ratio':args.ratio}).encode()
        headers['Content-Type'] = 'application/json'
        request = Request(BASE + path, data=body, headers=headers, method='POST')
        try:
            with opener.open(request, timeout=120) as response:
                result = json.load(response)
        except (URLError, TimeoutError):
            print('Submit was not retried. A task may already exist; check the console before submitting again.', file=sys.stderr)
            raise
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return
    if args.command == 'download':
        output = Path(args.output)
        partial = output.with_name(output.name + '.part')
        if output.exists() or partial.exists():
            raise FileExistsError('Output or .part file exists; choose another --output')
        with opener.open(Request(BASE+path,headers=headers),timeout=30) as response:
            task=json.load(response)
        if task.get('status') != 'completed':
            raise ValueError('Task is not completed; query the same task again later')
        # Exclusive creation also prevents overwriting an existing local file.
        with opener.open(Request(BASE+path+'/content',headers=headers),timeout=120) as response:
            with partial.open('xb') as stream:
                while chunk := response.read(1024*1024):stream.write(chunk)
        with output.open('xb') as stream, partial.open('rb') as source:
            while chunk := source.read(1024*1024):stream.write(chunk)
        partial.unlink()
        print(str(output))
        return
    deadline = time.monotonic() + (args.minutes * 60 if args.command == 'wait' else 30)
    while time.monotonic() < deadline:
        try:
            with opener.open(Request(BASE+path,headers=headers),timeout=30) as response:
                task=json.load(response)
        except HTTPError as error:
            if args.command != 'wait' or error.code not in (429,500,502,503,504):raise
            time.sleep(30)
            continue
        except (URLError,TimeoutError):
            if args.command != 'wait':raise
            time.sleep(30)
            continue
        print(json.dumps(task,ensure_ascii=False,indent=2),flush=True)
        status=task.get('status')
        if args.command=='query' or status=='completed':return
        if status=='failed':raise RuntimeError('Task failed; inspect the task log and refund record')
        if status not in ('queued','in_progress'):raise ValueError('Unknown status; retain the original task ID')
        time.sleep(30)
    raise TimeoutError('Local waiting ended. Retain the task ID; do not resubmit.')

if __name__ == '__main__':
    main()