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

# Set up the Artec Studuino sketchbook and its command-line workflow.

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ARDUINO_USER_OVERRIDE="${ARDUINO_USER:-}"
ARDUINO_USER="${ARDUINO_USER:-$HOME/Arduino}"
RESOURCE_DIR="$SCRIPT_DIR/Resources"
LIBRARY_ARCHIVE="$RESOURCE_DIR/studuino_v1_lib.zip"
HARDWARE_ARCHIVE="$RESOURCE_DIR/arduino_IDE_hardware_settings_ver17x.zip"
FQBN="studuino:avr:studuino"
BUILD_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/arduino/studuino-compile-db"
VERIFY_BUILD_DIR="$BUILD_DIR/verify"

FORCE=0
CLEAN_LIBRARIES=0
SKIP_PACKAGES=0
SKIP_PERMISSIONS=0
SKIP_COMPILE_DB=0
WITH_MONITOR=0
CONFIG_FILE="${ARDUINO_CONFIG_FILE:-}"
TEMP_DIRS=()

cleanup() {
    local temporary
    for temporary in "${TEMP_DIRS[@]}"; do
        [[ -d "$temporary" ]] && rm -rf -- "$temporary"
    done
}
trap cleanup EXIT

usage() {
    cat <<'EOF'
Usage: ./setup.sh [options]

Options:
  --force              Replace existing Studuino library/hardware files.
  --clean              Remove user libraries not supplied by the Studuino bundle.
  --skip-packages      Do not install Arch packages.
  --skip-permissions   Do not modify serial-port group membership.
  --skip-compile-db    Do not generate compile_commands.json.
  --with-monitor       Also install picocom for a serial monitor.
  -h, --help           Show this help.

Environment:
  ARDUINO_USER         Arduino sketchbook directory (default: ~/Arduino).
  ARDUINO_CONFIG_FILE  Optional arduino-cli configuration file.
EOF
}

while (($# > 0)); do
    case "$1" in
        --force)
            FORCE=1
            ;;
        --clean)
            CLEAN_LIBRARIES=1
            ;;
        --skip-packages)
            SKIP_PACKAGES=1
            ;;
        --skip-permissions)
            SKIP_PERMISSIONS=1
            ;;
        --skip-compile-db)
            SKIP_COMPILE_DB=1
            ;;
        --with-monitor)
            WITH_MONITOR=1
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            printf 'Unknown option: %s\n' "$1" >&2
            usage >&2
            exit 2
            ;;
    esac
    shift
done

if ((EUID == 0)); then
    ROOT_CMD=()
    TARGET_USER="${SUDO_USER:-$(id -un)}"
else
    ROOT_CMD=(sudo)
    TARGET_USER="${USER:-$(id -un)}"
fi

CLI=(arduino-cli)
if [[ -n "$CONFIG_FILE" ]]; then
    CLI+=(--config-file "$CONFIG_FILE")
fi

info() {
    printf '[studuino] %s\n' "$*"
}

warn() {
    printf '[studuino] WARNING: %s\n' "$*" >&2
}

die() {
    printf '[studuino] ERROR: %s\n' "$*" >&2
    exit 1
}

as_root() {
    if ((${#ROOT_CMD[@]} > 0)); then
        "${ROOT_CMD[@]}" "$@"
    else
        "$@"
    fi
}

run_cli() {
    "${CLI[@]}" "$@"
}

require_command() {
    command -v "$1" >/dev/null 2>&1 || die "Required command not found: $1"
}

install_packages() {
    if ((SKIP_PACKAGES)); then
        info "Skipping Arch package installation."
    else
        require_command pacman
        local -a packages=(
            arduino-cli
            clangd
            make
            neovim
            python
            unzip
        )
        if ((WITH_MONITOR)); then
            packages+=(picocom)
        fi
        info "Installing required Arch packages."
        as_root pacman -S --needed "${packages[@]}"
    fi

    require_command arduino-cli
    require_command python3
    require_command unzip
}

ensure_serial_permissions() {
    if ((SKIP_PERMISSIONS)); then
        info "Skipping serial-port group membership changes."
        return
    fi

    require_command id
    require_command getent

    local current_groups
    current_groups="$(id -nG "$TARGET_USER" 2>/dev/null || true)"
    local changed=0
    local group

    for group in uucp dialout; do
        if ! getent group "$group" >/dev/null 2>&1; then
            warn "Serial-port group '$group' does not exist; skipping it."
            continue
        fi
        if ! grep -qw -- "$group" <<<"$current_groups"; then
            info "Adding $TARGET_USER to serial-port group '$group'."
            as_root usermod -a -G "$group" "$TARGET_USER"
            changed=1
        fi
    done

    if ((changed)); then
        warn "Log out and back in before uploading to a serial device."
    fi
}

init_cli_config() {
    if [[ -n "$CONFIG_FILE" ]]; then
        mkdir -p -- "$(dirname -- "$CONFIG_FILE")"
        if [[ ! -f "$CONFIG_FILE" ]]; then
            info "Creating arduino-cli config: $CONFIG_FILE"
            arduino-cli config init --dest-file "$CONFIG_FILE" >/dev/null
        fi
    elif [[ ! -f "$HOME/.arduino15/arduino-cli.yaml" ]]; then
        info "Creating the default arduino-cli config."
        arduino-cli config init >/dev/null
    fi
}

resolve_arduino_user() {
    if [[ -n "$ARDUINO_USER_OVERRIDE" ]]; then
        run_cli config set directories.user "$ARDUINO_USER" >/dev/null
        info "Using configured Arduino sketchbook: $ARDUINO_USER"
        return
    fi

    local configured_user
    configured_user="$(run_cli config get directories.user 2>/dev/null | sed -n '1p' || true)"
    if [[ -n "$configured_user" && "$configured_user" != "null" ]]; then
        ARDUINO_USER="$configured_user"
    fi
}

ensure_hardware_layout() {
    local vendor_dir="$ARDUINO_USER/hardware/studuino"
    local hardware_dir="$vendor_dir/avr"
    local nested_dir="$vendor_dir/studuino/avr"

    # Older versions of the setup instructions could create this extra level.
    if [[ -d "$nested_dir" ]]; then
        if [[ ! -e "$hardware_dir" ]]; then
            info "Moving the legacy nested Studuino hardware path."
            mv -- "$nested_dir" "$hardware_dir"
            rmdir -- "$vendor_dir/studuino" 2>/dev/null || true
        elif ((FORCE)); then
            info "Removing the redundant nested Studuino hardware path."
            rm -rf -- "$nested_dir"
            rmdir -- "$vendor_dir/studuino" 2>/dev/null || true
        else
            die "Both $hardware_dir and $nested_dir exist; re-run with --force to remove the nested copy."
        fi
    fi

    if [[ -f "$hardware_dir/boards.txt" && -f "$hardware_dir/platform.txt" ]] && (( FORCE == 0 )); then
        info "Studuino hardware files are already installed; keeping them."
    else
        [[ -f "$HARDWARE_ARCHIVE" ]] || die "Hardware archive not found: $HARDWARE_ARCHIVE"
        mkdir -p -- "$hardware_dir"

        # The official archive was created on Windows and uses backslashes in ZIP
        # entry names. Extract only the Linux/other Studuino platform files.
    python3 - "$HARDWARE_ARCHIVE" "$hardware_dir" "$FORCE" <<'PY'
from pathlib import Path
import sys
from zipfile import ZipFile

archive = Path(sys.argv[1])
destination = Path(sys.argv[2])
force = sys.argv[3] == "1"

wanted = {
    "arduino_ide_settings_ver17x/other/hardware/studuino/avr/boards.txt": "boards.txt",
    "arduino_ide_settings_ver17x/other/hardware/studuino/avr/platform.txt": "platform.txt",
}
found = set()

with ZipFile(archive) as zipped:
    for info in zipped.infolist():
        normalized = info.filename.replace("\\", "/")
        target_name = wanted.get(normalized)
        if target_name is None:
            continue
        target = destination / target_name
        if target.exists() and not force:
            print(f"[studuino] Keeping existing {target}")
        else:
            target.write_bytes(zipped.read(info))
            print(f"[studuino] Installed {target}")
        found.add(target_name)

missing = set(wanted.values()) - found
if missing:
    raise SystemExit("Missing hardware files in archive: " + ", ".join(sorted(missing)))
PY
    fi

    # The archive targets the legacy IDE's tool paths. These overrides make it
    # work with the AVR tools installed by arduino-cli.
    local local_properties="$hardware_dir/platform.local.txt"
    cat >"$local_properties" <<'EOF'
# Generated by setup.sh for modern arduino-cli AVR tool paths.
compiler.path={runtime.tools.avr-gcc.path}/bin/
archive_file_path={build.path}/{archive_file}
recipe.ar.pattern="{compiler.path}{compiler.ar.cmd}" {compiler.ar.flags} {compiler.ar.extra_flags} "{archive_file_path}" "{object_file}"
tools.avrdude.path={runtime.tools.avrdude.path}
tools.avrdude.cmd.path={path}/bin/avrdude
tools.avrdude.config.path={path}/etc/avrdude.conf
EOF
    info "Wrote modern CLI compatibility overrides: $local_properties"

    grep -q '^studuino.name=' "$hardware_dir/boards.txt" \
        || die "The installed Studuino boards.txt is missing its board definition."
    grep -q '^name=' "$hardware_dir/platform.txt" \
        || die "The installed Studuino platform.txt is missing its platform name."
}

ensure_avr_core() {
    if run_cli core list | awk '$1 == "arduino:avr" { found = 1 } END { exit !found }'; then
        info "AVR core is already installed."
        return
    fi

    info "Installing the Arduino AVR core."
    run_cli core update-index
    run_cli core install arduino:avr
}

install_studuino_libraries() {
    [[ -f "$LIBRARY_ARCHIVE" ]] || die "Library archive not found: $LIBRARY_ARCHIVE"

    local temporary
    temporary="$(mktemp -d)"
    TEMP_DIRS+=("$temporary")

    unzip -q "$LIBRARY_ARCHIVE" -d "$temporary"
    local bundle_dir="$temporary/studuino_v1_lib"
    [[ -d "$bundle_dir/Studuino" ]] || die "Unexpected Studuino library archive layout."

    mkdir -p -- "$ARDUINO_USER/libraries"
    local source_dir destination name
    for source_dir in "$bundle_dir"/*; do
        [[ -d "$source_dir" ]] || continue
        name="${source_dir##*/}"
        destination="$ARDUINO_USER/libraries/$name"

        if [[ -L "$destination" ]]; then
            die "Refusing to replace symlinked library: $destination"
        fi
        if [[ -e "$destination" ]] && ((FORCE == 0)); then
            info "Keeping existing Studuino library: $destination"
        else
            if [[ -e "$destination" ]]; then
                rm -rf -- "$destination"
            fi
            cp -a -- "$source_dir" "$destination"
            info "Installed Studuino library: $destination"
        fi
    done

    if ((CLEAN_LIBRARIES)); then
        local -A keep=()
        for source_dir in "$bundle_dir"/*; do
            [[ -d "$source_dir" ]] || continue
            keep["${source_dir##*/}"]=1
        done
        # Servo is required by Studuino even though it is not in the ZIP.
        keep[Servo]=1

        local installed_dir installed_name
        for installed_dir in "$ARDUINO_USER/libraries"/*; do
            [[ -d "$installed_dir" && ! -L "$installed_dir" ]] || continue
            installed_name="${installed_dir##*/}"
            if [[ -z "${keep[$installed_name]+present}" ]]; then
                rm -rf -- "$installed_dir"
                info "Removed non-Studuino library: $installed_dir"
            fi
        done
    fi
}

ensure_servo() {
    if [[ -n "$(find "$ARDUINO_USER/libraries" -type f -name Servo.h -print -quit 2>/dev/null)" ]]; then
        info "Servo library is already installed."
        return
    fi

    if run_cli lib list | awk '$1 == "Servo" { found = 1 } END { exit !found }'; then
        info "Servo library is already registered."
        return
    fi

    info "Installing Servo, required by Studuino.h."
    run_cli lib install Servo
}

ensure_project_files() {
    local metadata="$SCRIPT_DIR/sketch.yaml"
    if [[ ! -f "$metadata" ]]; then
        cat >"$metadata" <<EOF
default_fqbn: $FQBN
EOF
        info "Created $metadata"
    elif grep -q '^default_fqbn:' "$metadata"; then
        sed -i -E "s/^default_fqbn:.*/default_fqbn: $FQBN/" "$metadata"
    elif grep -q '^fqbn:' "$metadata"; then
        sed -i -E "s/^fqbn:.*/default_fqbn: $FQBN/" "$metadata"
    else
        printf '\ndefault_fqbn: %s\n' "$FQBN" >>"$metadata"
    fi

    local sketch="$SCRIPT_DIR/Arduino.ino"
    if [[ ! -f "$sketch" ]]; then
        if [[ -f "$SCRIPT_DIR/sketch.ino" ]]; then
            info "Renaming sketch.ino to Arduino.ino."
            mv -- "$SCRIPT_DIR/sketch.ino" "$sketch"
        else
            cat >"$sketch" <<'EOF'
#include <Arduino.h>

void setup() {
  // Add Studuino initialization here.
}

void loop() {
  // Add Studuino code here.
}
EOF
            info "Created a minimal sketch at $sketch"
        fi
    elif [[ -f "$SCRIPT_DIR/sketch.ino" ]]; then
        warn "Both Arduino.ino and sketch.ino exist; leaving sketch.ino untouched."
    fi
}

generate_compile_database() {
    if ((SKIP_COMPILE_DB)); then
        info "Skipping compile_commands.json generation."
        return
    fi

    mkdir -p -- "$BUILD_DIR"
    info "Generating the Neovim/clangd compilation database."
    if run_cli compile --clean --fqbn "$FQBN" --only-compilation-database \
        --build-path "$BUILD_DIR" "$SCRIPT_DIR"; then
        cp -- "$BUILD_DIR/compile_commands.json" "$SCRIPT_DIR/compile_commands.json"
        info "Wrote $SCRIPT_DIR/compile_commands.json"
    else
        warn "Could not generate compile_commands.json; the Studuino setup is otherwise complete."
    fi
}

verify_setup() {
    info "Verifying board discovery..."
    run_cli board listall | grep -F "$FQBN" >/dev/null \
        || die "arduino-cli cannot discover $FQBN."

    info "Verifying the Studuino and Servo libraries..."
    run_cli lib list | grep -E '^Studuino[[:space:]]' >/dev/null \
        || die "The Studuino library is not visible to arduino-cli."
    run_cli lib list | grep -E '^Servo[[:space:]]' >/dev/null \
        || die "The Servo library is not visible to arduino-cli."

    info "Compiling the project for $FQBN..."
    run_cli compile --clean --fqbn "$FQBN" \
        --build-path "$VERIFY_BUILD_DIR" "$SCRIPT_DIR"
}

install_packages
ensure_serial_permissions
init_cli_config
resolve_arduino_user
ensure_hardware_layout
ensure_avr_core
install_studuino_libraries
ensure_servo
ensure_project_files
generate_compile_database
verify_setup

info "Studuino setup complete."
info "Compile with: arduino-cli compile"
info "Check the port with: arduino-cli board list"
info "Upload after connecting the board with: arduino-cli upload"