#!/usr/bin/env python3
"""asc_api.py — minimal App Store Connect API helper for ad-hoc device registration.

Used by scripts/dist-ios-adhoc.sh to register friends' device UDIDs before building an
ad-hoc IPA. Talks to the App Store Connect API directly: signs an ES256 JWT with your
Team Key .p8 (via the `cryptography` package — no PyJWT/requests needed) and calls the
REST endpoints with urllib from the standard library.

Subcommands:
  register   register one device UDID (idempotent — an already-registered UDID is OK)
  list       list all registered devices (and the count, against the 100/year cap)

Auth is the same for both, supplied via flags (the wrapper script passes them from the
ASC_KEY_ID / ASC_ISSUER_ID / ASC_KEY_PATH env vars):
  --key-id      the Key ID of the App Store Connect API key
  --issuer-id   the Issuer ID (Users and Access -> Integrations)
  --key         path to the AuthKey_XXXXXXXXXX.p8 file

Examples:
  python3 scripts/asc_api.py list \
      --key-id ABC123 --issuer-id 11111111-2222-... --key ~/.appstoreconnect/AuthKey_ABC123.p8
  python3 scripts/asc_api.py register --udid 00008110-0011... --name "My iPhone" \
      --key-id ABC123 --issuer-id 11111111-2222-... --key ~/.appstoreconnect/AuthKey_ABC123.p8

The .p8 is created at App Store Connect -> Users and Access -> Integrations ->
App Store Connect API -> Team Keys, with Admin or App Manager access. Keep it out of the
repo.
"""

import argparse
import base64
import json
import sys
import time
import urllib.error
import urllib.request

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, utils

API_BASE = "https://api.appstoreconnect.apple.com"


def _b64url(data: bytes) -> str:
    """base64url without padding, as JWT requires."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def make_jwt(key_id: str, issuer_id: str, key_path: str) -> str:
    """Build a short-lived ES256 JWT for the App Store Connect API."""
    with open(key_path, "rb") as fh:
        private_key = serialization.load_pem_private_key(fh.read(), password=None)
    if not isinstance(private_key, ec.EllipticCurvePrivateKey):
        sys.exit(f"error: {key_path} is not an EC private key (.p8 from App Store Connect)")

    now = int(time.time())
    header = {"alg": "ES256", "kid": key_id, "typ": "JWT"}
    # exp must be <= 20 minutes out; 10 minutes is comfortable.
    payload = {"iss": issuer_id, "iat": now, "exp": now + 600, "aud": "appstoreconnect-v1"}

    signing_input = f"{_b64url(json.dumps(header).encode())}.{_b64url(json.dumps(payload).encode())}"
    der_sig = private_key.sign(signing_input.encode("ascii"), ec.ECDSA(hashes.SHA256()))
    # JWS wants raw r||s (two 32-byte big-endian ints), not the ASN.1/DER openssl emits.
    r, s = utils.decode_dss_signature(der_sig)
    raw_sig = r.to_bytes(32, "big") + s.to_bytes(32, "big")
    return f"{signing_input}.{_b64url(raw_sig)}"


def _request(method: str, path: str, token: str, body: dict | None = None):
    """Perform an authenticated API call. Returns (status_code, parsed_json|None)."""
    url = path if path.startswith("http") else f"{API_BASE}{path}"
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Authorization", f"Bearer {token}")
    if data is not None:
        req.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req) as resp:
            raw = resp.read()
            return resp.status, (json.loads(raw) if raw else None)
    except urllib.error.HTTPError as exc:
        raw = exc.read()
        try:
            parsed = json.loads(raw) if raw else None
        except json.JSONDecodeError:
            parsed = {"_raw": raw.decode("utf-8", "replace")}
        return exc.code, parsed


def _errors_text(payload) -> str:
    if isinstance(payload, dict) and payload.get("errors"):
        return "; ".join(
            f"{e.get('title', '')}: {e.get('detail', '')}".strip(": ")
            for e in payload["errors"]
        )
    return json.dumps(payload)


def cmd_register(args, token: str) -> int:
    body = {
        "data": {
            "type": "devices",
            "attributes": {
                "name": args.name or args.udid,
                "platform": "IOS",
                "udid": args.udid,
            },
        }
    }
    status, payload = _request("POST", "/v1/devices", token, body)
    if status in (200, 201):
        print(f"  registered: {args.udid}  ({args.name or args.udid})")
        return 0

    # A UDID that already exists comes back as a 409 conflict, or a 422 with an error
    # detail mentioning the device already exists. Either way it's fine — idempotent.
    text = _errors_text(payload)
    if status == 409 or "already exist" in text.lower() or "already been taken" in text.lower():
        print(f"  already registered: {args.udid}")
        return 0

    print(f"error: failed to register {args.udid} (HTTP {status}): {text}", file=sys.stderr)
    return 1


def cmd_list(args, token: str) -> int:
    path = "/v1/devices?limit=200&sort=name"
    rows = []
    while path:
        status, payload = _request("GET", path, token)
        if status != 200:
            print(f"error: list failed (HTTP {status}): {_errors_text(payload)}", file=sys.stderr)
            return 1
        for d in payload.get("data", []):
            a = d.get("attributes", {})
            rows.append((a.get("platform", "?"), a.get("status", "?"),
                         a.get("udid", "?"), a.get("name", "")))
        path = (payload.get("links") or {}).get("next")

    ios = [r for r in rows if r[0] == "IOS"]
    for platform, dev_status, udid, name in rows:
        print(f"  [{platform:7}] {dev_status:8} {udid}  {name}")
    print(f"\n  {len(rows)} device(s) total, {len(ios)} iOS (cap is 100 iOS/membership year)")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)

    def add_auth(p):
        p.add_argument("--key-id", required=True)
        p.add_argument("--issuer-id", required=True)
        p.add_argument("--key", required=True, help="path to AuthKey_*.p8")

    p_reg = sub.add_parser("register", help="register one device UDID (idempotent)")
    add_auth(p_reg)
    p_reg.add_argument("--udid", required=True)
    p_reg.add_argument("--name", default=None)

    p_list = sub.add_parser("list", help="list registered devices")
    add_auth(p_list)

    args = parser.parse_args()
    token = make_jwt(args.key_id, args.issuer_id, args.key)
    if args.command == "register":
        return cmd_register(args, token)
    if args.command == "list":
        return cmd_list(args, token)
    return 2


if __name__ == "__main__":
    sys.exit(main())
