#!/usr/bin/env bash
set -euo pipefail

# En Croissant + lc0 + Maia + Stockfish setup script for Arch Linux.
#
# Idempotent: re-runs skip healthy components, replace broken ones, and
# preserve tuned settings and engine IDs in engines.json. See
# setup-chess-README.md in this directory for usage and known issues.

APPIMAGE_DIR="$HOME/.local/bin"
APPIMAGE_PATH="$APPIMAGE_DIR/en-croissant.AppImage"
LC0_BIN="$HOME/.local/bin/lc0"
LC0_VERSION="v0.32.1"
MAIA_WEIGHTS_DIR="$HOME/.local/share/maia"
MAIA_BASE_URL="https://github.com/CSSLab/maia-chess/releases/download/v1.0"
ENCROISSANT_VERSION="0.15.0"
ENCROISSANT_APPIMAGE="en-croissant_${ENCROISSANT_VERSION}_amd64.AppImage"
ENCROISSANT_SHA256="fc3afc9fcace62a1b2dfefd06f81038991ff828e7afa6d0ab7f6f8e26018aff3"
ENCROISSANT_URL="https://github.com/franciscoBSalgueiro/en-croissant/releases/download/v${ENCROISSANT_VERSION}/${ENCROISSANT_APPIMAGE}"
STOCKFISH_VERSION="sf_18"
STOCKFISH_RELEASE_URL="https://github.com/official-stockfish/Stockfish/releases/download/${STOCKFISH_VERSION}"
STOCKFISH_DIR="$HOME/.local/share/org.encroissant.app/engines/stockfish"
ENGINES_JSON="$HOME/.local/share/org.encroissant.app/engines/engines.json"
DESKTOP_FILE="$HOME/.local/share/applications/en-croissant.desktop"

STOCKFISH_FLAVOR=""
STOCKFISH_TAR=""
STOCKFISH_TAR_URL=""
STOCKFISH_TAR_SHA256=""
STOCKFISH_BIN=""
FORCE=0
COMPONENT=""
LC0_BUILD_DIR=""
TMP_PATHS=()

declare -A STOCKFISH_SHA256=(
    [vnni512]="91d89e0e387faa78607fa00c0ed783d00023b4fd45181230c6b58624a5a65463"
    [avx512]="9ddf29d8589d65cbf4b444c29241e711fccac24e6e9f5452480bb9a2e0fddd83"
    [avxvnni]="57fb4139ea786c5cb606099cd901bf363d57e0e00c5a3dc7cd98269a6baf041a"
    [bmi2]="7b200a3cd8ae6e2b07386cd213058edc91faf05ff77db68604d2f5143c56b69e"
    [avx2]="536c0c2c0cf06450df0bfb5e876ef0d3119950703a8f143627f990c7b5417964"
    [sse41-popcnt]="dea5016a6d9ab705e5697b093d882fca4677d84d8828f470ee33e76de33cf962"
    [x86-64]="5c6f38b02a4da5f3ffe763f27da6c3e743eebefd92b50cb3661623b96696adff"
)

declare -A MAIA_SHA256=(
    [1100]="e1cf1cd0c96b8a4fa6a275f4b9fd54ed1ffebf9fe44641b9fceded310e9619c4"
    [1200]="ead4ba953f233ae732999ebc1e2b675378148527ebcfad2f0acbc5e4c224d98e"
    [1300]="36195f87bf4761834baa0bf87472b18509a7261a9d7d6f1a8443261369a733f2"
    [1400]="d5353ea6766356dad2d28920c6692f37a5f30963767f1a3105d33b4d0af011e8"
    [1500]="35ab6f20421d59e1df3b17c5a5016947af4c6761368ef84044a9a9c7619a9a00"
    [1600]="d2c9e5948581acf4b9fc0b1e720c5dc0fe64ce80cfc4a239d3f8a42e1176c876"
    [1700]="d277eacd792d340a30abb464dc65127254e65cac57abca17facc469889b96478"
    [1800]="0031ad7c4256b1fd09fbebd28418d644d68b26cd2a45df4967ccf5c7ec9c4965"
    [1900]="e2f565f42d7cd9f122557e6dc4eb84e5bbaedceda1d404dc485d3611c7c97a12"
)

VALID_COMPONENTS=(prereqs en-croissant lc0 maia stockfish engines-json desktop)

info()  { printf '\033[1;34m==> %s\033[0m\n' "$*"; }
ok()    { printf '\033[1;32m  -> %s\033[0m\n' "$*"; }
skip()  { printf '\033[1;33m  -> %s (skipped)\033[0m\n' "$*"; }
warn()  { printf '\033[1;33mWARN: %s\033[0m\n' "$*" >&2; }
err()   { printf '\033[1;31mERROR: %s\033[0m\n' "$*" >&2; }

usage() {
    cat <<EOF
Usage: ${0##*/} [--force] [--component <name>]

Install En Croissant, lc0, Maia weights, and Stockfish for the current user.
Idempotent: existing healthy components are skipped; settings tuned by the
user are preserved across re-runs. See setup-chess-README.md for details.

Options:
  -f, --force              Re-download/rebuild even if components look healthy
      --component <name>   Run only one section; one of:
                           ${VALID_COMPONENTS[*]}
  -h, --help               Show this help
EOF
}

cleanup() {
    if [[ -n "${LC0_BUILD_DIR:-}" && -d "$LC0_BUILD_DIR" ]]; then
        rm -rf "$LC0_BUILD_DIR"
    fi

    if (( ${#TMP_PATHS[@]} )); then
        for path in "${TMP_PATHS[@]}"; do
            [[ -e "$path" ]] && rm -rf "$path"
        done
    fi

    return 0
}
trap cleanup EXIT

download_file() {
    local url=$1
    local dest=$2
    local sha256=${3:-}
    local tmp

    mkdir -p "$(dirname "$dest")"
    tmp=$(mktemp "$(dirname "$dest")/.download.XXXXXX")
    TMP_PATHS+=("$tmp")

    curl -fL --retry 3 --retry-delay 2 -o "$tmp" "$url"
    if [[ -n "$sha256" ]]; then
        printf '%s  %s\n' "$sha256" "$tmp" | sha256sum -c -
    fi
    mv "$tmp" "$dest"
}

has_cpu_flag() {
    local flag=$1

    [[ -r /proc/cpuinfo ]] && grep -qm1 -E "(^| )${flag}( |$)" /proc/cpuinfo
}

# UCI smoke test: send `uci` + `quit`, expect the engine's `id name <prefix>`
# line on stdout within a few seconds. Validates the binary actually loads
# and responds, not just that the file is executable.
#
# Output is buffered to a variable rather than piped into `grep -q` so the
# engine writes its full response (closing cleanly on `quit`) — `grep -q`
# closes stdin on first match and would SIGPIPE the engine, which pipefail
# would then count as a failure.
smoke_test_uci() {
    local bin=$1
    local id_prefix=$2
    local out

    [[ -x "$bin" ]] || return 1
    out=$({ printf 'uci\nquit\n'; } | timeout 5 "$bin" 2>/dev/null) || return 1
    grep -qE "^id name ${id_prefix}" <<< "$out"
}

aur_helper() {
    command -v yay 2>/dev/null || command -v paru 2>/dev/null || true
}

select_stockfish() {
    if has_cpu_flag avx512_vnni; then
        STOCKFISH_FLAVOR="vnni512"
    elif has_cpu_flag avx512f; then
        STOCKFISH_FLAVOR="avx512"
    elif has_cpu_flag avx_vnni; then
        STOCKFISH_FLAVOR="avxvnni"
    elif has_cpu_flag avx2; then
        STOCKFISH_FLAVOR="avx2"
    elif has_cpu_flag bmi2; then
        STOCKFISH_FLAVOR="bmi2"
    elif has_cpu_flag sse4_1 && has_cpu_flag popcnt; then
        STOCKFISH_FLAVOR="sse41-popcnt"
    else
        STOCKFISH_FLAVOR="x86-64"
    fi

    if [[ "$STOCKFISH_FLAVOR" == "x86-64" ]]; then
        STOCKFISH_TAR="stockfish-ubuntu-x86-64.tar"
    else
        STOCKFISH_TAR="stockfish-ubuntu-x86-64-${STOCKFISH_FLAVOR}.tar"
    fi

    STOCKFISH_TAR_URL="${STOCKFISH_RELEASE_URL}/${STOCKFISH_TAR}"
    STOCKFISH_TAR_SHA256="${STOCKFISH_SHA256[$STOCKFISH_FLAVOR]}"
    STOCKFISH_BIN="$STOCKFISH_DIR/${STOCKFISH_TAR%.tar}"
}

should_run() {
    [[ -z "$COMPONENT" || "$COMPONENT" == "$1" ]]
}

valid_component() {
    local c=$1 v
    for v in "${VALID_COMPONENTS[@]}"; do
        [[ "$v" == "$c" ]] && return 0
    done
    return 1
}

while (($#)); do
    case "$1" in
        -f|--force)
            FORCE=1
            ;;
        --component)
            shift
            if (($# == 0)); then
                err "--component requires an argument"
                exit 2
            fi
            if ! valid_component "$1"; then
                err "Unknown component: $1. Valid: ${VALID_COMPONENTS[*]}"
                exit 2
            fi
            COMPONENT=$1
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            err "Unknown option: $1"
            usage
            exit 2
            ;;
    esac
    shift
done

select_stockfish

# ---------------------------------------------------------------------------
# 1. Check prerequisites
# ---------------------------------------------------------------------------
if should_run prereqs; then
info "Checking prerequisites"

missing=()
declare -A missing_seen=()

add_missing_package() {
    local pkg=$1

    if [[ -z "${missing_seen[$pkg]:-}" ]]; then
        missing+=("$pkg")
        missing_seen[$pkg]=1
    fi
}

for pkg in git curl tar coreutils grep desktop-file-utils fuse2; do
    pacman -Qi "$pkg" &>/dev/null || add_missing_package "$pkg"
done

# Only require the build toolchain when a healthy system lc0 is unavailable.
# The AUR path is tried first when possible, but it can fail; these packages
# keep the source-build fallback ready.
need_source_build=1
if smoke_test_uci /usr/bin/lc0 'Lc0'; then
    need_source_build=0
fi

if (( need_source_build )); then
    for pkg in base-devel meson ninja python pkgconf protobuf zlib openblas eigen; do
        pacman -Qi "$pkg" &>/dev/null || add_missing_package "$pkg"
    done
fi

if (( ${#missing[@]} )); then
    info "Installing missing packages: ${missing[*]}"
    if (( EUID == 0 )); then
        pacman -S --needed --noconfirm "${missing[@]}"
    else
        sudo pacman -S --needed --noconfirm "${missing[@]}"
    fi
fi
ok "All prerequisites found"
fi

# ---------------------------------------------------------------------------
# 2. Install En Croissant AppImage
# ---------------------------------------------------------------------------
if should_run en-croissant; then
info "Installing En Croissant AppImage"

if [[ -f "$APPIMAGE_PATH" && "$FORCE" -eq 0 ]]; then
    skip "Already exists at $APPIMAGE_PATH"
else
    mkdir -p "$APPIMAGE_DIR"
    echo "  Downloading: $ENCROISSANT_URL"
    download_file "$ENCROISSANT_URL" "$APPIMAGE_PATH" "$ENCROISSANT_SHA256"
    chmod +x "$APPIMAGE_PATH"
    ok "Installed to $APPIMAGE_PATH"
fi
fi

# ---------------------------------------------------------------------------
# 3. Install lc0
#
# Resolution order, each gated by a smoke test:
#   a) Existing $LC0_BIN passes UCI handshake — keep it.
#   b) System /usr/bin/lc0 passes — symlink it into $LC0_BIN.
#   c) AUR helper available — install lc0 + lc0-network-sm via AUR.
#   d) Source build from upstream tag.
# ---------------------------------------------------------------------------
if should_run lc0; then
info "Installing lc0"

if smoke_test_uci "$LC0_BIN" 'Lc0' && [[ "$FORCE" -eq 0 ]]; then
    skip "Already healthy at $LC0_BIN"
else
    mkdir -p "$HOME/.local/bin"

    # Attempt AUR install when a system lc0 isn't already healthy.
    if ! smoke_test_uci /usr/bin/lc0 'Lc0'; then
        helper=$(aur_helper)
        if [[ -n "$helper" ]]; then
            # Surface the BLAS-provider conflict before pacman aborts the
            # transaction non-interactively. See setup-chess-README.md.
            if pacman -Qi blas &>/dev/null && ! pacman -Qi blas-openblas &>/dev/null; then
                warn "lc0 requires the blas-openblas BLAS provider, which conflicts with the installed 'blas' package."
                warn "Resolve manually before re-running, then re-invoke this script:"
                warn "  sudo pacman -S blas-openblas    # prompts to remove blas/cblas/lapack"
                warn "Falling through to source build path for this run."
            else
                info "Trying AUR install: $helper -S lc0 lc0-network-sm"
                if "$helper" -S --needed --noconfirm lc0 lc0-network-sm; then
                    ok "Installed lc0 via $helper"
                else
                    warn "AUR install failed; falling back to source build"
                fi
            fi
        fi
    fi

    if smoke_test_uci /usr/bin/lc0 'Lc0'; then
        rm -f "$LC0_BIN"
        ln -s /usr/bin/lc0 "$LC0_BIN"
        ok "Linked $LC0_BIN -> /usr/bin/lc0"
    else
        info "Building lc0 from source (this takes several minutes)"
        LC0_BUILD_DIR=$(mktemp -d /tmp/lc0-build.XXXXXX)
        echo "  Cloning to $LC0_BUILD_DIR"
        git clone --branch "$LC0_VERSION" --recurse-submodules \
            https://github.com/LeelaChessZero/lc0.git "$LC0_BUILD_DIR/lc0"
        pushd "$LC0_BUILD_DIR/lc0" >/dev/null
        INSTALL_PREFIX="$HOME/.local" ./build.sh
        popd >/dev/null
        rm -f "$LC0_BIN"
        cp "$LC0_BUILD_DIR/lc0/build/release/lc0" "$LC0_BIN"
        chmod +x "$LC0_BIN"
        ok "Built and installed $LC0_BIN"
        cleanup
        LC0_BUILD_DIR=""
    fi

    if ! smoke_test_uci "$LC0_BIN" 'Lc0'; then
        err "lc0 at $LC0_BIN failed its UCI handshake after install."
        err "Run '$LC0_BIN' manually to see the load-time error."
        exit 1
    fi
fi
fi

# ---------------------------------------------------------------------------
# 4. Download Maia weights (with SHA256 pinning)
# ---------------------------------------------------------------------------
if should_run maia; then
info "Downloading Maia weights"
mkdir -p "$MAIA_WEIGHTS_DIR"

for elo in 1100 1200 1300 1400 1500 1600 1700 1800 1900; do
    file="maia-${elo}.pb.gz"
    dest="$MAIA_WEIGHTS_DIR/$file"
    if [[ -f "$dest" && "$FORCE" -eq 0 ]]; then
        # Validate existing file matches the pinned hash; redownload if not.
        if printf '%s  %s\n' "${MAIA_SHA256[$elo]}" "$dest" | sha256sum -c --status -; then
            skip "$file already exists (sha256 OK)"
            continue
        else
            warn "$file sha256 mismatch; re-downloading"
        fi
    fi
    echo "  Downloading $file"
    download_file "$MAIA_BASE_URL/$file" "$dest" "${MAIA_SHA256[$elo]}"
    ok "$file"
done

# Maia wrapper scripts — thin shell wrappers that exec lc0 with the
# matching weights file. Regenerated unconditionally so they always point
# at the current $LC0_BIN.
info "Creating Maia wrapper scripts"
mkdir -p "$HOME/.local/bin"

for elo in 1100 1200 1300 1400 1500 1600 1700 1800 1900; do
    wrapper="$HOME/.local/bin/maia-${elo}"
    cat > "$wrapper" <<SCRIPT
#!/bin/sh
exec "$LC0_BIN" --weights="$MAIA_WEIGHTS_DIR/maia-${elo}.pb.gz" "\$@"
SCRIPT
    chmod +x "$wrapper"
done
ok "Created 9 wrapper scripts in ~/.local/bin/"
fi

# ---------------------------------------------------------------------------
# 5. Install Stockfish
# ---------------------------------------------------------------------------
if should_run stockfish; then
info "Installing Stockfish"
echo "  Selected build: $STOCKFISH_FLAVOR"

if smoke_test_uci "$STOCKFISH_BIN" 'Stockfish' && [[ "$FORCE" -eq 0 ]]; then
    skip "Already healthy at $STOCKFISH_BIN"
else
    STOCKFISH_TMP_TAR=$(mktemp /tmp/stockfish.XXXXXX.tar)
    STOCKFISH_TMP_DIR=$(mktemp -d /tmp/stockfish.XXXXXX)
    TMP_PATHS+=("$STOCKFISH_TMP_TAR" "$STOCKFISH_TMP_DIR")
    echo "  Downloading Stockfish"
    download_file "$STOCKFISH_TAR_URL" "$STOCKFISH_TMP_TAR" "$STOCKFISH_TAR_SHA256"
    tar -xf "$STOCKFISH_TMP_TAR" -C "$STOCKFISH_TMP_DIR" --strip-components=1
    if [[ ! -f "$STOCKFISH_TMP_DIR/$(basename "$STOCKFISH_BIN")" ]]; then
        err "Stockfish archive did not contain $(basename "$STOCKFISH_BIN")"
        exit 1
    fi
    mkdir -p "$STOCKFISH_DIR"
    cp -a "$STOCKFISH_TMP_DIR"/. "$STOCKFISH_DIR"/
    chmod +x "$STOCKFISH_BIN"

    if ! smoke_test_uci "$STOCKFISH_BIN" 'Stockfish'; then
        err "Stockfish at $STOCKFISH_BIN failed its UCI handshake after install."
        exit 1
    fi
    ok "Installed to $STOCKFISH_BIN"
fi
fi

# ---------------------------------------------------------------------------
# 6. Write engines.json
#
# Preserves existing engine IDs and user-tuned settings on re-runs. Unknown
# entries in the existing file are kept too so manually-added engines (e.g.
# Komodo) survive a re-run.
# ---------------------------------------------------------------------------
if should_run engines-json; then
info "Writing engines.json"
mkdir -p "$(dirname "$ENGINES_JSON")"
if [[ -f "$ENGINES_JSON" ]]; then
    backup="$ENGINES_JSON.bak.$(date +%Y%m%d%H%M%S)"
    cp -p "$ENGINES_JSON" "$backup"
    ok "Backed up existing engines.json to $backup"
fi

python3 - "$STOCKFISH_BIN" "$STOCKFISH_VERSION" "$STOCKFISH_TAR_URL" \
    "$HOME/.local/bin" "$ENGINES_JSON" <<'PYEOF' > "$ENGINES_JSON.new"
import json, os, sys, uuid

stockfish_bin, stockfish_version, stockfish_download_link, \
    maia_bin_dir, engines_json_path = sys.argv[1:6]

existing = {}
existing_order = []
if os.path.exists(engines_json_path):
    try:
        with open(engines_json_path) as f:
            for e in json.load(f):
                name = e.get("name")
                if name:
                    existing[name] = e
                    existing_order.append(name)
    except (json.JSONDecodeError, OSError):
        pass

def merge(name, defaults):
    prev = existing.get(name, {})
    out = dict(defaults)
    out["name"] = name
    out["id"] = prev.get("id", str(uuid.uuid4()))
    # Preserve user-tuned settings if a prior entry exists; otherwise use
    # the defaults this script ships.
    if "settings" in prev:
        out["settings"] = prev["settings"]
    return out

managed = []
managed_names = set()

stockfish_defaults = {
    "type": "local",
    "version": stockfish_version,
    "path": stockfish_bin,
    "image": "https://upload.wikimedia.org/wikipedia/commons/3/3a/NewLogoSF.png",
    "elo": 3635,
    "downloadSize": 79953920,
    "downloadLink": stockfish_download_link,
    "loaded": True,
    "go": {"t": "Infinite"},
    "settings": [
        {"name": "Threads", "value": 1},
        {"name": "Hash", "value": 16},
        {"name": "MultiPV", "value": 1},
    ],
}
managed.append(merge("Stockfish", stockfish_defaults))
managed_names.add("Stockfish")

for elo in range(1100, 2000, 100):
    name = f"Maia {elo}"
    maia_defaults = {
        "type": "local",
        "version": "1.0",
        "path": f"{maia_bin_dir}/maia-{elo}",
        "image": None,
        "elo": elo,
        "loaded": True,
        "go": {"t": "Nodes", "c": 1},
        "settings": [{"name": "MultiPV", "value": 1}],
    }
    managed.append(merge(name, maia_defaults))
    managed_names.add(name)

# Append any non-managed engines from the existing file (manually-added
# user engines) so a re-run never silently drops them.
extras = [existing[n] for n in existing_order if n not in managed_names]

print(json.dumps(managed + extras, indent=4))
PYEOF

mv "$ENGINES_JSON.new" "$ENGINES_JSON"
ok "Wrote $ENGINES_JSON"
fi

# ---------------------------------------------------------------------------
# 7. Create .desktop file
# ---------------------------------------------------------------------------
if should_run desktop; then
info "Creating .desktop file"
mkdir -p "$(dirname "$DESKTOP_FILE")"

cat > "$DESKTOP_FILE" <<EOF
[Desktop Entry]
Name=En Croissant
Comment=Chess GUI for analysis, puzzles, and play
Exec=$APPIMAGE_PATH
Icon=chess
Type=Application
Categories=Game;BoardGame;
StartupWMClass=En Croissant
EOF

update-desktop-database "$HOME/.local/share/applications/" 2>/dev/null || true
ok "Created $DESKTOP_FILE"
fi

# ---------------------------------------------------------------------------
# Done
# ---------------------------------------------------------------------------
echo ""
info "Setup complete!"
echo "  En Croissant: $APPIMAGE_PATH"
echo "  lc0:          $LC0_BIN"
echo "  Stockfish:    $STOCKFISH_BIN"
echo "  Maia weights: $MAIA_WEIGHTS_DIR/"
echo "  engines.json: $ENGINES_JSON"
echo ""
echo "  Launch with:  $APPIMAGE_PATH"
echo ""
echo "  See setup-chess-README.md for known issues (BLAS provider conflict,"
echo "  Maia-as-Black wtime quirk in En Croissant)."