#!/usr/bin/env bash
#
# sprinter-cc — wrapper around SDCC for building Sprinter .EXE files.
#
# Usage:
#   sprinter-cc -o foo.exe foo.c [more.c ...] [options]
#
# Options:
#   -o NAME          output executable name (required)
#   --crt0=TYPE      crt0 variant: default | minimal | banked   (default: default)
#   --memory MODE    memory layout: tiny | small | big | huge | manual
#                    tiny  (default): CODE+DATA+stack in W2 (single page, ≤14 KB)
#                    small: CODE in W1, DATA+BSS continue into W2, heap spans
#                           remainder, stack in W2.  crt0_small reads port
#                           0xA2 to auto-detect whether DSS already mapped W2
#                           (program > 16 KB) or whether we need to allocate
#                           it ourselves (program < 16 KB).  Covers 0..30 KB.
#                    big:   tiny + banked code in W1     [TODO]
#                    huge:  small + banked code in W3    [TODO]
#                    manual: explicit via --memory-manual / --code-loc / --data-loc
#   --memory-manual SPEC   placement spec, only with --memory manual.
#                    SPEC = comma-separated KEY=VAL list:
#                      CODE=W1|W2  DATA=W1|W2|SAME  BANKED=W1|W3
#                    Example: --memory-manual CODE=W2,DATA=W2,BANKED=W3
#   -I PATH          additional include path (repeatable)
#   -L 0xADDR        code load address (default: derived from --memory mode)
#   -E 0xADDR        entry address       (default: same as -L)
#   -S 0xADDR        stack address       (default: 0xBFFE)
#   --stack-size N   bytes reserved for the stack (default: ~1278, sets
#                    HEAP_TOP = stack_top + 1 - N to constrain heap growth)
#   --code-loc 0xN   override SDCC --code-loc (default: derived from --memory)
#   --data-loc 0xN   override SDCC --data-loc (default: derived from --memory)
#   -Wl FLAG         extra linker flag (repeatable)
#   --bank N=FILE.c  compile FILE.c as bank N; repeatable; pulls crt0_banked
#                    automatically and adds -Wl-b_BANKN=0x{N}C000
#   --mkexe FLAG     extra mkexe flag (repeatable; e.g. --mkexe -p --mkexe 0)
#   --debug          enable runtime diagnostics — defines DEBUG_RT for both
#                    sdcc (-DDEBUG_RT) and the crt0 assembly (prepended
#                    `DEBUG_RT = 1`).  Exposes runtime introspection symbols
#                    such as `_w2_self_allocated` (see <sprinter.h>).
#   -v               verbose (echo every command)
#   -h, --help       show this help
#
# Example:
#   sprinter-cc -o hello.exe hello.c
#   sprinter-cc -o app.exe main.c --bank 1=engine.c
#   sprinter-cc -o tiny.exe tiny.c --crt0=minimal

set -eo pipefail

# ------- Locate the toolchain ------------------------------------------------
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
PROJ_ROOT="$( cd "$SCRIPT_DIR/.." && pwd )"

SDCC_BIN="$PROJ_ROOT/third_party/sdcc/bin"
SDCC="$SDCC_BIN/sdcc"
SDASZ80="$SDCC_BIN/sdasz80"
MKEXE="$PROJ_ROOT/toolchain/mkexe/mkexe"
CHECK_BANKS="$PROJ_ROOT/toolchain/check_banks.py"

RUNTIME="$PROJ_ROOT/runtime"
LIB_DIR="$PROJ_ROOT/lib"
INC_DIR="$PROJ_ROOT/libc/include"

# ------- Defaults ------------------------------------------------------------
OUT=""
CRT0_TYPE="default"
MEMORY_MODE="tiny"
MEMORY_MANUAL=""
# CODE_LOC / DATA_LOC are derived from MEMORY_MODE unless overridden via
# explicit --code-loc / --data-loc; capture user intent in USER_*.
USER_CODE_LOC=""
USER_DATA_LOC=""
LOAD_ADDR=""
ENTRY_ADDR=""
STACK_ADDR="0xBFFE"
STACK_SIZE=""        # if set, used to compute HEAP_TOP = STACK_ADDR + 1 - STACK_SIZE
DEBUG_RT=0           # if 1, prepend `DEBUG_RT = 1` to crt0 + pass -DDEBUG_RT to sdcc
VERBOSE=0
USER_INCS=()
SOURCES=()
LD_EXTRA=()
MKEXE_EXTRA=()
BANK_SPECS=()    # entries like "1=engine.c"

# ------- Parse args ----------------------------------------------------------
usage() {
    sed -n '3,28p' "$0" | sed 's/^# \{0,1\}//'
    exit "${1:-0}"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        -o)            OUT="$2";          shift 2;;
        --crt0=*)      CRT0_TYPE="${1#*=}"; shift;;
        -I)            USER_INCS+=("-I" "$2"); shift 2;;
        -L)            LOAD_ADDR="$2";    shift 2;;
        -E)            ENTRY_ADDR="$2";   shift 2;;
        -S)            STACK_ADDR="$2";   shift 2;;
        --code-loc)    USER_CODE_LOC="$2"; shift 2;;
        --data-loc)    USER_DATA_LOC="$2"; shift 2;;
        --memory)      MEMORY_MODE="$2";  shift 2;;
        --memory-manual) MEMORY_MANUAL="$2"; shift 2;;
        --stack-size)  STACK_SIZE="$2";   shift 2;;
        -Wl)           LD_EXTRA+=("$2");  shift 2;;
        --bank)        BANK_SPECS+=("$2"); shift 2;;
        --mkexe)       MKEXE_EXTRA+=("$2"); shift 2;;
        --debug)       DEBUG_RT=1;        shift;;
        -v)            VERBOSE=1;         shift;;
        -h|--help)     usage 0;;
        -*)            echo "sprinter-cc: unknown option: $1" >&2; usage 1;;
        *)             SOURCES+=("$1");   shift;;
    esac
done

[[ -z "$OUT" ]]         && { echo "sprinter-cc: -o NAME is required" >&2; exit 1; }
[[ ${#SOURCES[@]} -eq 0 ]] && { echo "sprinter-cc: no input files" >&2; exit 1; }

# ------- Resolve memory mode → CODE_LOC / DATA_LOC ---------------------------
# tiny  : CODE in W2 (0x8100), DATA auto after code (= W2)
# small : CODE in W1 (0x4100), DATA in W2 (0x8000)        — crt0 must alloc W2
# big   : tiny  + banked code in W1 (--bank stays in HOME-style trampolines)
# huge  : small + banked code in W3 — same as today's crt0_banked layout
# manual: parsed from --memory-manual SPEC
case "$MEMORY_MODE" in
    tiny)
        MODE_CODE_LOC="0x8100"; MODE_DATA_LOC="0"
        ;;
    small)
        # CODE at 0x4100; DATA/BSS/HEAP chained auto by linker.  Supports
        # any size up to ~30 KB code+data combined.  crt0_small reads port
        # 0xA2 at startup to detect if DSS already mapped W2 (program > 16
        # KB) and skips SETWIN2 in that case; otherwise it allocates W2
        # itself via ESTEX $3D + $3A.
        MODE_CODE_LOC="0x4100"; MODE_DATA_LOC="0"
        ;;
    big)
        # tiny layout (CODE+DATA in W2) + banked code in W1.  crt0_banked
        # gets BANK_W1=1 prepended so its trampoline + bank loader use
        # port 0xA2 (W1) and load address 0x4000.  mkexe gets -B 0x4000.
        MODE_CODE_LOC="0x8100"; MODE_DATA_LOC="0"
        ;;
    huge)
        # small layout + banked code in W3.  crt0_banked auto-detects W2
        # the same way crt0_small does, then loads banks from the .EXE.
        MODE_CODE_LOC="0x4100"; MODE_DATA_LOC="0x8000"
        ;;
    manual)
        # Defaults if SPEC omits a key.
        MODE_CODE_LOC="0x4100"
        MODE_DATA_LOC="0"
        IFS=',' read -ra _parts <<< "$MEMORY_MANUAL"
        for _p in "${_parts[@]}"; do
            [[ -z "$_p" ]] && continue
            _k="${_p%%=*}"; _v="${_p#*=}"
            case "$_k=$_v" in
                CODE=W1)    MODE_CODE_LOC="0x4100";;
                CODE=W2)    MODE_CODE_LOC="0x8100";;
                DATA=W1)    MODE_DATA_LOC="0x4000";;
                DATA=W2)    MODE_DATA_LOC="0x8000";;
                DATA=SAME)  MODE_DATA_LOC="0";;
                BANKED=W1|BANKED=W3) : ;;  # recorded; banked window only matters for --bank trampolines
                *)
                    echo "sprinter-cc: bad --memory-manual entry: $_p" >&2
                    echo "  valid: CODE=W1|W2, DATA=W1|W2|SAME, BANKED=W1|W3" >&2
                    exit 1
                    ;;
            esac
        done
        ;;
    *)
        echo "sprinter-cc: unknown --memory mode: $MEMORY_MODE" >&2
        echo "  valid: tiny | small | big | huge | manual" >&2
        exit 1
        ;;
esac

# Explicit --code-loc / --data-loc win over the mode defaults.
CODE_LOC="${USER_CODE_LOC:-$MODE_CODE_LOC}"
DATA_LOC="${USER_DATA_LOC:-$MODE_DATA_LOC}"

LOAD_ADDR="${LOAD_ADDR:-$CODE_LOC}"
ENTRY_ADDR="${ENTRY_ADDR:-$LOAD_ADDR}"

[[ $VERBOSE -eq 1 ]] && echo "  memory: mode=$MEMORY_MODE  code=$CODE_LOC  data=$DATA_LOC  load=$LOAD_ADDR"

# Pick crt0 source.  Memory mode drives default crt0 selection; explicit
# --crt0=TYPE still wins (e.g. --crt0=minimal for tiny w/o argv).
#
# big/huge ALWAYS use crt0_banked, even without explicit --bank flags:
#   big  — banks in W1 (BANK_W1=1 prepended) — code stays in W2 = tiny layout
#   huge — banks in W3 — code in W1 = small layout, crt0_banked auto-detects W2
# crt0_banked also exports startup-info symbols (estex_file_handle etc.).
case "$MEMORY_MODE" in
    small)      DEFAULT_CRT0="small";;
    big|huge)   DEFAULT_CRT0="banked";;
    *)          DEFAULT_CRT0="default";;   # tiny, manual
esac
if [[ "$CRT0_TYPE" == "default" && "$DEFAULT_CRT0" != "default" ]]; then
    CRT0_TYPE="$DEFAULT_CRT0"
fi

case "$CRT0_TYPE" in
    default) CRT0_SRC="$RUNTIME/crt0.s";        ;;
    minimal) CRT0_SRC="$RUNTIME/crt0_minimal.s";;
    small)   CRT0_SRC="$RUNTIME/crt0_small.s";  ;;
    banked)  CRT0_SRC="$RUNTIME/crt0_banked.s"; ;;
    *) echo "sprinter-cc: bad --crt0 type: $CRT0_TYPE" >&2; exit 1;;
esac

# Banked builds always need crt0_banked, regardless of what the user asked.
if [[ ${#BANK_SPECS[@]} -gt 0 ]] && [[ "$CRT0_TYPE" != "banked" ]]; then
    [[ $VERBOSE -eq 1 ]] && echo "  --bank present → switching to crt0_banked"
    CRT0_SRC="$RUNTIME/crt0_banked.s"
fi

# ------- Build in a per-output workdir ---------------------------------------
WORK="$(dirname "$OUT")/.sprinter-cc-$(basename "$OUT" .exe)"
mkdir -p "$WORK"

run() {
    [[ $VERBOSE -eq 1 ]] && echo "  $*"
    "$@"
}

# Decide whether the build needs the W1 banking window (BIG mode).  This
# affects crt0_banked.s and bank.s — both honour `BANK_W1 = 1` prepended
# by us.  HUGE / default (no BANK_W1) → banks live in W3 at port 0xE2.
BANK_W1=0
[[ "$MEMORY_MODE" == "big" ]] && BANK_W1=1

# Helper: assemble a runtime .s with optional `DEBUG_RT = 1` and/or
# `BANK_W1 = 1` prepended.  No-op prefix path keeps the original source
# (zero overhead for hot non-debug, non-banked builds).
asm_runtime() {
    local src="$1" out="$2"
    local prefix=""
    [[ $DEBUG_RT -eq 1 ]] && prefix+="DEBUG_RT = 1"$'\n'
    [[ $BANK_W1   -eq 1 ]] && prefix+="BANK_W1 = 1"$'\n'
    if [[ -n "$prefix" ]]; then
        local patched="$WORK/$(basename "$src" .s)_patched.s"
        { printf "%s" "$prefix"; cat "$src"; } > "$patched"
        run "$SDASZ80" -o "$out" "$patched"
    else
        run "$SDASZ80" -o "$out" "$src"
    fi
}

# 1. crt0 — patched with DEBUG_RT / BANK_W1 if requested.
CRT0_REL="$WORK/$(basename "$CRT0_SRC" .s).rel"
asm_runtime "$CRT0_SRC" "$CRT0_REL"

# 1b. heap_top — either default from runtime/, or generate custom if user
# specified --stack-size.  HEAP_TOP = stack_top + 1 - stack_size, e.g.
# --stack-size 2048 with default stack=0xBFFE → HEAP_TOP = 0xB7FF.
HEAP_TOP_REL="$WORK/heap_top.rel"
if [[ -n "$STACK_SIZE" ]]; then
    # Resolve numeric values (accept hex 0xNNN or decimal).
    _stack_top=$(( STACK_ADDR ))      # bash auto-detects 0x prefix
    _stack_sz=$(( STACK_SIZE ))
    _heap_top=$(( _stack_top + 1 - _stack_sz ))
    printf "    .module sprinter_heap_top\n    ___sdcc_heap_end = 0x%04X\n    .globl ___sdcc_heap_end\n" \
        "$_heap_top" > "$WORK/heap_top.s"
    HEAP_TOP_SRC="$WORK/heap_top.s"
    [[ $VERBOSE -eq 1 ]] && echo "  heap_top: custom HEAP_TOP=$(printf '0x%04X' $_heap_top) (stack reserve $_stack_sz bytes)"
else
    HEAP_TOP_SRC="$RUNTIME/heap_top.s"
fi
run "$SDASZ80" -o "$HEAP_TOP_REL" "$HEAP_TOP_SRC"

# 2. user sources → .rel (HOME)
USER_RELS=()
CC_FLAGS=(-mz80 --no-std-crt0 --std-c99 --opt-code-size -I "$INC_DIR" "${USER_INCS[@]}")
[[ $DEBUG_RT -eq 1 ]] && CC_FLAGS+=(-DDEBUG_RT)
for src in "${SOURCES[@]}"; do
    rel="$WORK/$(basename "$src" .c).rel"
    run "$SDCC" "${CC_FLAGS[@]}" -c -o "$rel" "$src"
    USER_RELS+=("$rel")
done

# 3. bank infrastructure — required whenever crt0_banked is in play.
#    Always assemble bank.s for _bank_pages + the bcall/bjump trampolines.
#    If the user did not pass any --bank, generate a tiny stub providing
#    n_banks = 0 (crt0_banked.s imports this symbol unconditionally).
#    In BIG mode banks live in W1 at low16 = 0x4000; in HUGE / default
#    mode they live in W3 at low16 = 0xC000.
BANK_RELS=()
BANK_LD_FLAGS=()
if [[ "$CRT0_TYPE" == "banked" ]]; then
    if [[ $BANK_W1 -eq 1 ]]; then
        BANK_LOW16=0x4000
    else
        BANK_LOW16=0xC000
    fi
    # Trampoline + _bank_pages table — depend on BANK_W1, so assemble per build.
    BANK_TRAMP_REL="$WORK/bank_trampoline.rel"
    asm_runtime "$RUNTIME/bank.s" "$BANK_TRAMP_REL"
    BANK_RELS+=("$BANK_TRAMP_REL")

    if [[ ${#BANK_SPECS[@]} -eq 0 ]]; then
        # User has no banks — supply n_banks=0 so crt0_banked links.
        STUB_C="$WORK/_n_banks_stub.c"
        STUB_REL="$WORK/_n_banks_stub.rel"
        echo "const unsigned char n_banks = 0;" > "$STUB_C"
        run "$SDCC" "${CC_FLAGS[@]}" -c -o "$STUB_REL" "$STUB_C"
        BANK_RELS+=("$STUB_REL")
    else
        for spec in "${BANK_SPECS[@]}"; do
            bank_n="${spec%%=*}"
            bank_src="${spec#*=}"
            rel="$WORK/bank${bank_n}_$(basename "$bank_src" .c).rel"
            run "$SDCC" "${CC_FLAGS[@]}" \
                --codeseg "BANK${bank_n}" --constseg "BANK${bank_n}" --dataseg "BANK${bank_n}" \
                -c -o "$rel" "$bank_src"
            BANK_RELS+=("$rel")
            # Virtual address: bank_n in upper byte, BANK_LOW16 in low half.
            addr=$(printf "0x%X" $(( (bank_n << 16) | BANK_LOW16 )))
            BANK_LD_FLAGS+=("-Wl-b_BANK${bank_n}=${addr}")
        done
    fi
fi

# 4. link → .ihx
IHX="$WORK/$(basename "$OUT" .exe).ihx"
LINK_FLAGS=(-mz80 --no-std-crt0 --std-c99 --opt-code-size
            --code-loc "$CODE_LOC" --data-loc "$DATA_LOC")
LINK_FLAGS+=("${BANK_LD_FLAGS[@]}")
for f in "${LD_EXTRA[@]}"; do LINK_FLAGS+=("$f"); done
# libsprinter.lib via -l/-L (sdcc passes -lsprinter through to sdldz80).
#
# Filter benign "?ASlink-Warning-Definition ... found more than once"
# warnings.  These come from intentional sprinter.lib overrides of the
# SDCC z80.lib defaults (puts, ___sdcc_heap, asctime/localtime, etc.).
# Linker picks the first found = our version, so behaviour is correct;
# the warning is just noise.  In verbose mode show everything.
if [[ $VERBOSE -eq 1 ]]; then
    run "$SDCC" "${LINK_FLAGS[@]}" -o "$IHX" \
        "$CRT0_REL" "$HEAP_TOP_REL" "${USER_RELS[@]}" "${BANK_RELS[@]}" \
        "-L$LIB_DIR" "-lsprinter"
else
    # Drop the warning line + its two follow-up "Library:" lines.
    run "$SDCC" "${LINK_FLAGS[@]}" -o "$IHX" \
        "$CRT0_REL" "$HEAP_TOP_REL" "${USER_RELS[@]}" "${BANK_RELS[@]}" \
        "-L$LIB_DIR" "-lsprinter" 2>&1 \
        | awk '
            /^\?ASlink-Warning-Definition of public symbol/ { skip = 3 }
            skip > 0 { skip--; next }
            { print }
        '
fi

# Quick bank-size check (only meaningful if there are banks).
if [[ ${#BANK_SPECS[@]} -gt 0 ]] && [[ -f "${IHX%.ihx}.map" ]]; then
    python3 "$CHECK_BANKS" "${IHX%.ihx}.map" || true
fi

# 5. mkexe → .exe.  In BIG mode tell mkexe banks live at 0x4000 (W1).
# IHX is the positional input — it must come last.  Build the prefix
# (options + extras) first, then append the positional.
MK_PREFIX=()
[[ $VERBOSE -eq 1 ]] && MK_PREFIX+=(-v)
[[ $BANK_W1 -eq 1 ]] && MK_PREFIX+=(-B 0x4000)
MK_PREFIX+=("${MKEXE_EXTRA[@]}")
MK_PREFIX+=(-L "$LOAD_ADDR" -E "$ENTRY_ADDR" -S "$STACK_ADDR" -o "$OUT")
run "$MKEXE" "${MK_PREFIX[@]}" "$IHX"

echo "sprinter-cc: wrote $OUT"