#!/usr/bin/env python3
"""Download 3D models (STEP/OBJ) from LCEDA by LCSC part code (Cxxxxx) or component UUID.

Usage:
  python3 lceda_step.py C25744 C2590
  python3 lceda_step.py --obj --outdir ./models C25744 C2590
  python3 lceda_step.py 98c88bc2b9094f218f3a3e30a7012738
"""

import argparse
import json
import re
import sys
import urllib.request

PRODUCT_SEARCH = "https://pro.lceda.cn/api/v2/eda/product/search"
COMPONENT_API = "https://pro.lceda.cn/api/v2/components/{uuid}"
CDN = "https://modules.lceda.cn"
STEP_TOKEN = "qAxj6KHrDKw4blvCG8QJPs7Y"
UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
UUID_RE = re.compile(r"^[0-9a-f]{32}$")


def http_get(url, cookie=None, timeout=60):
    headers = {"User-Agent": UA}
    if cookie:
        headers["Cookie"] = cookie
    req = urllib.request.Request(url, headers=headers)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


def http_post(url, payload, cookie=None, timeout=60):
    headers = {"User-Agent": UA, "Content-Type": "application/json"}
    if cookie:
        headers["Cookie"] = cookie
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers=headers)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read())


def search_products(codes, cookie=None):
    data = http_post(PRODUCT_SEARCH, {"codes": ",".join(codes)}, cookie=cookie)
    return data.get("result", {}).get("productList", [])


def get_component(uuid, cookie=None):
    data = json.loads(http_get(COMPONENT_API.format(uuid=uuid), cookie=cookie))
    if not data.get("success"):
        raise RuntimeError("{}".format(data.get("message", "failed")))
    return data["result"]


def safe_name(name):
    return re.sub(r'[\\/:*?"<>|\s]+', "_", name).strip("_") or "model"


def download(url, dest, cookie=None):
    body = http_get(url, cookie=cookie)
    with open(dest, "wb") as f:
        f.write(body)
    return len(body)


def fetch_step(code, name, model_doc_uuid, owner, outdir, want_obj, cookie):
    try:
        comp = get_component(model_doc_uuid, cookie)
        model = comp.get("3d_model_uuid") or model_doc_uuid
        if not owner:
            owner = comp.get("path", "")
    except Exception:
        model = model_doc_uuid
    if not model:
        print("  {}: no 3D model uuid for {}".format(code, model_doc_uuid))
        return
    fname = safe_name(name)
    step_url = "{}/{}/{}?path={}".format(CDN, STEP_TOKEN, model, owner)
    try:
        n = download(step_url, "{}/{}.step".format(outdir, fname), cookie)
        print("  -> {} STEP {} bytes -> {}.step".format(code, n, fname))
    except Exception as e:
        print("  -> {} STEP download failed: {}".format(code, e))
    if want_obj:
        obj_url = "{}/3dmodel/{}?path={}".format(CDN, model, owner)
        try:
            n = download(obj_url, "{}/{}.obj".format(outdir, fname), cookie)
            print("  -> {} OBJ  {} bytes -> {}.obj".format(code, n, fname))
        except Exception as e:
            print("  -> {} OBJ download failed: {}".format(code, e))


def main():
    ap = argparse.ArgumentParser(description="Download LCEDA 3D models by LCSC code / UUID")
    ap.add_argument("items", nargs="+", help="LCSC part codes (Cxxxxx) or 32-hex component UUIDs")
    ap.add_argument("--obj", action="store_true", help="also download the source OBJ model")
    ap.add_argument("--outdir", default=".", help="output directory (default: current)")
    ap.add_argument("--cookie", default="", help="optional Cookie header for private components")
    args = ap.parse_args()

    codes = [i for i in args.items if not UUID_RE.match(i)]
    uuids = [i for i in args.items if UUID_RE.match(i)]

    if codes:
        try:
            products = search_products(codes, args.cookie)
        except Exception as e:
            print("product search failed: {}".format(e))
            products = []
        found = {}
        for p in products:
            di = p.get("device_info") or {}
            fi = di.get("footprint_info") or {}
            m3 = fi.get("model_3d") or {}
            uri = m3.get("uri") or (di.get("attributes") or {}).get("3D Model")
            owner = (di.get("owner") or {}).get("uuid", "")
            found[p["code"]] = {
                "name": p.get("model") or p.get("code"),
                "uri": uri,
                "owner": owner,
            }
        for code in codes:
            info = found.get(code)
            if not info or not info["uri"]:
                print("{}: no device/3D model info".format(code))
                continue
            print("{} | {} | model_doc={}".format(code, info["name"], info["uri"]))
            try:
                fetch_step(code, info["name"], info["uri"], info["owner"], args.outdir, args.obj, args.cookie)
            except Exception as e:
                print("  {}: {}".format(code, e))

    for uuid in uuids:
        name, model, owner = uuid, uuid, ""
        try:
            comp = get_component(uuid, args.cookie)
            name = comp.get("display_title") or comp.get("title") or uuid
            model = comp.get("3d_model_uuid")
            owner = comp.get("path", "")
        except Exception as e:
            print("{}: component lookup failed ({}), trying uuid as model uuid".format(uuid, e))
        print("{} | {} | model={} | owner={}".format(uuid, name, model, owner))
        if not model:
            print("  no 3D model available")
            continue
        try:
            fetch_step(uuid, name, uuid, owner, args.outdir, args.obj, args.cookie)
        except Exception as e:
            print("  {}: {}".format(uuid, e))


if __name__ == "__main__":
    main()