#!/usr/bin/env python3
"""Download every relevant public API artifact for one Onshape Part Studio.

The output is a self-describing directory for developing and validating an
Onshape-feature-tree-to-CDSL converter.  It intentionally collects both the
editable source representation (features, sketches, parameters and queries)
and independent reconstruction targets (STEP, Parasolid, meshes, topology and
mass properties).  Credentials are read from environment variables or hidden
terminal prompts and are never written to the output directory.
"""

from __future__ import annotations

import argparse
import base64
import getpass
import hashlib
import json
import os
import ssl
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Callable

from download_onshape_samples import PartStudioRef, parse_part_studio_url


API_VERSION = "v17"
DEFAULT_OUTPUT = Path("json_to_cdsl/input/onshape_complete")
DEFAULT_SAMPLE_ID = "00000352"
DEFAULT_SAMPLE_URL = (
    "https://cad.onshape.com/documents/4185972a944744d8a7a0f2b4/"
    "w/d82d7eef8edf4342b7e49732/e/b6d6b562e8b64e7ea50d8325"
)
JSON_ACCEPT = "application/json;charset=UTF-8; qs=0.09"
USER_AGENT = "cdsl-cad-onshape-complete-downloader/1.0"


@dataclass
class DownloadResult:
    name: str
    status: str
    path: str | None = None
    url: str | None = None
    bytes: int | None = None
    sha256: str | None = None
    http_status: int | None = None
    content_type: str | None = None
    message: str | None = None


class OnshapeClient:
    def __init__(self, ref: PartStudioRef, authorization: str, timeout: float) -> None:
        self.ref = ref
        self.authorization = authorization
        self.timeout = timeout
        self.context = ssl.create_default_context(cafile=self._ca_bundle())

    @staticmethod
    def _ca_bundle() -> str | None:
        try:
            import certifi
        except ImportError:
            return None
        return certifi.where()

    def url(self, path: str, query: dict[str, Any] | None = None) -> str:
        encoded = urllib.parse.urlencode(query or {}, doseq=True)
        suffix = f"?{encoded}" if encoded else ""
        return f"https://{self.ref.stack}/api/{API_VERSION}{path}{suffix}"

    def request(
        self,
        method: str,
        path: str,
        query: dict[str, Any] | None = None,
        body: dict[str, Any] | None = None,
        accept: str = JSON_ACCEPT,
    ) -> tuple[bytes, str | None, str]:
        payload = None if body is None else json.dumps(body).encode("utf-8")
        request = urllib.request.Request(
            self.url(path, query),
            data=payload,
            method=method,
            headers={
                "Accept": accept,
                "Authorization": self.authorization,
                "Content-Type": JSON_ACCEPT,
                "User-Agent": USER_AGENT,
            },
        )
        with urllib.request.urlopen(request, timeout=self.timeout, context=self.context) as response:
            return response.read(), response.headers.get_content_type(), response.geturl()

    def get_json(self, path: str, query: dict[str, Any] | None = None) -> Any:
        raw, _, _ = self.request("GET", path, query)
        return json.loads(raw.decode("utf-8"))


def credentials() -> tuple[str, str]:
    access_key = os.environ.get("ONSHAPE_ACCESS_KEY") or getpass.getpass("Onshape access key: ")
    secret_key = os.environ.get("ONSHAPE_SECRET_KEY") or getpass.getpass("Onshape secret key: ")
    if not access_key or not secret_key:
        raise ValueError("both Onshape access and secret keys are required")
    return access_key, secret_key


def write_json(path: Path, value: Any) -> None:
    path.write_text(json.dumps(value, ensure_ascii=True, indent=2) + "\n", encoding="utf-8")


def file_info(path: Path) -> tuple[int, str]:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return path.stat().st_size, digest.hexdigest()


def relative_path(path: Path, root: Path) -> str:
    return path.relative_to(root).as_posix()


def save_json(
    client: OnshapeClient,
    root: Path,
    name: str,
    path: str,
    filename: str,
    query: dict[str, Any] | None = None,
) -> DownloadResult:
    try:
        value = client.get_json(path, query)
        destination = root / filename
        destination.parent.mkdir(parents=True, exist_ok=True)
        write_json(destination, value)
        size, digest = file_info(destination)
        return DownloadResult(name, "downloaded", relative_path(destination, root), client.url(path, query), size, digest)
    except urllib.error.HTTPError as exc:
        return DownloadResult(name, "http_error", url=client.url(path, query), http_status=exc.code, message=exc.read().decode("utf-8", "replace")[:1000])
    except (urllib.error.URLError, TimeoutError, OSError, ValueError) as exc:
        return DownloadResult(name, "error", url=client.url(path, query), message=str(exc))


def save_binary(
    client: OnshapeClient,
    root: Path,
    name: str,
    path: str,
    filename: str,
    query: dict[str, Any] | None = None,
    accept: str = "application/octet-stream",
) -> DownloadResult:
    try:
        payload, content_type, final_url = client.request("GET", path, query, accept=accept)
        destination = root / filename
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_bytes(payload)
        size, digest = file_info(destination)
        return DownloadResult(name, "downloaded", relative_path(destination, root), final_url, size, digest, content_type=content_type)
    except urllib.error.HTTPError as exc:
        return DownloadResult(name, "http_error", url=client.url(path, query), http_status=exc.code, message=exc.read().decode("utf-8", "replace")[:1000])
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        return DownloadResult(name, "error", url=client.url(path, query), message=str(exc))


def part_studio_path(ref: PartStudioRef, suffix: str) -> str:
    return f"/partstudios/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}{suffix}"


def collect_sketch_artifacts(
    client: OnshapeClient,
    root: Path,
    sketches: Any,
) -> list[DownloadResult]:
    if not isinstance(sketches, list):
        return []
    results: list[DownloadResult] = []
    seen_ids: set[str] = set()
    for sketch in sketches:
        if not isinstance(sketch, dict):
            continue
        sketch_id = sketch.get("sketchId") or sketch.get("featureId")
        if not isinstance(sketch_id, str) or not sketch_id or sketch_id in seen_ids:
            continue
        seen_ids.add(sketch_id)
        quoted_id = urllib.parse.quote(sketch_id, safe="")
        base = part_studio_path(client.ref, f"/sketches/{quoted_id}")
        results.append(
            save_json(client, root, f"sketch/{sketch_id}/bounding_box", f"{base}/boundingboxes", f"sketches/{sketch_id}/bounding_box.json")
        )
        results.append(
            save_json(client, root, f"sketch/{sketch_id}/tessellation", f"{base}/tessellatedentities", f"sketches/{sketch_id}/tessellated_entities.json")
        )
    return results


def export_step(client: OnshapeClient, root: Path, sample_id: str, poll_seconds: float, poll_limit: int) -> DownloadResult:
    path = part_studio_path(client.ref, "/export/step")
    body = {
        "destinationName": f"{sample_id}.step",
        "grouping": True,
        "notifyUser": False,
        "storeInDocument": False,
        "triggerAutoDownload": False,
        "stepUnit": "METER",
        "stepVersionString": "AP242",
    }
    try:
        raw, _, _ = client.request("POST", path, body=body)
        translation = json.loads(raw.decode("utf-8"))
        translation_id = translation.get("id")
        if not isinstance(translation_id, str) or not translation_id:
            return DownloadResult("step", "error", url=client.url(path), message="STEP export response has no translation id")
        request_path = "/translations/" + urllib.parse.quote(translation_id, safe="")
        for _ in range(poll_limit):
            state = client.get_json(request_path)
            request_state = state.get("requestState") if isinstance(state, dict) else None
            if request_state == "DONE":
                external_ids = state.get("resultExternalDataIds", [])
                if not isinstance(external_ids, list) or not external_ids:
                    return DownloadResult("step", "error", url=client.url(request_path), message="completed STEP export has no external data id")
                external_id = external_ids[0]
                if not isinstance(external_id, str):
                    return DownloadResult("step", "error", url=client.url(request_path), message="invalid STEP external data id")
                return save_binary(
                    client,
                    root,
                    "step",
                    f"/documents/d/{client.ref.did}/externaldata/{urllib.parse.quote(external_id, safe='')}",
                    "model.step",
                    accept="application/step, application/octet-stream",
                )
            if request_state == "FAILED":
                return DownloadResult("step", "http_error", url=client.url(request_path), message=str(state.get("failureReason", "translation failed")))
            time.sleep(poll_seconds)
        return DownloadResult("step", "error", url=client.url(request_path), message=f"STEP translation did not finish after {poll_limit} polls")
    except urllib.error.HTTPError as exc:
        return DownloadResult("step", "http_error", url=client.url(path), http_status=exc.code, message=exc.read().decode("utf-8", "replace")[:1000])
    except (urllib.error.URLError, TimeoutError, OSError, ValueError) as exc:
        return DownloadResult("step", "error", url=client.url(path), message=str(exc))


def export_stl(client: OnshapeClient, root: Path) -> DownloadResult:
    """Export a Part Studio STL, retrying the two API-supported encodings.

    Some older public documents reject one encoding even though the Part Studio
    can otherwise be read.  The retry is recorded in the URL and the manifest
    retains a useful error if neither server-side export succeeds.
    """
    path = part_studio_path(client.ref, "/stl")
    attempts = (
        {"mode": "binary", "grouping": "true", "units": "METER"},
        {"mode": "text", "grouping": "true", "units": "METER"},
        {"mode": "binary", "grouping": "false", "units": "METER"},
    )
    failures: list[str] = []
    for query in attempts:
        result = save_binary(client, root, "stl", path, "model.stl", query)
        if result.status == "downloaded":
            return result
        detail = result.message or result.status
        failures.append(f"{urllib.parse.urlencode(query)}: {detail}")
    return DownloadResult("stl", "error", url=client.url(path), message=" | ".join(failures))


def save_shaded_preview(root: Path) -> DownloadResult:
    """Extract the API's base64 PNG view as a locally inspectable preview."""
    shaded_path = root / "shaded_views.json"
    try:
        value = json.loads(shaded_path.read_text(encoding="utf-8"))
        images = value.get("images", []) if isinstance(value, dict) else []
        encoded = images[0] if isinstance(images, list) and images else None
        if not isinstance(encoded, str):
            return DownloadResult("shaded_preview", "error", message="shaded view response contains no PNG")
        payload = base64.b64decode(encoded, validate=True)
        if not payload.startswith(b"\x89PNG\r\n\x1a\n"):
            return DownloadResult("shaded_preview", "error", message="shaded view response is not a PNG")
        destination = root / "shaded_preview.png"
        destination.write_bytes(payload)
        size, digest = file_info(destination)
        return DownloadResult("shaded_preview", "downloaded", relative_path(destination, root), "shaded_views.json#images[0]", size, digest, content_type="image/png")
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        return DownloadResult("shaded_preview", "error", message=str(exc))


def complete_download(ref: PartStudioRef, authorization: str, output: Path, timeout: float, poll_seconds: float, poll_limit: int) -> dict[str, Any]:
    root = output / ref.sample_id
    root.mkdir(parents=True, exist_ok=True)
    client = OnshapeClient(ref, authorization, timeout)
    request_info = {
        **asdict(ref),
        "api_version": API_VERSION,
        "download_contract": "all relevant data exposed by the public API for this Part Studio, not an internal Onshape document backup",
    }
    write_json(root / "request.json", request_info)

    common = {"rollbackBarIndex": -1}
    resources: list[tuple[str, str, str, dict[str, Any] | None]] = [
        ("document", f"/documents/{ref.did}", "document.json", None),
        ("workspaces", f"/documents/d/{ref.did}/workspaces", "workspaces.json", None),
        ("elements", f"/documents/d/{ref.did}/{ref.wvm}/{ref.wvmid}/elements", "elements.json", None),
        ("unit_info", f"/documents/d/{ref.did}/{ref.wvm}/{ref.wvmid}/unitinfo", "unit_info.json", None),
        ("configuration", f"/elements/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}/configuration", "configuration.json", None),
        ("parts", f"/parts/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}", "parts.json", {"withThumbnails": "true", "includeFlatParts": "true"}),
        ("features", part_studio_path(ref, "/features"), "features.json", {**common, "includeGeometryIds": "true", "noSketchGeometry": "false"}),
        ("featurescript_representation", part_studio_path(ref, "/featurescriptrepresentation"), "featurescript_representation.json", common),
        ("feature_specs", part_studio_path(ref, "/featurespecs"), "feature_specs.json", None),
        ("body_details", part_studio_path(ref, "/bodydetails"), "body_details.json", {**common, "includeSurfaces": "true", "includeCompositeParts": "true", "includeGeometricData": "true"}),
        ("bounding_boxes", part_studio_path(ref, "/boundingboxes"), "bounding_boxes.json", {"includeHidden": "true", "includeWireBodies": "true"}),
        ("mass_properties", part_studio_path(ref, "/massproperties"), "mass_properties.json", {**common, "massAsGroup": "true"}),
        ("sketches", part_studio_path(ref, "/sketches"), "sketches.json", {"includeGeometry": "true", "output3D": "true", "curvePoints": "true"}),
        ("named_views", f"/partstudios/d/{ref.did}/e/{ref.eid}/namedViews", "named_views.json", None),
        ("tessellated_faces", part_studio_path(ref, "/tessellatedfaces"), "tessellated_faces.json", {**common, "outputVertexNormals": "true", "outputFacetNormals": "true", "outputIndexTable": "true", "outputErrorFaces": "true"}),
        ("tessellated_edges", part_studio_path(ref, "/tessellatededges"), "tessellated_edges.json", common),
        ("shaded_views", part_studio_path(ref, "/shadedviews"), "shaded_views.json", {"viewMatrix": "front", "outputWidth": 512, "outputHeight": 512, "edges": "show", "showAllParts": "true", "includeSurfaces": "true", "useAntiAliasing": "true"}),
    ]
    results = [save_json(client, root, *resource) for resource in resources]
    sketches_result = next((item for item in results if item.name == "sketches" and item.status == "downloaded"), None)
    if sketches_result and sketches_result.path:
        sketches = json.loads((root / sketches_result.path).read_text(encoding="utf-8"))
        results.extend(collect_sketch_artifacts(client, root, sketches))

    results.extend(
        [
            save_binary(client, root, "parasolid", part_studio_path(ref, "/parasolid"), "model.x_t", {"version": "0", "includeExportIds": "true", "binaryExport": "false"}, "text/plain, application/octet-stream"),
            export_stl(client, root),
            save_binary(client, root, "gltf", part_studio_path(ref, "/gltf"), "model.gltf", {**common, "outputSeparateFaceNodes": "true", "outputFaceAppearances": "true"}, "model/gltf+json, model/gltf-binary, application/octet-stream"),
            save_binary(client, root, "thumbnail", f"/thumbnails/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}/s/512x512", "thumbnail.png", {"rejectEmpty": "true"}, "image/png, image/*, application/octet-stream"),
        ]
    )
    results.append(save_shaded_preview(root))
    results.append(export_step(client, root, ref.sample_id, poll_seconds, poll_limit))
    manifest = {
        "schema": "onshape.complete_sample.v1",
        "source": request_info,
        "resource_count": len(results),
        "downloaded_count": sum(item.status == "downloaded" for item in results),
        "unavailable_count": sum(item.status != "downloaded" for item in results),
        "resources": [asdict(item) for item in results],
    }
    write_json(root / "manifest.json", manifest)
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default=DEFAULT_SAMPLE_URL, help="Onshape Part Studio URL")
    parser.add_argument("--id", default=DEFAULT_SAMPLE_ID, help="local sample identifier")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUTPUT, help="output root")
    parser.add_argument("--timeout", type=float, default=60.0, help="per-request timeout in seconds")
    parser.add_argument("--poll-seconds", type=float, default=3.0, help="STEP translation poll interval")
    parser.add_argument("--poll-limit", type=int, default=40, help="maximum STEP translation polls")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.poll_seconds <= 0 or args.poll_limit < 1:
        raise ValueError("poll interval must be positive and poll limit must be at least 1")
    ref = parse_part_studio_url(args.id, args.url)
    access_key, secret_key = credentials()
    token = base64.b64encode(f"{access_key}:{secret_key}".encode("utf-8")).decode("ascii")
    manifest = complete_download(ref, f"Basic {token}", args.out, args.timeout, args.poll_seconds, args.poll_limit)
    print(json.dumps({key: manifest[key] for key in ("resource_count", "downloaded_count", "unavailable_count")}, indent=2))
    return 0 if manifest["downloaded_count"] else 1


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        raise SystemExit(2)