#!/usr/bin/env python3
"""outworld_assets.py: single-file Python client for the Outworld Open Assets API (v1).

Outworld Open Assets is a free archive of robotics-ready 3D objects (OpenUSD, NVIDIA SimReady
Robotics-Prop 4.0.0, plus a deformable tier), licensed CC BY 4.0. Docs: https://outworld-open-assets.pages.dev/developers

Standard library only (Python 3.8+). Use as a module or from the command line:

    import outworld_assets as oa
    hits = oa.search("ceramic mug", semantic=True, mechanism="rigid")
    meta = oa.get(hits["items"][0]["id"])
    folder = oa.download(meta["id"], "assets/")      # -> assets/<id>/simready_usd/sm_<id>_01.usd

    python outworld_assets.py search "cabinet with doors" --mechanism articulated
    python outworld_assets.py get kitchenware_mug_a01
    python outworld_assets.py download kitchenware_mug_a01 --dest assets/
    python outworld_assets.py categories | stats | status <submission_id>

Set OUTWORLD_ASSETS_URL to use another server (for example https://assets.outworldlabs.com).
Attribution (CC BY 4.0): keep the LICENSE.txt that ships in every asset folder, or credit
"<asset name>" by Outworld Labs (outworldlabs.com), from Outworld Open Assets, CC BY 4.0.
"""
from __future__ import annotations

import hashlib
import json
import os
import random
import shutil
import sys
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path

__version__ = "1.0.0"
DEFAULT_BASE = "https://outworld-open-assets.pages.dev"
USER_AGENT = f"outworld-assets-python/{__version__}"
RETRY_STATUS = {408, 425, 429, 500, 502, 503, 504, 520, 521, 522, 523, 524}


class ApiError(Exception):
    def __init__(self, status: int, message: str, url: str = ""):
        super().__init__(f"HTTP {status}: {message}" + (f" ({url})" if url else ""))
        self.status = status
        self.message = message
        self.url = url


class Client:
    """Client for /api/v1. All methods return parsed JSON (dicts)."""

    def __init__(self, base_url: str | None = None, timeout: float = 60.0, retries: int = 5, backoff: float = 1.0):
        self.base_url = (base_url or os.environ.get("OUTWORLD_ASSETS_URL") or DEFAULT_BASE).rstrip("/")
        self.timeout = timeout
        self.retries = retries
        self.backoff = backoff

    # ---- HTTP with retries -------------------------------------------------
    def _url(self, path: str) -> str:
        return path if path.startswith("http") else self.base_url + path

    def _open(self, url: str, method: str = "GET", data: bytes | None = None, headers: dict | None = None):
        h = {"user-agent": USER_AGENT, "accept": "application/json"}
        h.update(headers or {})
        last = None
        for attempt in range(self.retries + 1):
            req = urllib.request.Request(url, data=data, method=method, headers=h)
            try:
                return urllib.request.urlopen(req, timeout=self.timeout)
            except urllib.error.HTTPError as e:
                body = e.read().decode("utf-8", "replace")
                try:
                    j = json.loads(body)
                    msg = j.get("error", {}).get("message") if isinstance(j.get("error"), dict) else j.get("error")
                except Exception:
                    msg = body[:200]
                last = ApiError(e.code, msg or e.reason, url)
                if e.code not in RETRY_STATUS or attempt == self.retries:
                    raise last from None
                ra = e.headers.get("retry-after")
                wait = float(ra) if ra and ra.replace(".", "", 1).isdigit() else self.backoff * (2 ** attempt)
            except (urllib.error.URLError, TimeoutError, ConnectionError, OSError) as e:
                last = e
                if attempt == self.retries:
                    raise
                wait = self.backoff * (2 ** attempt)
            time.sleep(min(60.0, wait) + random.uniform(0, 0.25))
        raise last  # pragma: no cover

    def _get(self, path: str, params: dict | None = None):
        q = {k: (str(v).lower() if isinstance(v, bool) else v) for k, v in (params or {}).items() if v is not None}
        url = self._url(path) + ("?" + urllib.parse.urlencode(q) if q else "")
        with self._open(url) as r:
            return json.loads(r.read().decode("utf-8"))

    def _post(self, path: str, body: dict):
        with self._open(self._url(path), "POST", json.dumps(body).encode(), {"content-type": "application/json"}) as r:
            return json.loads(r.read().decode("utf-8"))

    # ---- API ---------------------------------------------------------------
    def search(self, q: str | None = None, *, semantic: bool = False, category: str | None = None,
               subcategory: str | None = None, mechanism: str | None = None, articulated: bool | None = None,
               size: str | None = None, size_min: float | None = None, size_max: float | None = None,
               mass_min: float | None = None, mass_max: float | None = None, contributor: str | None = None,
               sort: str | None = None, page: int = 1, per_page: int = 20) -> dict:
        """Search assets. mechanism: rigid|articulated|deformable; size: small|medium|large|xlarge (longest side)."""
        return self._get("/api/v1/search", dict(q=q, semantic=semantic or None, category=category, subcategory=subcategory,
                                                mechanism=mechanism, articulated=articulated, size=size, size_min=size_min,
                                                size_max=size_max, mass_min=mass_min, mass_max=mass_max,
                                                contributor=contributor, sort=sort, page=page, per_page=per_page))

    def iter_search(self, q: str | None = None, **kw):
        """Yield every matching asset summary across all pages."""
        kw.setdefault("per_page", 100)
        page = 1
        while True:
            r = self.search(q, page=page, **kw)
            yield from r["items"]
            if page >= r["pages"]:
                return
            page += 1

    def get(self, asset_id: str) -> dict:
        """Full metadata.json of one asset (plus links)."""
        return self._get(f"/api/v1/assets/{urllib.parse.quote(asset_id)}")

    def downloads(self, asset_id: str) -> dict:
        """Absolute download URLs, file sizes and checksums, and the attribution text."""
        return self._get(f"/api/v1/assets/{urllib.parse.quote(asset_id)}/downloads")

    def categories(self) -> dict:
        return self._get("/api/v1/categories")

    def stats(self) -> dict:
        return self._get("/api/v1/stats")

    def submission_status(self, submission_id: str) -> dict:
        return self._get(f"/api/v1/submissions/{urllib.parse.quote(submission_id)}")

    def download(self, asset_id: str, dest: str | os.PathLike = ".", unpack: bool = True, overwrite: bool = False,
                 verify: bool = True, quiet: bool = False) -> Path:
        """Download an asset zip into dest. With unpack=True returns dest/<id>/ (standard layout:
        <id>/simready_usd/sm_<id>_01.usd, textures/, metadata.json, LICENSE.txt); otherwise the zip path."""
        dest = Path(dest)
        dest.mkdir(parents=True, exist_ok=True)
        info = self.downloads(asset_id)
        target = dest / asset_id
        if unpack and target.exists() and not overwrite and any(target.rglob("*.usd*")):
            return target
        zpath = dest / f"{asset_id}.zip"
        self._fetch(info["zip"]["url"], zpath, info["zip"].get("bytes"), quiet)
        if not unpack:
            return zpath
        tmp = Path(tempfile.mkdtemp(prefix=f".{asset_id}_", dir=dest))
        try:
            with zipfile.ZipFile(zpath) as z:
                for m in z.infolist():
                    p = Path(m.filename)
                    if p.is_absolute() or ".." in p.parts:
                        raise ApiError(0, f"unsafe path in zip: {m.filename}")
                z.extractall(tmp)
            entries = [p for p in tmp.iterdir() if not p.name.startswith("__MACOSX")]
            # Zips hold either <id>/... or the asset folder contents at the root.
            root = entries[0] if len(entries) == 1 and entries[0].is_dir() and entries[0].name != "simready_usd" else tmp
            if target.exists():
                shutil.rmtree(target)
            shutil.move(str(root), str(target))
        finally:
            shutil.rmtree(tmp, ignore_errors=True)
        zpath.unlink()
        # metadata.json and LICENSE.txt are normally in the zip; fetch them if missing.
        for name, url in (("metadata.json", info.get("metadata_json")), ("LICENSE.txt", info.get("license_txt"))):
            if url and not (target / name).exists():
                try:
                    self._fetch(url, target / name, None, True)
                except ApiError:
                    pass
        if verify:
            bad = []
            for f in info.get("files") or []:
                p = target / f["path"]
                if f.get("sha256") and len(f["sha256"]) == 64 and p.exists():
                    if hashlib.sha256(p.read_bytes()).hexdigest() != f["sha256"]:
                        bad.append(f["path"])
            if bad:
                raise ApiError(0, f"checksum mismatch: {bad}")
        return target

    def _fetch(self, url: str, path: Path, expected: int | None, quiet: bool):
        part = path.with_suffix(path.suffix + ".part")
        with self._open(url, headers={"accept": "*/*"}) as r, open(part, "wb") as f:
            total = expected or int(r.headers.get("content-length") or 0)
            done = 0
            while True:
                chunk = r.read(1 << 20)
                if not chunk:
                    break
                f.write(chunk)
                done += len(chunk)
                if not quiet and total and sys.stderr.isatty():
                    sys.stderr.write(f"\r{path.name}: {done / 1e6:.1f}/{total / 1e6:.1f} MB")
            if not quiet and total and sys.stderr.isatty():
                sys.stderr.write("\n")
        if expected and done != expected:
            part.unlink(missing_ok=True)
            raise ApiError(0, f"incomplete download of {url}: {done} of {expected} bytes")
        os.replace(part, path)


_default = None


def _client() -> Client:
    global _default
    if _default is None:
        _default = Client()
    return _default


def search(q=None, **kw): return _client().search(q, **kw)
def get(asset_id): return _client().get(asset_id)
def downloads(asset_id): return _client().downloads(asset_id)
def download(asset_id, dest=".", **kw): return _client().download(asset_id, dest, **kw)
def categories(): return _client().categories()
def stats(): return _client().stats()
def submission_status(submission_id): return _client().submission_status(submission_id)


def _main(argv=None):
    import argparse
    ap = argparse.ArgumentParser(prog="outworld_assets", description="Outworld Open Assets API client")
    ap.add_argument("--base-url", default=None, help=f"default {DEFAULT_BASE} or $OUTWORLD_ASSETS_URL")
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("search", help="search assets")
    s.add_argument("q", nargs="?")
    s.add_argument("--semantic", action="store_true")
    for k in ("category", "subcategory", "mechanism", "size", "contributor", "sort"):
        s.add_argument(f"--{k}")
    s.add_argument("--articulated", choices=["true", "false"])
    s.add_argument("--mass-max", type=float)
    s.add_argument("--page", type=int, default=1)
    s.add_argument("--per-page", type=int, default=20)
    s.add_argument("--json", action="store_true", help="print raw JSON")
    for name in ("get", "downloads"):
        sub.add_parser(name).add_argument("id")
    d = sub.add_parser("download", help="download and unpack one or more assets")
    d.add_argument("ids", nargs="+")
    d.add_argument("--dest", default=".")
    d.add_argument("--zip-only", action="store_true")
    d.add_argument("--overwrite", action="store_true")
    sub.add_parser("categories")
    sub.add_parser("stats")
    sub.add_parser("status").add_argument("id")
    a = ap.parse_args(argv)
    c = Client(a.base_url)
    if a.cmd == "search":
        r = c.search(a.q, semantic=a.semantic, category=a.category, subcategory=a.subcategory, mechanism=a.mechanism,
                     articulated=None if a.articulated is None else a.articulated == "true", size=a.size,
                     mass_max=a.mass_max, contributor=a.contributor, sort=a.sort, page=a.page, per_page=a.per_page)
        if a.json:
            print(json.dumps(r, indent=1))
        else:
            mode = "semantic" if r.get("query", {}).get("semantic") else "keyword"
            print(f"{r['total']} result(s), page {r['page']}/{r['pages']} ({mode})" + (f"  note: {r['note']}" if r.get("note") else ""))
            for it in r["items"]:
                dims = "x".join(f"{x:.3f}" for x in (it.get("dimensions_m") or []))
                print(f"{it['id']:<40} {it['mechanism']:<12} {dims:<20} {it.get('mass_kg') or '':>8} kg  {it['name']}")
    elif a.cmd == "download":
        for i in a.ids:
            print(c.download(i, a.dest, unpack=not a.zip_only, overwrite=a.overwrite))
    else:
        fn = {"get": lambda: c.get(a.id), "downloads": lambda: c.downloads(a.id), "categories": c.categories,
              "stats": c.stats, "status": lambda: c.submission_status(a.id)}[a.cmd]
        print(json.dumps(fn(), indent=1))


if __name__ == "__main__":
    try:
        _main()
    except ApiError as e:
        sys.exit(f"error: {e}")