#!/usr/bin/env python3
"""Llmer local reference upload and resumable downloads (Python standard library)."""
import argparse
from email.utils import parsedate_to_datetime
import hashlib
import http.client
import json
import os
from pathlib import Path
import re
import random
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

RETRY_ATTEMPTS = 5
RETRY_STATUSES = (408, 429, 500, 502, 503, 504)


def retry_wait(attempt, headers=None, status=None):
    """Bounded retries; never shorten a server's Retry-After instruction."""
    if attempt >= RETRY_ATTEMPTS - 1:
        raise RuntimeError('Retry limit reached; resume the same upload/download later, not generation')
    raw = (headers or {}).get('Retry-After', '').strip()
    delay = None
    if raw:
        try:
            if raw.isdigit():
                delay = int(raw)
            else:
                when = parsedate_to_datetime(raw)
                if when.tzinfo is not None:
                    delay = max(0, when.timestamp() - time.time())
        except (ValueError, TypeError, OverflowError):
            pass
    if delay is None:
        # Existing upload limiter uses a 60-second window, sometimes without a header.
        delay = 60 if status == 429 else min(30, 2 ** attempt)
    if delay > 300:
        raise RuntimeError('Server requested a wait over 300 seconds; retry this operation later, not generation')
    delay += random.uniform(0, 0.5)
    print('Transient media error; retrying the same operation in %.1fs' % delay, file=sys.stderr)
    time.sleep(delay)


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

class LlmerMedia:
    def __init__(self, base=None, key=None):
        self.base = (base or os.environ.get('LLMER_BASE_URL', 'https://ry.live')).rstrip('/')
        if self.base.endswith('/v1'):
            self.base = self.base[:-3]
        parsed = urllib.parse.urlsplit(self.base)
        if parsed.scheme != 'https' or not parsed.netloc or parsed.username or parsed.query or parsed.fragment:
            raise ValueError('LLMER_BASE_URL must be an HTTPS service origin')
        self.key = key or os.environ.get('LLMER_API_KEY', '')
        if not self.key:
            raise ValueError('Set LLMER_API_KEY in the environment')
        self.opener = urllib.request.build_opener(NoRedirect())

    def api(self, path, body=None):
        request = urllib.request.Request(self.base + path,
            data=None if body is None else json.dumps(body).encode(),
            headers={'Authorization': 'Bearer ' + self.key, 'Content-Type': 'application/json'})
        route = urllib.parse.urlsplit(path).path
        create_upload = body is not None and route == '/v1/assets/uploads'
        download_auth = bool(re.fullmatch(r'/v1/(?:assets/[^/]+|videos/[^/]+|tasks/[^/]+/artifacts/[^/]+)/download-url', route))
        retry_safe = body is None or download_auth or bool(re.fullmatch(r'/v1/assets/uploads/[^/]+/complete', route))
        # Every other POST, including generation, gets exactly one network attempt.
        for attempt in range(RETRY_ATTEMPTS):
            try:
                with self.opener.open(request, timeout=90) as response:
                    result = json.load(response)
                    if response.status == 202 and download_auth:
                        retry_wait(attempt, response.headers, 202)
                        continue
                    return result
            except urllib.error.HTTPError as exc:
                status, headers = exc.code, exc.headers
                exc.close()
                if (retry_safe and status in RETRY_STATUSES) or (create_upload and status == 429):
                    retry_wait(attempt, headers, status)
                    continue
                raise RuntimeError('Llmer API returned HTTP ' + str(status) + '; no generation was resubmitted') from None
            except (OSError, EOFError, http.client.HTTPException):
                if retry_safe:
                    retry_wait(attempt)
                    continue
                raise RuntimeError('Llmer API connection failed; reconcile any submission before retrying') from None
            except ValueError:
                raise RuntimeError('Invalid Llmer API response; reconcile any submission before retrying') from None

    def upload_reference(self, local_path):
        path = Path(local_path)
        if path.stat().st_size > 10 << 20:
            raise ValueError('Reference exceeds 10 MiB')
        data = path.read_bytes()
        mime = 'image/png' if data.startswith(b'\x89PNG\r\n\x1a\n') else 'image/jpeg' if data.startswith(b'\xff\xd8\xff') else None
        if not mime:
            raise ValueError('Use a PNG or JPEG reference')
        upload = self.api('/v1/assets/uploads', {'size_bytes': len(data), 'mime_type': mime})
        if urllib.parse.urlsplit(upload['upload_url']).scheme != 'https':
            raise RuntimeError('Upload endpoint must use HTTPS')
        # This request carries only its dedicated upload authorization, no Llmer key.
        request = urllib.request.Request(upload['upload_url'], data=data, method='PUT', headers=upload['headers'])
        for attempt in range(RETRY_ATTEMPTS):
            try:
                with self.opener.open(request, timeout=120) as response:
                    if response.status != 200:
                        raise RuntimeError('Reference upload failed')
                break
            except urllib.error.HTTPError as exc:
                status, headers = exc.code, exc.headers
                exc.close()
                if status in RETRY_STATUSES:
                    retry_wait(attempt, headers, status)
                    continue
                raise RuntimeError('Reference PUT returned HTTP ' + str(status) + '; no generation was submitted') from None
            except (OSError, EOFError, http.client.HTTPException):
                retry_wait(attempt)
        asset = self.api('/v1/assets/uploads/' + upload['upload_id'] + '/complete', {})
        if asset['sha256'] != hashlib.sha256(data).hexdigest():
            raise RuntimeError('Uploaded reference integrity check failed')
        return asset

    def prepare_h3(self, prompt, first=None, last=None, references=None, resolution='768P'):
        references = references or []
        if references and (first or last):
            raise ValueError('Choose endpoint frames or reference images')
        if len(references) > 4:
            raise ValueError('At most four reference images')
        paths = [p for p in [first, last] if p] + references
        if sum(Path(p).stat().st_size for p in paths) > 20 << 20:
            raise ValueError('Combined references exceed 20 MiB')
        body = {'model': 'minimax-h3-local/image-to-video' if paths else 'minimax-h3-local/text-to-video',
                'prompt': prompt, 'seconds': 5, 'steps': 8, 'resolution': resolution}
        if first:
            body['first_frame_asset_id'] = self.upload_reference(first)['id']
        if last:
            body['last_frame_asset_id'] = self.upload_reference(last)['id']
        if references:
            body['reference_asset_ids'] = [self.upload_reference(p)['id'] for p in references]
        return body

    def submit_once(self, body, receipt):
        receipt = Path(receipt)
        # An uncertain submission requires reconciliation, never automatic resubmission.
        with receipt.open('x') as f:
            json.dump({'submission_state': 'unconfirmed', 'created_at': int(time.time())}, f)
            f.flush(); os.fsync(f.fileno())
        result = self.api('/v1/videos', body)
        if not result.get('id'):
            raise RuntimeError('No confirmed task ID; reconcile the receipt before retrying')
        save(receipt, result)
        return result

    def get_video_status(self, task_id):
        if not re.fullmatch(r'[A-Za-z0-9_.~-]+', task_id):
            raise ValueError('Invalid task ID')
        return self.api('/v1/videos/' + task_id)

    def download(self, task_id, output, artifact='video'):
        if not re.fullmatch(r'[A-Za-z0-9_.~-]+', task_id) or not re.fullmatch(r'[A-Za-z0-9_.~-]+', artifact):
            raise ValueError('Invalid task or artifact ID')
        output = Path(output); part = output.with_name(output.name + '.part'); meta = output.with_name(output.name + '.part.json')
        identity = {'task_id': task_id, 'artifact': artifact}
        if part.exists() and not meta.exists():
            raise RuntimeError('Partial file has no identity record; move it aside before retrying')
        previous = json.loads(meta.read_text()) if meta.exists() else identity
        if any(previous.get(k) != v for k, v in identity.items()):
            raise RuntimeError('Partial download belongs to another artifact')
        for attempt in range(RETRY_ATTEMPTS):
            info = self.api('/v1/tasks/' + task_id + '/artifacts/' + artifact + '/download-url?disposition=attachment', {})
            if info.get('artifact_status') != 'ready':
                raise RuntimeError('Result is still publishing; retry download later, not generation')
            if urllib.parse.urlsplit(info['url']).scheme != 'https':
                raise RuntimeError('Download endpoint must use HTTPS')
            size = info['size_bytes']; offset = part.stat().st_size if part.exists() else 0
            if offset > size:
                raise RuntimeError('Partial file is larger than the result')
            headers = {}
            if offset:
                headers['Range'] = 'bytes=' + str(offset) + '-'
                if previous.get('etag'):
                    headers['If-Range'] = previous['etag']
            try:
                # Separate request: never forward the Llmer Authorization header.
                with self.opener.open(urllib.request.Request(info['url'], headers=headers), timeout=120) as response:
                    if response.status == 206:
                        if not response.headers.get('Content-Range', '').startswith('bytes ' + str(offset) + '-'):
                            raise RuntimeError('Invalid resumed byte range')
                        mode = 'ab'
                    elif response.status == 200:
                        mode = 'wb'
                    else:
                        raise RuntimeError('Unexpected media response')
                    previous = {**identity, 'etag': response.headers.get('ETag'), 'size_bytes': size}
                    save(meta, previous)
                    with part.open(mode) as f:
                        while True:
                            try:
                                chunk = response.read(1024 * 1024)
                            except http.client.IncompleteRead as exc:
                                f.write(exc.partial)
                                raise
                            if not chunk:
                                break
                            f.write(chunk)
                        f.flush(); os.fsync(f.fileno())
            except urllib.error.HTTPError as exc:
                status, error_headers = exc.code, exc.headers
                exc.close()
                if status == 416 and offset == size and previous.get('size_bytes') == size:
                    pass
                elif status == 403 or status in RETRY_STATUSES:
                    retry_wait(attempt, error_headers, status)
                    continue
                else:
                    raise RuntimeError('Media download returned HTTP ' + str(status)) from None
            except (OSError, EOFError, http.client.HTTPException):
                retry_wait(attempt)
                continue
            if part.stat().st_size != size:
                retry_wait(attempt)
                continue
            sha = hashlib.sha256()
            with part.open('rb') as f:
                while chunk := f.read(1024 * 1024): sha.update(chunk)
            os.replace(part, output); meta.unlink(missing_ok=True)
            return {'path': str(output), 'size_bytes': size, 'sha256': sha.hexdigest()}
        raise RuntimeError('Incomplete download retained; rerun download to resume')

def save(path, data):
    path = Path(path); temporary = path.with_name(path.name + '.tmp')
    with temporary.open('w') as f:
        json.dump(data, f, ensure_ascii=False, indent=2); f.flush(); os.fsync(f.fileno())
    os.replace(temporary, path)

def main():
    os.umask(0o077)
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest='command', required=True)
    upload = commands.add_parser('upload'); upload.add_argument('file')
    prepare = commands.add_parser('prepare'); prepare.add_argument('--prompt', required=True)
    prepare.add_argument('--first-frame'); prepare.add_argument('--last-frame'); prepare.add_argument('--reference', action='append', default=[])
    prepare.add_argument('--resolution', choices=['480P', '768P', '1080P'], default='768P'); prepare.add_argument('--output', default='request.json')
    submit = commands.add_parser('submit'); submit.add_argument('request'); submit.add_argument('--receipt', required=True)
    status = commands.add_parser('status'); status.add_argument('task_id')
    download = commands.add_parser('download'); download.add_argument('task_id'); download.add_argument('--output', required=True); download.add_argument('--artifact', default='video')
    args = parser.parse_args(); client = LlmerMedia()
    if args.command == 'upload': result = client.upload_reference(args.file)
    elif args.command == 'prepare':
        result = client.prepare_h3(args.prompt, args.first_frame, args.last_frame, args.reference, args.resolution)
        save(args.output, result); result = {'request_file': args.output, 'generation_submitted': False}
    elif args.command == 'submit': result = client.submit_once(json.loads(Path(args.request).read_text()), args.receipt)
    elif args.command == 'status': result = client.get_video_status(args.task_id)
    else: result = client.download(args.task_id, args.output, args.artifact)
    print(json.dumps(result, ensure_ascii=False, indent=2))

if __name__ == '__main__':
    main()