#!/usr/bin/env python3
"""ModelKiwi async video. Submit once, poll, download; resume with --task ID."""
import argparse,json,os,re,sys,time
from pathlib import Path
from urllib.request import Request,build_opener,HTTPRedirectHandler
class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self,*a,**kw):return None

def video_body(model,prompt,duration=8,resolution='720p'):
    if model not in ['grok-imagine-video','grok-imagine-video-1.5'] or not prompt or not 1<=duration<=15:raise ValueError('Invalid model/prompt/duration')
    if resolution not in (['480p','720p','1080p'] if model.endswith('-1.5') else ['480p','720p']):raise ValueError('Invalid resolution')
    return dict(model=model,prompt=prompt,duration=duration,resolution=resolution)

def main():
    p=argparse.ArgumentParser();p.add_argument('--model',default='grok-imagine-video');p.add_argument('--prompt');p.add_argument('--duration',type=int,default=8);p.add_argument('--resolution',default='720p');p.add_argument('--task');p.add_argument('--out');args=p.parse_args()
    body=None if args.task else video_body(args.model,args.prompt,args.duration,args.resolution)
    key=os.environ['MODELKIWI_API_KEY'];opener=build_opener(NoRedirect)
    def request(suffix='',data=None):
        headers={'Authorization':'Bearer '+key}
        if data is not None:headers['Content-Type']='application/json'
        req=Request('https://api.modelkiwi.com/v1/videos'+suffix,data=json.dumps(data).encode() if data is not None else None,headers=headers,method='POST' if data is not None else 'GET')
        return opener.open(req,timeout=180)
    if args.task:task_id=args.task
    else:
        with request(data=body) as r:task_id=json.load(r).get('id')
    if not isinstance(task_id,str) or not re.fullmatch(r'[-A-Za-z0-9_]+',task_id):raise ValueError('Missing task ID')
    print('Task ID:',task_id,'(resume: --task '+task_id+')',flush=True)
    complete=False
    for _ in range(120):
        with request('/'+task_id) as r:status=json.load(r).get('status')
        print('Status:',status,flush=True)
        if status in ['completed','succeeded','done','SUCCESS']:complete=True;break
        if status in ['failed','error','cancelled','failure','FAILURE']:raise RuntimeError('Video failed')
        time.sleep(5)
    if not complete:raise TimeoutError('Polling timed out; resume with task ID')
    with request('/'+task_id+'/content') as r:data=r.read()
    if len(data)<12 or data[4:8]!=b'ftyp':raise ValueError('Expected MP4')
    output=Path(args.out or task_id+'.mp4')
    with output.open('xb') as f:f.write(data)
    print('Saved:',output.resolve())
if __name__=='__main__':
    try:main()
    except Exception:
        print('Video request stopped. Keep the task ID to resume; do not submit another task blindly.',file=sys.stderr);sys.exit(1)