#!/usr/bin/env python3
"""Generate one MP4 through QDivZero. Python 3.11+, no third-party packages.

Set QDIV0_API_KEY and QDIV0_ACCOUNT_ID. The deployment must be running.
One request may incur compute charges. This client never retries automatically.
"""
import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import shutil
import time
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
from uuid import uuid4

SHAPES = {'16:9': (1024, 576), '9:16': (576, 1024), '1:1': (768, 768)}


def multipart(fields):
    boundary = 'qdivzero-' + uuid4().hex
    chunks = []
    for name, value in fields.items():
        chunks.append((f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n'
                       f'\r\n{value}\r\n').encode('utf-8'))
    chunks.append(f'--{boundary}--\r\n'.encode())
    return b''.join(chunks), f'multipart/form-data; boundary={boundary}'


def generate(fields, output, timeout=1800):
    output = Path(output)
    if output.suffix.lower() != '.mp4':
        raise ValueError('The output filename must end in .mp4.')
    metadata_path = output.with_suffix('.json')
    if output.exists() or metadata_path.exists():
        raise FileExistsError('Choose a new output filename; the video or metadata already exists.')
    api_key = os.environ.get('QDIV0_API_KEY')
    account_id = os.environ.get('QDIV0_ACCOUNT_ID')
    if not api_key or not account_id:
        raise ValueError('Set QDIV0_API_KEY and QDIV0_ACCOUNT_ID first.')
    body, content_type = multipart(fields)
    request = Request('https://api.qdiv0.com/v1/videos', data=body, method='POST', headers={
        'Authorization': f'Bearer {api_key}', 'X-Account-Id': account_id,
        'Content-Type': content_type, 'Accept': 'video/mp4',
    })
    output.parent.mkdir(parents=True, exist_ok=True)
    # Reserve the output before making a billable request, including concurrent runs.
    with output.open('xb') as file:
        try:
            requested_at = datetime.now(timezone.utc).isoformat()
            started = time.monotonic()
            with urlopen(request, timeout=timeout) as response:
                media_type = response.headers.get('Content-Type', '').split(';')[0].strip()
                prefix = response.read(12)
                if media_type not in {'video/mp4', 'application/octet-stream'} or prefix[4:8] != b'ftyp':
                    raise ValueError('The response is not an MP4. Inspect the deployment before retrying.')
                file.write(prefix)
                shutil.copyfileobj(response, file)
                expected_length = response.headers.get('Content-Length')
                if expected_length is not None and file.tell() != int(expected_length):
                    raise ValueError('The MP4 download is incomplete. Inspect the deployment before retrying.')
        except Exception:
            file.close()
            output.unlink()  # Remove only the incomplete file created by this request.
            raise
    elapsed = time.monotonic() - started
    metadata = {'request': fields, 'requested_at': requested_at,
                'request_seconds': round(elapsed, 3), 'file': output.name}
    with metadata_path.open('x', encoding='utf-8') as file:
        json.dump(metadata, file, ensure_ascii=False, indent=2)
        file.write('\n')
    return metadata


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model', required=True, help='QDivZero serving name')
    parser.add_argument('--example', choices=['tea', 'river'], default='tea')
    parser.add_argument('--prompts', type=Path, default=Path(__file__).with_name('prompts.json'))
    parser.add_argument('--duration', type=int, choices=range(4, 16), default=5, metavar='4..15')
    parser.add_argument('--aspect-ratio', choices=SHAPES, default='16:9')
    parser.add_argument('--steps', type=int, choices=range(1, 201), default=50, metavar='1..200')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--timeout', type=int, default=1800, help='Socket timeout in seconds')
    parser.add_argument('--output', type=Path, required=True, help='New .mp4 filename')
    args = parser.parse_args()
    if args.seed < 0 or args.timeout <= 0:
        parser.error('Seed must be non-negative and timeout must be positive.')
    prompts = json.loads(args.prompts.read_text(encoding='utf-8'))
    width, height = SHAPES[args.aspect_ratio]
    fields = {'model': args.model, 'prompt': prompts[args.example],
              'task': 't2va', 'duration': args.duration,
              'width': width, 'height': height, 'aspect_ratio': args.aspect_ratio,
              'num_inference_steps': args.steps, 'flow_shift': 12, 'audio_flow_shift': 3,
              'seed': args.seed}
    try:
        metadata = generate(fields, args.output, args.timeout)
    except HTTPError as error:
        detail = error.read(2000).decode('utf-8', errors='replace')
        raise SystemExit(f'HTTP {error.code}: {detail}\nInspect the deployment before repeating the request; it may still be generating.') from None
    except (URLError, TimeoutError, OSError, ValueError) as error:
        raise SystemExit(f'{error}\nNo automatic retry was made. Check QDivZero before submitting again.') from None
    print(json.dumps(metadata, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()