#!/usr/bin/env bash
# beagle-ast: project a strictly checked program as versioned JSON.
#
# Usage:
#   beagle ast [--module-root LOGICAL=PHYSICAL]...
#     [--interface-sha256-out PATH] <source.bgl|.bclj|.bjs|.bnix>
#   beagle ast --bundle [--module-root LOGICAL=PHYSICAL]... -- <source>...
#
# Parses and checks without executing or emitting the source program. Output is
# canonical JSON for deterministic external consumers.

set -euo pipefail
source "$(dirname "$0")/_beagle-racket"
source "$(dirname "$0")/_beagle-source-id"

bundle=0
interface_sha256_out=""
module_root_specs=()
sources=()

usage() {
    echo "usage: beagle ast [--module-root LOGICAL=PHYSICAL]... [--interface-sha256-out PATH] <source> | --bundle [--module-root LOGICAL=PHYSICAL]... -- <source>..." >&2
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --bundle)
            bundle=1
            shift
            ;;
        --interface-sha256-out)
            [[ $# -ge 2 ]] || {
                usage
                exit 2
            }
            interface_sha256_out="$2"
            shift 2
            ;;
        --module-root)
            [[ $# -ge 2 ]] || {
                usage
                exit 2
            }
            module_root_specs+=("$2")
            shift 2
            ;;
        --module-root=*)
            module_root_specs+=("${1#*=}")
            shift
            ;;
        --)
            shift
            while [[ $# -gt 0 ]]; do
                sources+=("$1")
                shift
            done
            ;;
        -*)
            echo "beagle ast: unknown option: $1" >&2
            usage
            exit 2
            ;;
        *)
            sources+=("$1")
            shift
            ;;
    esac
done

if [[ ${#sources[@]} -lt 1 || ( $bundle -eq 0 && ${#sources[@]} -ne 1 ) ]]; then
    usage
    exit 2
fi
[[ $bundle -eq 0 || -z "$interface_sha256_out" ]] || {
    echo "beagle ast: --interface-sha256-out is only valid for one source" >&2
    exit 2
}

for src in "${sources[@]}"; do
    if [[ ! -f "$src" ]]; then
        echo "beagle ast: source file not found: $src" >&2
        exit 1
    fi
done

# Require by file path (not collection) so the WORKTREE's beagle-lib — which
# carries ast-json.rkt — is what loads, instead of the canonical checkout's
# collection link (which lacks it).
BEAGLE_ROOT="$(cd "$(dirname "$0")/.." && pwd)"

reuse=0
[[ "$bundle" == 1 &&
   "${BEAGLE_CHECKED_AST_REUSE:-1}" != 0 &&
   "${BEAGLE_FACT_REUSE_FORBIDDEN:-0}" != 1 ]] && reuse=1
ast_profile="beagle-ast-v2"
ast_unit="$(if [[ "$bundle" == 1 ]]; then echo bundle; else echo single; fi)"
tmp_dir="$(mktemp -d "${TMPDIR:-/tmp}/beagle-ast.XXXXXX")"
trap 'rm -rf "${tmp_dir:?}"' EXIT

run_checked_ast_store() {
    (
        flock --exclusive 9
        bb -cp "$BEAGLE_ROOT/store/out" \
            "$BEAGLE_ROOT/bin/_beagle-checked-ast-store.clj" "$@"
    ) 9>"$ast_store.publish.lock"
}

closure_fields="$tmp_dir/closure-fields"
if [[ "$reuse" == 1 ]]; then
    if [[ -n "${BEAGLE_AST_CLOSURE_FIELDS:-}" ]]; then
        cp -- "${BEAGLE_AST_CLOSURE_FIELDS}" "$closure_fields"
    else
        closure_cli="$BEAGLE_ROOT/beagle-lib/private/module-source-root-cli.rkt"
        closure_args=()
        for spec in "${module_root_specs[@]}"; do closure_args+=(--module-root "$spec"); done
        for src in "${sources[@]}"; do
            src_abs="$(realpath "$src")"
            closure_args+=(--source "$src_abs" "$(beagle_source_id "$src_abs" "${module_root_specs[@]}")")
        done
        "$RACKET" "$closure_cli" "${closure_args[@]}" >"$closure_fields"
    fi
    source_hashes="$tmp_dir/source-hashes"
    : >"$source_hashes"
    while IFS= read -r -d '' physical && IFS= read -r -d '' logical; do
        printf '%s\0%s\0' "$logical" "$(sha256sum -- "$physical" | awk '{print $1}')" >>"$source_hashes"
    done <"$closure_fields"
    # A checked bundle is controlled by the parser/checker closure, not by
    # unrelated build scripts or documentation.  Keep this list at the
    # parse/check normalization seam: it names the Racket modules whose source
    # can change the checked AST or its public interface.
    compiler_context="sha256:$(sha256sum "$0" \
        "$BEAGLE_ROOT/beagle-lib/private/ast-json.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/ast.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/parse.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/check.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/module-interface.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/module-overlay-check.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/module-source-root.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/effect-normalization-v1.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/scope-resolve.rkt" \
        "$BEAGLE_ROOT/beagle-lib/private/types.rkt" | awk '{print $1}' | sha256sum | awk '{print $1}')"
    key_input="$tmp_dir/key-input"
    printf 'checked-ast-bundle-v2\0%s\0' "$compiler_context" >"$key_input"
    for spec in "${module_root_specs[@]}"; do
        printf 'module-root\0%s\0' "${spec%%=*}" >>"$key_input"
    done
    cat "$source_hashes" >>"$key_input"
    result_key="sha256:$(sha256sum "$key_input" | awk '{print $1}')"
    checked_ast_cache_root="${BEAGLE_CHECKED_AST_CACHE:-${XDG_CACHE_HOME:-${HOME:?HOME is required}/.cache}/beagle/checked-ast}"
    ast_store="${BEAGLE_CHECKED_AST_STORE:-$checked_ast_cache_root/${result_key#sha256:}.storelog}"
    mkdir -p "$(dirname "$ast_store")"
    if run_checked_ast_store query "$ast_store" "$result_key" "$compiler_context" "$ast_profile" "$ast_unit" >"$tmp_dir/hit" 2>/dev/null; then
        cat "$tmp_dir/hit"
        exit 0
    fi
fi

source_args=()
for src in "${sources[@]}"; do
    src_abs="$(realpath "$src")"
    source_args+=(
        "$src_abs"
        "$(beagle_source_id "$src_abs" "${module_root_specs[@]}")"
    )
done

projection="$tmp_dir/projection"
"$RACKET" -e "
(require json
         racket/list
         (file \"$BEAGLE_ROOT/beagle-lib/private/ast-json.rkt\")
         (file \"$BEAGLE_ROOT/beagle-lib/private/module-interface.rkt\")
         (file \"$BEAGLE_ROOT/beagle-lib/private/module-overlay-check.rkt\")
         (file \"$BEAGLE_ROOT/beagle-lib/private/module-source-root.rkt\")
         (file \"$BEAGLE_ROOT/beagle-lib/private/typescript-foreign-resolver-v1.rkt\"))

(define arguments (vector->list (current-command-line-arguments)))
(unless (>= (length arguments) 3)
  (error 'beagle-ast \"missing internal invocation arguments\"))
(define bundle? (string=? (car arguments) \"1\"))
(define interface-sha256-out (cadr arguments))
(define root-count (string->number (caddr arguments)))
(unless (and (exact-nonnegative-integer? root-count)
             (<= root-count (- (length arguments) 3)))
  (error 'beagle-ast \"invalid internal module-root count\"))
(define payload (cdddr arguments))
(define roots
  (map parse-module-source-root (take payload root-count)))
(define input-arguments (drop payload root-count))
(unless (and (positive? (length input-arguments))
             (even? (length input-arguments)))
  (error 'beagle-ast
         \"inputs must pair a physical path with a logical source id\"))
(define inputs
  (for/list ([index (in-range 0 (length input-arguments) 2)])
    (module-source-input
     (list-ref input-arguments (add1 index))
     (list-ref input-arguments index))))
(define closure (resolve-production-module-source-closure inputs roots))
(define checked
  (check-module-source-closure
   closure
   #:emit? #f
   #:capture-types? #t
   #:check-sources
   (and (not bundle?)
        (module-source-closure-explicit-source-ids closure))))
(unless (overlay-check-result-ok? checked)
  (for ([diagnostic (in-list (overlay-check-result-diagnostics checked))])
    (eprintf \"beagle ast: ~a: ~a~n\"
             (or (overlay-diagnostic-source diagnostic) \"module closure\")
             (overlay-diagnostic-message diagnostic)))
  (exit 1))
(define checked-modules
  (sort
   (overlay-check-result-modules checked)
   string<?
   #:key (lambda (module)
           (format \"~a\" (checked-overlay-module-source module)))))
(unless bundle?
  (unless (= (length checked-modules) 1)
    (error 'beagle-ast
           \"single-source projection selected ~a checked modules\"
           (length checked-modules)))
  (define module (car checked-modules))
  (unless (string=? interface-sha256-out \"\")
    (call-with-output-file
     interface-sha256-out
     #:exists 'truncate/replace
     (lambda (interface-out)
       (display
        (module-interface-digest (checked-overlay-module-interface module))
        interface-out)
       (newline interface-out))))
  (write-checked-program-json
   (checked-overlay-module-program module)
   (current-output-port)
   #:source-id (format \"~a\" (checked-overlay-module-source module)))
  (exit 0))
(define modules
  (for/list ([module (in-list checked-modules)])
    (define output (open-output-string))
    (write-checked-program-json
     (checked-overlay-module-program module)
     output
     #:source-id (format \"~a\" (checked-overlay-module-source module)))
    (hasheq
     'source (format \"~a\" (checked-overlay-module-source module))
     'interfaceSha256
     (module-interface-digest (checked-overlay-module-interface module))
     'program (string->jsexpr (get-output-string output)))))
(write-json (hasheq 'schemaVersion 2 'modules modules))
(newline)
" -- "$bundle" "$interface_sha256_out" "${#module_root_specs[@]}" \
    "${module_root_specs[@]}" "${source_args[@]}" >"$projection"

if [[ "$reuse" == 1 ]]; then
    source_hashes_after="$tmp_dir/source-hashes-after"
    : >"$source_hashes_after"
    while IFS= read -r -d '' physical && IFS= read -r -d '' logical; do
        [[ -f "$physical" ]] || break
        printf '%s\0%s\0' "$logical" "$(sha256sum -- "$physical" | awk '{print $1}')" >>"$source_hashes_after"
    done <"$closure_fields"
    if cmp -s "$source_hashes" "$source_hashes_after"; then
        run_checked_ast_store append "$ast_store" "$result_key" "$compiler_context" "$ast_profile" "$ast_unit" "$projection"
    fi
fi
cat "$projection"
