"""Mumu API examples. No automatic POST retries. Install: pip install requests."""
import argparse
from contextlib import ExitStack
import getpass
import mimetypes
import os
from pathlib import Path
import re
import time
from urllib.parse import urljoin, urlsplit

import requests

ORIGIN = "https://api.mumugofe.com"
BASE = ORIGIN + "/v1"


def checked_task_id(value):
    if not isinstance(value, str) or not re.fullmatch(r"task_[A-Za-z0-9_-]{1,180}", value):
        raise ValueError("Response has no valid task ID; check the task log before submitting again.")
    return value


def json_result(response):
    try:
        body = response.json()
    except ValueError:
        raise RuntimeError(f"HTTP {response.status_code}: non-JSON response; do not resubmit blindly.") from None
    if not 200 <= response.status_code < 300:
        task_id = body.get("task_id") if isinstance(body, dict) else None
        if task_id:
            print("Task ID (resume this task):", checked_task_id(task_id))
        raise RuntimeError(f"HTTP {response.status_code}: check the console task log. No POST retry performed.")
    if not isinstance(body, dict):
        raise RuntimeError("Unexpected response shape")
    return body


def make_request(args):
    """Build documented public parameters without making requests."""
    if args.kind == "text":
        return "/chat/completions", {"model": args.model, "messages": [{"role": "user", "content": args.prompt}], "max_tokens": 256, "stream": False}
    payload = {"model": args.model, "prompt": args.prompt}
    if args.kind in ("image", "edit"):
        payload.update({"n": 1, "async": True, "response_format": "url"})
        if args.size:
            payload["size"] = args.size
        if args.kind == "image" and (args.image or args.image_url):
            raise ValueError("Use --kind edit for reference images.")
        if args.image_url:
            raise ValueError("This image-edit example uses local files; URL editing has not been accepted for every model.")
        if args.kind == "edit" and not args.image:
            raise ValueError("--kind edit requires --image")
        return "/images/edits" if args.kind == "edit" else "/images/generations", payload
    if args.seconds is None or not 1 <= args.seconds <= 3600:
        raise ValueError("Video requests require an explicit positive --seconds value.")
    if args.model == "sd25-30s":
        if args.seconds != 30 or len(args.image) + len(args.image_url) > 30:
            raise ValueError("sd25-30s requires 30 seconds and at most 30 reference images.")
        if args.size:
            payload["size"] = args.size
        if args.image_url:
            payload["images"] = args.image_url
    else:
        if len(args.image) + len(args.image_url) > 1:
            raise ValueError("This general-video example only supports one reference image; use sd25-30s for documented multi-image input.")
        payload["size"] = args.size or "1280x720"
        if args.image_url:
            payload["input_reference"] = args.image_url[0]
    payload["seconds"] = args.seconds
    return "/videos", payload


def submit(args, headers):
    path, payload = make_request(args)
    with ExitStack() as stack:
        if args.image:
            paths = [Path(name) for name in args.image]
            if sum(p.stat().st_size for p in paths) >= 30 * 1024 * 1024:
                raise ValueError("Keep uploaded files below 30 MiB combined (gateway body limit is 32 MiB).")
            if args.model == "sd25-30s" and any(p.stat().st_size > 20 * 1024 * 1024 for p in paths):
                raise ValueError("Each sd25-30s reference image must be at most 20 MiB.")
            field = "image[]" if args.kind == "edit" else "input_reference"
            files = [(field, (p.name, stack.enter_context(p.open("rb")), mimetypes.guess_type(p.name)[0] or "application/octet-stream")) for p in paths]
            # Do not combine URL arrays with multipart in this sample.
            if args.image_url:
                raise ValueError("Choose local files OR public URLs in this example.")
            form = {k: str(v).lower() if isinstance(v, bool) else str(v) for k, v in payload.items()}
            response = requests.post(BASE + path, headers=headers, data=form, files=files, timeout=(15, 180), allow_redirects=False)
        else:
            response = requests.post(BASE + path, headers=headers, json=payload, timeout=(15, 180), allow_redirects=False)
    return json_result(response)


def poll(kind, task_id, headers, wait_seconds):
    route = "/images/generations/tasks/" if kind in ("image", "edit") else "/videos/"
    deadline = time.monotonic() + wait_seconds
    while time.monotonic() < deadline:
        try:
            response = requests.get(BASE + route + task_id, headers=headers, timeout=(15, 60), allow_redirects=False)
            if response.status_code == 429 or response.status_code >= 500:
                time.sleep(10)
                continue
            body = json_result(response)
        except requests.RequestException:
            time.sleep(10)
            continue
        status = body.get("status")
        print("Status:", status, "Progress:", body.get("progress"))
        if status == "completed":
            return body
        if status in ("failed", "cancelled"):
            raise RuntimeError("Task failed. Check settlement in the console; this script does not resubmit.")
        time.sleep(5)
    raise TimeoutError("Stopped polling, not cancelled. Resume with --task-id " + task_id)


def download(kind, task_id, result, headers, filename):
    if kind in ("image", "edit"):
        url = urljoin(ORIGIN, result["data"][0]["url"])
        # Signed image URLs are credentials: do not print or publish them.
        download_headers = {}
    else:
        url = BASE + "/videos/" + task_id + "/content"
        download_headers = headers
    parsed = urlsplit(url)
    if parsed.scheme != "https" or parsed.netloc != "api.mumugofe.com" or parsed.username:
        raise ValueError("Refusing a non-site result URL; no API key sent.")
    with requests.get(url, headers=download_headers, timeout=(15, 180), stream=True, allow_redirects=False) as response:
        if response.status_code != 200:
            raise RuntimeError(f"Download HTTP {response.status_code}. Query the same task again; do not regenerate.")
        if not response.headers.get("Content-Type", "").startswith(("image/", "video/")):
            raise RuntimeError("Response is not image/video media; refusing to save it as a finished product.")
        # xb refuses to overwrite existing files. An interrupted file may be partial.
        with open(filename, "xb") as output:
            for chunk in response.iter_content(1024 * 1024):
                output.write(chunk)
    print("Saved:", filename)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--kind", choices=["text", "image", "edit", "video"], required=True)
    parser.add_argument("--model")
    parser.add_argument("--prompt")
    parser.add_argument("--image", action="append", default=[])
    parser.add_argument("--image-url", action="append", default=[])
    parser.add_argument("--seconds", type=int)
    parser.add_argument("--size")
    parser.add_argument("--task-id", help="Resume observation only; never creates a new generation")
    parser.add_argument("--out", default="result.bin", help="Use an unused filename; extension does not convert media")
    parser.add_argument("--wait", type=int, default=1800)
    args = parser.parse_args()
    if args.wait <= 0:
        parser.error("--wait must be positive")
    if args.kind != "text" and Path(args.out).exists():
        parser.error("Output file already exists; choose a new --out before submitting")
    if not args.task_id and (not args.model or not args.prompt):
        parser.error("New submissions require --model and --prompt")
    if args.task_id and args.kind == "text":
        parser.error("Text chat does not use these media task endpoints")
    if not args.task_id:
        make_request(args)
    key = os.environ.get("MU_API_KEY") or getpass.getpass("Your relay API key (hidden): ")
    if not key.strip():
        parser.error("API key is empty")
    headers = {"Authorization": "Bearer " + key.strip()}
    if args.task_id:
        task_id = checked_task_id(args.task_id)
    else:
        result = submit(args, headers)
        if args.kind == "text":
            print(result["choices"][0]["message"]["content"])
            return
        task_id = checked_task_id(result.get("task_id") or result.get("id"))
        print("Task ID (save this):", task_id)
    result = poll(args.kind, task_id, headers, args.wait)
    download(args.kind, task_id, result, headers, args.out)


if __name__ == "__main__":
    try:
        main()
    except requests.RequestException:
        raise SystemExit("Network error. Submission may have succeeded: check the task log before retrying. No automatic POST retry.") from None
    except (ValueError, RuntimeError, TimeoutError, OSError, KeyError, IndexError) as error:
        raise SystemExit(str(error)) from None