"""Python 3.10+: upload reference media; this script does not generate a video."""
import argparse
import base64
import json
import os
import sys
import urllib.parse
import urllib.request
from pathlib import Path

ORIGINS = ("https://openpix.cc", "https://api.manxiaobai.online")
MODELS = ("dreamina-seedance-2-0-mini-260615", "dreamina-seedance-2-0-fast-260128", "MiniMax-H3", "SD-2.0-MZ1")
MIMES = {
    ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp",
    ".mp4": "video/mp4", ".mov": "video/quicktime",
    ".mp3": "audio/mpeg", ".wav": "audio/wav", ".m4a": "audio/mp4", ".aac": "audio/aac",
}


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


def prepare(kind, model, path):
    mime = MIMES.get(path.suffix.lower(), "")
    if not mime.startswith(kind + "/"):
        raise ValueError("File extension does not match the selected kind")
    if kind != "image" and model not in MODELS:
        raise ValueError("Video/audio upload requires a documented model")
    if model == "SD-2.0-MZ1" and kind == "video":
        raise ValueError("This model does not accept reference video")
    if kind == "audio" and model != "MiniMax-H3" and path.suffix.lower() not in (".mp3", ".wav"):
        raise ValueError("Use MP3 or PCM WAV for this model")
    limit = 20 * 1024 * 1024 if kind == "image" else (15_000_000 if kind == "audio" else (50_000_000 if model == "MiniMax-H3" else 200_000_000))
    if not 0 < path.stat().st_size <= limit:
        raise ValueError("Empty file or exceeds this upload example's size limit")
    with path.open("rb") as stream:
        data = stream.read(limit + 1)
    if len(data) > limit:
        raise ValueError("File grew beyond the size limit")
    if kind == "image":
        data_url = "data:" + mime + ";base64," + base64.b64encode(data).decode("ascii")
        return "/v1/video-reference-images", "application/json", json.dumps({"image": data_url}).encode()
    return "/v1/video-reference-media/" + urllib.parse.quote(model, safe="") + "/" + kind, mime, data


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base", choices=ORIGINS, default=ORIGINS[0])
    parser.add_argument("--kind", choices=("image", "video", "audio"), required=True)
    parser.add_argument("--model", choices=MODELS)
    parser.add_argument("--file", type=Path, required=True)
    args = parser.parse_args()
    key = os.environ.get("MXB_API_KEY", "")
    if not key or "\r" in key or "\n" in key:
        raise ValueError("Set MXB_API_KEY in your server environment")
    path, mime, data = prepare(args.kind, args.model, args.file)
    request = urllib.request.Request(args.base + path, data=data, headers={
        "Authorization": "Bearer " + key, "Content-Type": mime,
    })
    with urllib.request.build_opener(NoRedirect()).open(request, timeout=120) as response:
        result = json.load(response)
    if not isinstance(result, dict) or not isinstance(result.get("url"), str):
        raise ValueError("Unexpected upload result")
    parsed = urllib.parse.urlsplit(result["url"])
    if parsed.username or parsed.password or parsed.scheme + "://" + parsed.netloc not in ORIGINS:
        raise ValueError("Unexpected upload URL origin")
    # URL is a capability: keep this output private, never publish logs containing it.
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    try:
        main()
    except (OSError, ValueError) as error:
        code = getattr(error, "code", None)
        print("Upload stopped" + (" (HTTP " + str(code) + ")" if code else "") + ". Check the file and model rules; no generation was submitted.", file=sys.stderr)
        sys.exit(1)