#!/usr/bin/env python3
"""Download AP242 STEP files for public Onshape Part Studios only."""

from __future__ import annotations

import argparse
import base64
import getpass
import hashlib
import json
import os
import re
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


API_VERSION = "v17"
DEFAULT_OUTPUT = Path("onshape_to_cdsl/input/raw")
DEFAULT_SAMPLES = (
    (
        "00000001",
        "https://cad.onshape.com/documents/1ffb81a71e5b402e966b9341/"
        "w/6e295017d1b34be684565c40/e/8df255ee6705423d8e85234e",
    ),
    (
        "00000002",
        "https://cad.onshape.com/documents/1ffb81a71e5b402e966b9341/"
        "w/6e295017d1b34be684565c40/e/bb398e4615fe4025b34ea8f0",
    ),
)
ONSHAPE_URL_RE = re.compile(
    r"^https://(?P<stack>[^/]+)/documents/(?P<did>[^/]+)/"
    r"(?P<wvm>w|v|m)/(?P<wvmid>[^/]+)/e/(?P<eid>[^/?#]+)"
)
JSON_ACCEPT = "application/json;charset=UTF-8; qs=0.09"
STEP_ACCEPT = "application/step, application/octet-stream"
USER_AGENT = "cdsl-cad-onshape-step-downloader/1.0"


@dataclass(frozen=True)
class PartStudioRef:
    sample_id: str
    source_url: str
    stack: str
    did: str
    wvm: str
    wvmid: str
    eid: str


@dataclass
class StepResult:
    status: str
    path: str | None = None
    bytes: int | None = None
    sha256: str | None = None
    translation_id: str | None = None
    http_status: int | None = None
    message: str | None = None


def parse_part_studio_url(sample_id: str, url: str) -> PartStudioRef:
    match = ONSHAPE_URL_RE.match(url.strip())
    if match is None:
        raise ValueError(f"not an Onshape Part Studio URL: {url}")
    return PartStudioRef(sample_id=sample_id, source_url=url.strip(), **match.groupdict())


def read_url_file(path: Path) -> list[tuple[str, str]]:
    """Read ABC objects YAML lines or plain '<id> <url>' records."""
    records: list[tuple[str, str]] = []
    url_pattern = re.compile(r"https://cad\.onshape\.com/documents/[^'\"\s]+")
    id_pattern = re.compile(r"^\s*['\"]?([A-Za-z0-9_-]+)['\"]?\s*[:\s]")
    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        url_match = url_pattern.search(line)
        if url_match is None:
            continue
        id_match = id_pattern.search(line)
        if id_match is None:
            raise ValueError(f"{path}:{line_number}: URL has no leading sample ID")
        records.append((id_match.group(1), url_match.group(0)))
    if not records:
        raise ValueError(f"no Onshape Part Studio URLs found in {path}")
    return records


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 ca_bundle() -> str | None:
    try:
        import certifi
    except ImportError:
        return None
    return certifi.where()


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=ca_bundle())

    def url(self, path: str) -> str:
        return f"https://{self.ref.stack}/api/{API_VERSION}{path}"

    def request(
        self,
        method: str,
        path: str,
        *,
        body: dict[str, Any] | None = None,
        accept: str = JSON_ACCEPT,
    ) -> bytes:
        payload = None if body is None else json.dumps(body).encode("utf-8")
        request = urllib.request.Request(
            self.url(path),
            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()

    def get_json(self, path: str) -> Any:
        return json.loads(self.request("GET", path).decode("utf-8"))


def sha256(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def export_step(
    client: OnshapeClient,
    root: Path,
    *,
    poll_seconds: float,
    poll_limit: int,
) -> StepResult:
    ref = client.ref
    export_path = f"/partstudios/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}/export/step"
    body = {
        "destinationName": f"{ref.sample_id}.step",
        "grouping": True,
        "notifyUser": False,
        "storeInDocument": False,
        "triggerAutoDownload": False,
        "stepUnit": "METER",
        "stepVersionString": "AP242",
    }
    try:
        translation = json.loads(client.request("POST", export_path, body=body).decode("utf-8"))
        translation_id = translation.get("id")
        if not isinstance(translation_id, str) or not translation_id:
            return StepResult("error", message="STEP export response has no translation id")
        translation_path = f"/translations/{urllib.parse.quote(translation_id, safe='')}"
        for _ in range(poll_limit):
            state = client.get_json(translation_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 or not isinstance(external_ids[0], str):
                    return StepResult("error", translation_id=translation_id, message="completed STEP export has no external data id")
                payload = client.request(
                    "GET",
                    f"/documents/d/{ref.did}/externaldata/{urllib.parse.quote(external_ids[0], safe='')}",
                    accept=STEP_ACCEPT,
                )
                if not payload.startswith(b"ISO-10303-21"):
                    return StepResult("invalid_step", translation_id=translation_id, message="download is empty or lacks the ISO-10303-21 STEP header")
                destination = root / "model.step"
                destination.write_bytes(payload)
                return StepResult(
                    "downloaded",
                    path=destination.name,
                    bytes=len(payload),
                    sha256=sha256(payload),
                    translation_id=translation_id,
                )
            if request_state == "FAILED":
                return StepResult("translation_failed", translation_id=translation_id, message=str(state.get("failureReason") or "translation failed"))
            time.sleep(poll_seconds)
        return StepResult("translation_timeout", translation_id=translation_id, message=f"STEP translation did not finish after {poll_limit} polls")
    except urllib.error.HTTPError as exc:
        return StepResult("http_error", http_status=exc.code, message=exc.read().decode("utf-8", "replace")[:1000])
    except (urllib.error.URLError, TimeoutError, OSError, ValueError, json.JSONDecodeError) as exc:
        return StepResult("error", message=str(exc))


def download_one(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)
    result = export_step(OnshapeClient(ref, authorization, timeout), root, poll_seconds=poll_seconds, poll_limit=poll_limit)
    manifest = {
        "schema": "onshape.step_sample.v1",
        "source": asdict(ref),
        "artifact": {"name": "step", **asdict(result)},
    }
    (root / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=True, indent=2) + "\n", encoding="utf-8")
    return {"sample_id": ref.sample_id, **asdict(result)}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", action="append", default=[], metavar="ID=URL", help="Onshape Part Studio URL; repeatable")
    parser.add_argument("--url-file", type=Path, help="ABC objects YAML or text file containing ID and URL records")
    parser.add_argument(
        "--count",
        type=int,
        help="maximum samples to download (defaults to the two built-in samples; custom URLs are appended)",
    )
    parser.add_argument("--out", type=Path, default=DEFAULT_OUTPUT, help="output directory")
    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.count is not None and args.count < 1) or args.poll_seconds <= 0 or args.poll_limit < 1:
        raise ValueError("count and poll limit must be positive; poll interval must be greater than zero")
    records: list[tuple[str, str]] = list(DEFAULT_SAMPLES)
    for value in args.url:
        if "=" not in value:
            raise ValueError("--url must use ID=URL syntax")
        records.append(tuple(value.split("=", 1)))
    if args.url_file:
        records.extend(read_url_file(args.url_file))
    limit = args.count if args.count is not None else len(records)
    if limit < 1:
        raise ValueError("count must be positive")
    refs = [parse_part_studio_url(sample_id, url) for sample_id, url in records[:limit]]
    access_key, secret_key = credentials()
    authorization = "Basic " + base64.b64encode(f"{access_key}:{secret_key}".encode("utf-8")).decode("ascii")
    results = [download_one(ref, authorization, args.out, timeout=args.timeout, poll_seconds=args.poll_seconds, poll_limit=args.poll_limit) for ref in refs]
    summary = {
        "schema": "onshape.step_download_summary.v1",
        "requested_count": len(refs),
        "downloaded_count": sum(item["status"] == "downloaded" for item in results),
        "failed_count": sum(item["status"] != "downloaded" for item in results),
        "results": results,
    }
    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "manifest.json").write_text(json.dumps(summary, ensure_ascii=True, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=True, indent=2))
    return 0 if summary["failed_count"] == 0 else 1


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