#!/usr/bin/env bash
#
# figma-screenshot.sh
#
# Download Figma frame PNGs for /multi-agent:analysis Phase 2b.
#
# Modes:
#   Single frame:
#     figma-screenshot.sh --file-key <fileKey> --node-id <nodeId> \
#                        [--output-dir <dir>] [--scale 2] [--format png]
#
#   Section drill:
#     figma-screenshot.sh --section <sectionUrl> \
#                        [--output-dir <dir>] [--scale 2]
#
# Section drill: parse the section URL, fetch metadata, walk
# document.children, keep nodes whose type is FRAME, render them in
# a single batched image call, then download each signed URL in
# parallel. Manifest is emitted alongside the PNGs.
#
# Tier chain (per ~/.claude/rules/figma-pipeline.md):
#   Tier 1 (Figma MCP)            -> N/A inside a bash script; noted on stderr
#   Tier 2 (Figma REST + PAT)     -> primary path here
#   Tier 3 (user screenshot)      -> not handled; exit 3 with guidance
#
# Token source: figma-token.sh, which resolves the canonical `figma` logical
# key through credential-store.sh (falling back to the legacy `figma_pat` key,
# then to the FIGMA_PAT environment variable). No literal keychain service name
# is embedded in this script.
#
# Exit codes:
#   0  success
#   1  usage error
#   2  auth failure (missing or rejected PAT)
#   3  Tier 3 required (REST exhausted)
#   4  frame not found (HTTP 404)
#   5  rate limited (HTTP 429 after retry)
#   6  network timeout or other fatal failure

set -euo pipefail

# --- Globals ----------------------------------------------------------------

# Locate the resolver with an existence check, not a `.`-chain: sourcing a missing file
# aborts the shell under `set -e`, `||` included, so a chain skips both its later
# candidates and its trailing `|| true`. A missing store is tolerated here - CRED_STORE
# simply stays empty and the caller falls back to an env-supplied token.
for _cred_resolver in \
  "$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)/credential-store-resolver.sh" \
  "$HOME/.claude/lib/credential-store-resolver.sh" \
  "$HOME/.copilot/lib/credential-store-resolver.sh" \
  "$HOME/.codex/lib/credential-store-resolver.sh"; do
  [ -f "$_cred_resolver" ] || continue
  # shellcheck source=/dev/null
  . "$_cred_resolver" 2>/dev/null || true
  # `if`, not `[ ... ] && break`: the latter is the loop body's last command and returns
  # 1 when CRED_STORE is still empty, which under `set -e` kills the loop on the first
  # candidate that does not resolve - the very case the loop exists to survive.
  if [ -n "${CRED_STORE:-}" ]; then break; fi
done
unset _cred_resolver
CRED_STORE="${CRED_STORE:-}"

# Tier 2 token lookup lives in one file so it cannot drift per fetcher. Same
# existence-check discipline as the resolver loop above, and the same tolerance:
# without it, resolution still works through the FIGMA_PAT env fallback.
for _figma_token_lib in \
  "$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)/figma-token.sh" \
  "$HOME/.claude/lib/figma-token.sh" \
  "$HOME/.copilot/lib/figma-token.sh" \
  "$HOME/.codex/lib/figma-token.sh"; do
  [ -f "$_figma_token_lib" ] || continue
  # shellcheck source=/dev/null
  . "$_figma_token_lib" 2>/dev/null || true
  if command -v resolve_figma_token >/dev/null 2>&1; then break; fi
done
unset _figma_token_lib
if ! command -v resolve_figma_token >/dev/null 2>&1; then
  resolve_figma_token() { [ -n "${FIGMA_PAT:-}" ] && printf '%s' "$FIGMA_PAT"; }
  figma_token_remediation() { printf 'figma-token.sh not found - reinstall the pipeline, or export FIGMA_PAT.'; }
fi

FIGMA_API="https://api.figma.com/v1"
PARALLEL_DL=4
HTTP_TIMEOUT=60
MAX_IDS_PER_CALL=50

FILE_KEY=""
NODE_ID=""
SECTION_URL=""
OUTPUT_DIR=""
SCALE="2"
FORMAT="png"
TIER="rest"

# --- Helpers ----------------------------------------------------------------

log() { printf '%s\n' "$*" >&2; }

die() {
  local code="$1"; shift
  log "ERR: $*"
  exit "$code"
}

print_help() {
  cat <<'HELP'
figma-screenshot.sh - download Figma frames as PNG

Usage:
  figma-screenshot.sh --file-key <fileKey> --node-id <nodeId> \
                     [--output-dir <dir>] [--scale 2] [--format png]

  figma-screenshot.sh --section <sectionUrl> \
                     [--output-dir <dir>] [--scale 2]

Options:
  --file-key       Figma file key (or branchKey for branch URLs).
  --node-id        Node id; may use either "-" or ":" separator.
  --section        Section URL; child FRAMEs are auto-discovered.
  --output-dir     Output directory. Default: /tmp/figma-screenshots-<ts>.
  --scale          Render scale (1-4). Default: 2.
  --format         Image format. Default: png.
  --help           Show this help text and exit.

Token resolution:
  Read via credential-store.sh using prefs.global.keychainMapping.figma
  (legacy figma_pat still honoured). Falls back to FIGMA_PAT env.

Output:
  PNG files plus manifest.json in the output directory.
HELP
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || die 6 "missing dependency: $1"
}

normalize_node_id() {
  # Figma node ids are sometimes carried with dashes (URL form) and sometimes
  # with colons (API form). The API requires the colon form.
  printf '%s' "$1" | tr '-' ':'
}

node_id_filename() {
  printf '%s' "$1" | tr ':' '-'
}

# Parse a Figma URL into "fileKey<TAB>nodeId". Handles:
#   figma.com/design/<fileKey>/<file>?node-id=<id>
#   figma.com/design/<fileKey>/branch/<branchKey>/<file>?node-id=<id>
parse_figma_url() {
  local url="$1"
  python3 - "$url" <<'PY'
import sys
from urllib.parse import urlparse, parse_qs

url = sys.argv[1]
u = urlparse(url)
parts = [p for p in u.path.split('/') if p]
file_key = ''
if len(parts) >= 2 and parts[0] in ('design', 'file', 'proto', 'board'):
    file_key = parts[1]
# branch URLs override the file key with the branch key
if len(parts) >= 4 and parts[2] == 'branch':
    file_key = parts[3]
q = parse_qs(u.query)
node = q.get('node-id', [''])[0]
print(f"{file_key}\t{node}")
PY
}

# curl_figma <output> <url>
# Honours retries: 429 retries once after a 5s pause; 5xx uses exponential
# backoff (1s, 2s, 4s). Writes the body to <output>, returns the HTTP code on
# stdout.
curl_figma() {
  local out="$1" url="$2"
  local code attempt sleep_s
  for attempt in 1 2 3; do
    # Token header goes through a curl config via process substitution so it
    # never appears in argv (argv is visible to ps).
    # curl already writes "000" via -w on a total connection failure (and
    # still exits non-zero for it, which the inner `|| echo "000"` used to
    # catch too - so `code` became the literal string "000000", matching
    # neither the `000)` retry branch below nor handle_http_code's "000)
    # network timeout" case, just the generic catch-all). `|| true` on its
    # own line keeps `set -e` from aborting on curl's exit status without
    # re-printing anything into $code; the `:-000` fallback only fires for
    # the rarer case where curl fails before -w ever runs and $code is
    # genuinely empty.
    code=$(curl -sS -o "$out" -w '%{http_code}' \
              --max-time "$HTTP_TIMEOUT" --connect-timeout 5 \
              -K <(printf 'header = "X-Figma-Token: %s"\n' "$FIGMA_TOKEN") \
              "$url") || true
    code="${code:-000}"
    case "$code" in
      2*) printf '%s' "$code"; return 0 ;;
      401|403) printf '%s' "$code"; return 0 ;;
      404)     printf '%s' "$code"; return 0 ;;
      429)
        if [ "$attempt" -eq 1 ]; then
          log "WARN: 429 from Figma, retrying after 5s"
          sleep 5
          continue
        fi
        printf '%s' "$code"; return 0 ;;
      5*)
        sleep_s=$((2 ** (attempt - 1)))
        log "WARN: HTTP $code from Figma, retry in ${sleep_s}s"
        sleep "$sleep_s"
        ;;
      000)
        log "WARN: network failure on Figma call (attempt $attempt)"
        sleep_s=$((2 ** (attempt - 1)))
        sleep "$sleep_s"
        ;;
      *) printf '%s' "$code"; return 0 ;;
    esac
  done
  printf '%s' "${code:-000}"
}

handle_http_code() {
  local code="$1" context="$2"
  case "$code" in
    2*) return 0 ;;
    401|403) die 2 "auth failed ($context), check FIGMA_PAT" ;;
    404)     die 4 "frame not found ($context)" ;;
    429)     die 5 "rate limited ($context), retry exhausted" ;;
    000)     die 6 "network timeout ($context)" ;;
    *)       die 6 "HTTP $code ($context)" ;;
  esac
}

# Fetch node metadata for a comma-joined list of node ids. Result written to
# the given path.
fetch_nodes_metadata() {
  local ids="$1" out="$2"
  local url="$FIGMA_API/files/$FILE_KEY/nodes?ids=$ids"
  local code
  code=$(curl_figma "$out" "$url")
  handle_http_code "$code" "metadata for $ids"
}

# Render images for a comma-joined list of node ids. Returns a JSON object on
# stdout: { "nodeId": "signed-url", ... }.
fetch_image_urls() {
  local ids="$1"
  local tmp
  tmp=$(mktemp)
  local url="$FIGMA_API/images/$FILE_KEY?ids=$ids&format=$FORMAT&scale=$SCALE"
  local code
  code=$(curl_figma "$tmp" "$url")
  handle_http_code "$code" "image render for $ids"
  # Surface Figma-level errors (the API can return 200 with { "err": "..." }).
  local err
  err=$(jq -r '.err // empty' "$tmp" 2>/dev/null || true)
  if [ -n "$err" ]; then
    rm -f "$tmp"
    die 6 "Figma image API error: $err"
  fi
  jq -c '.images' "$tmp"
  rm -f "$tmp"
}

# Render + download a batch of node ids in one shot. Caller passes the comma-
# joined id list and a TSV "nodeId<TAB>name<TAB>width<TAB>height" lookup file.
render_and_download() {
  local ids="$1" lookup="$2"
  local urls_json
  urls_json=$(fetch_image_urls "$ids")

  # Emit a job list (nodeId<TAB>url<TAB>outPath) for parallel download.
  local jobs
  jobs=$(mktemp)
  local key_prefix="${FILE_KEY:0:4}"
  printf '%s' "$urls_json" | jq -r 'to_entries[] | "\(.key)\t\(.value)"' \
    | while IFS=$'\t' read -r nid url; do
        local fname
        fname="frame-${key_prefix}-$(node_id_filename "$nid").$FORMAT"
        printf '%s\t%s\t%s/%s\n' "$nid" "$url" "$OUTPUT_DIR" "$fname"
      done > "$jobs"

  # Parallel download. xargs -P keeps the channel count bounded.
  local results
  results=$(mktemp)
  # shellcheck disable=SC2016
  awk -F'\t' '{printf "%s\037%s\037%s\n", $1,$2,$3}' "$jobs" \
    | xargs -P "$PARALLEL_DL" -I{} bash -c '
        IFS=$'"'"'\037'"'"' read -r nid url out <<<"$1"
        if [ -z "$url" ] || [ "$url" = "null" ]; then
          printf "%s\tFAIL\t0\n" "$nid"
          exit 0
        fi
        if curl -sS --max-time "$2" --connect-timeout 5 -o "$out" "$url"; then
          size=$(wc -c < "$out" | tr -d " ")
          printf "%s\t%s\t%s\n" "$nid" "$out" "$size"
        else
          printf "%s\tFAIL\t0\n" "$nid"
        fi
      ' _ {} "$HTTP_TIMEOUT" >> "$results"

  # Merge download results with the metadata lookup to build manifest rows.
  python3 - "$results" "$lookup" <<'PY'
import json, sys

results_path, lookup_path = sys.argv[1], sys.argv[2]

lookup = {}
with open(lookup_path) as f:
    for line in f:
        line = line.rstrip("\n")
        if not line:
            continue
        parts = line.split("\t")
        if len(parts) < 4:
            continue
        nid, name, w, h = parts[0], parts[1], parts[2], parts[3]
        lookup[nid] = {"name": name, "width": w, "height": h}

with open(results_path) as f:
    for line in f:
        line = line.rstrip("\n")
        if not line:
            continue
        parts = line.split("\t")
        if len(parts) < 3:
            continue
        nid, fpath, size = parts[0], parts[1], parts[2]
        meta = lookup.get(nid, {"name": "", "width": "", "height": ""})
        print(json.dumps({
            "nodeId": nid,
            "name": meta["name"],
            "width": meta["width"],
            "height": meta["height"],
            "file": fpath.split("/")[-1] if fpath != "FAIL" else "",
            "fileSize": int(size) if size.isdigit() else 0,
            "ok": fpath != "FAIL"
        }))
PY

  rm -f "$jobs" "$results"
}

write_manifest() {
  local manifest_path="$1" section_id="$2" rows_file="$3"
  python3 - "$manifest_path" "$FILE_KEY" "$section_id" "$SCALE" "$TIER" "$rows_file" <<'PY'
import json, sys
from datetime import datetime, timezone

manifest_path, file_key, section_id, scale, tier, rows_file = sys.argv[1:]
rows = []
with open(rows_file) as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        try:
            rows.append(json.loads(line))
        except json.JSONDecodeError:
            continue

doc = {
    "fileKey": file_key,
    "section": section_id or None,
    "scale": int(scale),
    "tier": tier,
    "generated": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    "frames": rows,
}
with open(manifest_path, "w") as f:
    json.dump(doc, f, indent=2)
PY
}

# --- Argument parsing -------------------------------------------------------

if [ "$#" -eq 0 ]; then
  print_help
  exit 1
fi

while [ "$#" -gt 0 ]; do
  case "$1" in
    --help|-h) print_help; exit 0 ;;
    --file-key) FILE_KEY="${2:-}"; shift 2 ;;
    --node-id)  NODE_ID="${2:-}"; shift 2 ;;
    --section)  SECTION_URL="${2:-}"; shift 2 ;;
    --output-dir) OUTPUT_DIR="${2:-}"; shift 2 ;;
    --scale)    SCALE="${2:-2}"; shift 2 ;;
    --format)   FORMAT="${2:-png}"; shift 2 ;;
    *) die 1 "unknown argument: $1" ;;
  esac
done

if [ -n "$SECTION_URL" ] && { [ -n "$FILE_KEY" ] || [ -n "$NODE_ID" ]; }; then
  die 1 "--section is mutually exclusive with --file-key/--node-id"
fi

if [ -z "$SECTION_URL" ] && { [ -z "$FILE_KEY" ] || [ -z "$NODE_ID" ]; }; then
  die 1 "either --section <url> or both --file-key and --node-id are required"
fi

require_cmd curl
require_cmd jq
require_cmd python3

# Tier 1 (Figma MCP) is not reachable from a bash script. Stay on Tier 2 (REST).
log "INFO: Tier 1 (Figma MCP) not available from shell; using Tier 2 (REST)."

FIGMA_TOKEN=$(resolve_figma_token || true)
if [ -z "${FIGMA_TOKEN:-}" ]; then
  log "Tier 3 (user-provided screenshot required): no PAT resolved."
  log "Hint: $(figma_token_remediation)"
  exit 3
fi

if [ -z "$OUTPUT_DIR" ]; then
  OUTPUT_DIR="/tmp/figma-screenshots-$(date +%Y%m%d-%H%M%S)"
fi
mkdir -p "$OUTPUT_DIR"

# --- Section mode -----------------------------------------------------------

run_section_mode() {
  local parsed file_key node_id
  parsed=$(parse_figma_url "$SECTION_URL")
  file_key=$(printf '%s' "$parsed" | cut -f1)
  node_id=$(printf '%s' "$parsed" | cut -f2)
  if [ -z "$file_key" ] || [ -z "$node_id" ]; then
    die 1 "could not parse --section URL: $SECTION_URL"
  fi
  FILE_KEY="$file_key"
  local section_id
  section_id=$(normalize_node_id "$node_id")

  local meta_path
  meta_path="$OUTPUT_DIR/_section-metadata.json"
  fetch_nodes_metadata "$section_id" "$meta_path"

  # Build the lookup: nodeId<TAB>name<TAB>width<TAB>height for every FRAME
  # child of the section.
  local lookup
  lookup="$OUTPUT_DIR/_lookup.tsv"
  python3 - "$meta_path" "$section_id" > "$lookup" <<'PY'
import json, sys
meta_path, section_id = sys.argv[1], sys.argv[2]
data = json.load(open(meta_path))
node = data.get("nodes", {}).get(section_id)
if not node:
    sys.exit(0)
doc = node.get("document", {})
for child in doc.get("children", []):
    if child.get("type") != "FRAME":
        continue
    box = child.get("absoluteBoundingBox") or {}
    print("\t".join([
        child.get("id", ""),
        (child.get("name") or "").replace("\t", " "),
        str(int(box.get("width", 0))),
        str(int(box.get("height", 0))),
    ]))
PY

  if [ ! -s "$lookup" ]; then
    write_manifest "$OUTPUT_DIR/manifest.json" "$section_id" "$lookup"
    log "WARN: no FRAME children found under section $section_id"
    log "OK: empty manifest at $OUTPUT_DIR/manifest.json"
    return 0
  fi

  # Collect ids and chunk to MAX_IDS_PER_CALL.
  local ids_file
  ids_file=$(mktemp)
  awk -F'\t' '{print $1}' "$lookup" > "$ids_file"

  local rows
  rows="$OUTPUT_DIR/_rows.ndjson"
  : > "$rows"

  local chunk=""
  local count=0
  while IFS= read -r nid; do
    if [ -z "$chunk" ]; then
      chunk="$nid"
    else
      chunk="$chunk,$nid"
    fi
    count=$((count + 1))
    if [ "$count" -ge "$MAX_IDS_PER_CALL" ]; then
      render_and_download "$chunk" "$lookup" >> "$rows"
      chunk=""
      count=0
    fi
  done < "$ids_file"
  if [ -n "$chunk" ]; then
    render_and_download "$chunk" "$lookup" >> "$rows"
  fi
  rm -f "$ids_file"

  write_manifest "$OUTPUT_DIR/manifest.json" "$section_id" "$rows"
  rm -f "$rows" "$lookup" "$meta_path"
  log "OK: manifest at $OUTPUT_DIR/manifest.json"
  printf '%s\n' "$OUTPUT_DIR/manifest.json"
}

# --- Single-frame mode ------------------------------------------------------

run_single_mode() {
  local node_id
  node_id=$(normalize_node_id "$NODE_ID")

  local meta_path lookup
  meta_path="$OUTPUT_DIR/_frame-metadata.json"
  lookup="$OUTPUT_DIR/_lookup.tsv"
  fetch_nodes_metadata "$node_id" "$meta_path"
  python3 - "$meta_path" "$node_id" > "$lookup" <<'PY'
import json, sys
meta_path, node_id = sys.argv[1], sys.argv[2]
data = json.load(open(meta_path))
node = data.get("nodes", {}).get(node_id)
if not node:
    sys.exit(0)
doc = node.get("document", {})
box = doc.get("absoluteBoundingBox") or {}
print("\t".join([
    doc.get("id", node_id),
    (doc.get("name") or "").replace("\t", " "),
    str(int(box.get("width", 0))),
    str(int(box.get("height", 0))),
]))
PY

  if [ ! -s "$lookup" ]; then
    die 4 "frame not found in metadata: $node_id"
  fi

  local rows
  rows="$OUTPUT_DIR/_rows.ndjson"
  render_and_download "$node_id" "$lookup" > "$rows"

  write_manifest "$OUTPUT_DIR/manifest.json" "" "$rows"
  rm -f "$rows" "$lookup" "$meta_path"
  log "OK: manifest at $OUTPUT_DIR/manifest.json"
  printf '%s\n' "$OUTPUT_DIR/manifest.json"
}

# --- Dispatch ---------------------------------------------------------------

if [ -n "$SECTION_URL" ]; then
  run_section_mode
else
  run_single_mode
fi