#!/bin/bash
# Sourced by install.sh after it has resolved the local or tagged repository.
# Keep this file declarative: function definitions only, with no top-level actions.

# Model integrity verification
#
# Models are 40MB-2GB downloads that end up being loaded by native code, so no
# model is installed without matching a digest pinned in this repository.
# src/vocalinux/utils/model_checksums.txt is the same manifest the application
# uses at runtime; it is regenerated by scripts/generate-model-checksums.py.
#
# A model that cannot be verified is deleted and the function returns 1, which
# callers already treat as "leave it for first run" rather than aborting the
# install. The app re-downloads and re-verifies it later.
# ---------------------------------------------------------------------------

# Print the digest of $1 using algorithm $2, trying the tools most likely present.
compute_file_digest() {
    local file="$1" algo="$2"

    case "$algo" in
        sha256)
            if command_exists sha256sum; then sha256sum "$file" | cut -d' ' -f1
            elif command_exists shasum; then shasum -a 256 "$file" | cut -d' ' -f1
            elif command_exists openssl; then openssl dgst -sha256 "$file" | awk '{print $NF}'
            else return 1; fi
            ;;
        *)
            return 1
            ;;
    esac
}

# Compare $1 against algorithm $2 and digest $3. $4 labels the file in messages.
verify_digest() {
    local file="$1" algo="$2" expected="$3" label="$4" actual

    if ! actual=$(compute_file_digest "$file" "$algo") || [ -z "$actual" ]; then
        print_error "Cannot compute the $algo digest of $label: no sha256sum, shasum or openssl found."
        return 1
    fi

    if [ "$actual" != "$expected" ]; then
        print_error "$label failed $algo verification."
        print_error "  expected: $expected"
        print_error "  actual:   $actual"
        print_error "The file does not match the digest pinned in this release."
        return 1
    fi

    print_success "$label verified ($algo)"
    return 0
}

# Verify $1 against the manifest entry named $2 (defaults to $1's basename).
verify_model_checksum() {
    local file="$1"
    local key="${2:-$(basename "$file")}"
    local algo expected size actual_size

    if [ ! -f "$MODEL_CHECKSUMS_FILE" ]; then
        print_error "Checksum manifest not found at $MODEL_CHECKSUMS_FILE."
        print_error "Cannot verify $key; refusing to install an unverified model."
        return 1
    fi

    # Manifest columns: filename  algorithm  digest  size-in-bytes
    algo=$(awk -v k="$key" '$1==k {print $2; exit}' "$MODEL_CHECKSUMS_FILE")
    expected=$(awk -v k="$key" '$1==k {print $3; exit}' "$MODEL_CHECKSUMS_FILE")
    size=$(awk -v k="$key" '$1==k {print $4; exit}' "$MODEL_CHECKSUMS_FILE")

    if [ -z "$algo" ] || [ -z "$expected" ]; then
        print_error "No checksum is pinned for $key in $MODEL_CHECKSUMS_FILE."
        print_error "Regenerate it with scripts/generate-model-checksums.py."
        return 1
    fi

    # Size first: it is free, and it reports a truncated download as truncation
    # rather than as a digest mismatch that reads like tampering.
    if [ -n "$size" ] && [ "$size" != "0" ]; then
        actual_size=$(stat -c%s "$file" 2>/dev/null || stat -f%z "$file" 2>/dev/null || echo "")
        if [ -n "$actual_size" ] && [ "$actual_size" != "$size" ]; then
            print_error "$key is $actual_size bytes, expected $size (truncated download?)."
            return 1
        fi
    fi

    verify_digest "$file" "$algo" "$expected" "$key"
}

# Verify an OpenAI Whisper checkpoint against the sha256 embedded in its own URL
# (they publish each file under a path segment that is its digest).
verify_openai_model_checksum() {
    local file="$1" url="$2" label="$3" expected

    expected=$(printf '%s\n' "$url" | grep -oE '/[0-9a-f]{64}/' | tr -d '/' | head -n1)
    if [ -z "$expected" ]; then
        print_error "$url carries no sha256 path segment; refusing to install an unverified model."
        return 1
    fi

    verify_digest "$file" "sha256" "$expected" "$label"
}

# Download $1 to $2, preferring wget and falling back to curl. $3 labels the
# model in error messages. Leaves no partial file behind on failure.
download_model_file() {
    local url="$1" dest="$2" label="$3"

    if command_exists wget; then
        if ! wget --progress=bar:force:noscroll --tries=3 --timeout=60 -O "$dest" "$url" 2>&1; then
            print_error "Failed to download $label with wget"
            rm -f "$dest"
            return 1
        fi
    elif command_exists curl; then
        # -f: fail on HTTP errors instead of saving the error page as the model
        if ! curl -fL --progress-bar --retry 3 --retry-delay 2 -o "$dest" "$url"; then
            print_error "Failed to download $label with curl"
            rm -f "$dest"
            return 1
        fi
    else
        print_error "Neither wget nor curl is available to download $label"
        return 1
    fi

    if [ ! -s "$dest" ]; then
        print_error "Downloaded $label is empty or missing"
        rm -f "$dest"
        return 1
    fi
}

# Print the digest pinned for manifest entry $1, or nothing when unpinned.
pinned_digest_for() {
    [ -f "$MODEL_CHECKSUMS_FILE" ] || return 1
    awk -v k="$1" '$1==k {print $3; exit}' "$MODEL_CHECKSUMS_FILE"
}

# The Hugging Face commit the whisper.cpp digests were taken at.
whispercpp_pinned_revision() {
    [ -f "$MODEL_CHECKSUMS_FILE" ] || return 1
    awk '/^#[[:space:]]*whispercpp-revision:/ {print $3; exit}' "$MODEL_CHECKSUMS_FILE"
}

# False when the cloned release predates the manifest. That app has no runtime
# verification either, so refusing to pre-download would protect nothing.
model_verification_available() {
    [ -f "$MODEL_CHECKSUMS_FILE" ]
}

# Function to download and install Whisper tiny model
install_whisper_model() {
    print_info "Installing Whisper tiny model (~75MB)..."

    # Create whisper models directory
    local WHISPER_DIR="$DATA_DIR/models/whisper"
    mkdir -p "$WHISPER_DIR"

    # Whisper tiny model URL and path
    local TINY_MODEL_URL="https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
    local TINY_MODEL_PATH="$WHISPER_DIR/tiny.pt"

    # An existing file is not a verified file: it may predate checksum
    # verification, or have been replaced since. Hash it before trusting it,
    # and re-download rather than keep something that does not match.
    if [ -f "$TINY_MODEL_PATH" ]; then
        if verify_openai_model_checksum "$TINY_MODEL_PATH" "$TINY_MODEL_URL" "Whisper tiny model"; then
            print_info "Whisper tiny model already exists at $TINY_MODEL_PATH"
            return 0
        fi
        print_warning "The existing Whisper tiny model does not match its pinned digest; replacing it."
        rm -f "$TINY_MODEL_PATH"
    fi

    # Check internet connectivity
    if ! command -v wget >/dev/null 2>&1 && ! command -v curl >/dev/null 2>&1; then
        print_warning "Neither wget nor curl found. Cannot download Whisper model."
        print_warning "Model will be downloaded on first application run."
        return 1
    fi

    # Test internet connectivity (HTTP probes; ICMP ping is often blocked)
    if ! check_connectivity; then
        print_warning "No internet connection detected (HTTP probes to pypi.org/api.github.com/huggingface.co failed)."
        print_warning "Whisper model will be downloaded on first application run."
        return 1
    fi

    print_info "Downloading Whisper tiny model..."
    print_info "This may take a few minutes depending on your internet connection."

    local TEMP_FILE="$VOCULINUX_TMP_DIR/tiny.pt"

    if ! download_model_file "$TINY_MODEL_URL" "$TEMP_FILE" "Whisper model"; then
        return 1
    fi

    # Verify before the file reaches its final location, so a failed download is
    # never installed. (The "already exists" path above hashes what it finds.)
    if ! verify_openai_model_checksum "$TEMP_FILE" "$TINY_MODEL_URL" "Whisper tiny model"; then
        rm -f "$TEMP_FILE"
        print_warning "Whisper model will be downloaded and verified on first application run."
        return 1
    fi

    # Move to final location
    mv "$TEMP_FILE" "$TINY_MODEL_PATH"

    # Verify the model file
    if [ -f "$TINY_MODEL_PATH" ]; then
        local MODEL_SIZE=$(du -h "$TINY_MODEL_PATH" | cut -f1)
        print_success "Whisper tiny model installed successfully ($MODEL_SIZE)"

        # Create a marker file to indicate this model was pre-installed
        echo "$(date)" > "$WHISPER_DIR/.voculinux_preinstalled"

        return 0
    else
        print_error "Whisper model installation failed"
        return 1
    fi
}

# Function to download and install VOSK models
install_vosk_models() {
    print_info "Installing VOSK speech recognition models..."

    # Create models directory
    local MODELS_DIR="$DATA_DIR/models"
    mkdir -p "$MODELS_DIR"

    # Define model information
    local SMALL_MODEL_URL="https://alphacephei.com/vosk/models/vosk-model-small-en-us-0.15.zip"
    local SMALL_MODEL_NAME="vosk-model-small-en-us-0.15"
    local SMALL_MODEL_PATH="$MODELS_DIR/$SMALL_MODEL_NAME"

    local SMALL_MODEL_ARCHIVE="$SMALL_MODEL_NAME.zip"
    # The extracted tree has no digest of its own — the pin covers the zip, which
    # is deleted after unpacking. So whoever unpacks it records the verified zip
    # digest in a stamp (here, and in the app's own downloader), and the directory
    # is trusted only while that stamp matches what we pin today.
    local SMALL_MODEL_STAMP="$SMALL_MODEL_PATH/.voculinux_verified"
    local EXPECTED_ZIP_DIGEST
    EXPECTED_ZIP_DIGEST=$(pinned_digest_for "$SMALL_MODEL_ARCHIVE" || true)
    if [ -d "$SMALL_MODEL_PATH" ]; then
        if [ -n "$EXPECTED_ZIP_DIGEST" ] && [ -f "$SMALL_MODEL_STAMP" ] &&
           [ "$(cat "$SMALL_MODEL_STAMP" 2>/dev/null)" = "$EXPECTED_ZIP_DIGEST" ]; then
            print_info "Small VOSK model already exists at $SMALL_MODEL_PATH (verified)"
            return 0
        fi
        # Unverified is not the same as known-bad, and this runs on every install:
        # a tree unpacked before stamps existed, or by the app itself, simply has
        # no stamp. Deleting it here would trade a working model for no model
        # whenever the download, unzip or network that follows fails. Fetch and
        # verify the replacement first; the swap below is what removes this tree.
        print_warning "The existing VOSK model carries no matching verification stamp; re-downloading it."
    fi

    # Refuse before spending the download, not after verification rejects it.
    if [ -z "$EXPECTED_ZIP_DIGEST" ]; then
        print_error "No checksum is pinned for $SMALL_MODEL_ARCHIVE in $MODEL_CHECKSUMS_FILE."
        print_warning "VOSK model will be downloaded and verified on first application run."
        return 1
    fi

    # Check internet connectivity
    if ! command -v wget >/dev/null 2>&1 && ! command -v curl >/dev/null 2>&1; then
        print_warning "Neither wget nor curl found. Cannot download VOSK models."
        print_warning "Models will be downloaded on first application run."
        return 1
    fi

    # Test internet connectivity (HTTP probes; ICMP ping is often blocked)
    if ! check_connectivity; then
        print_warning "No internet connection detected (HTTP probes to pypi.org/api.github.com/huggingface.co failed)."
        print_warning "VOSK models will be downloaded on first application run."
        return 1
    fi

    print_info "Downloading small VOSK model (approximately 40MB)..."
    print_info "This may take a few minutes depending on your internet connection."

    local TEMP_ZIP="$VOCULINUX_TMP_DIR/$(basename $SMALL_MODEL_URL)"

    if ! download_model_file "$SMALL_MODEL_URL" "$TEMP_ZIP" "VOSK model"; then
        return 1
    fi

    # Verify before extracting: a zip that fails its pinned digest must not get
    # as far as writing files into the models directory.
    if ! verify_model_checksum "$TEMP_ZIP"; then
        rm -f "$TEMP_ZIP"
        print_warning "VOSK model will be downloaded and verified on first application run."
        return 1
    fi

    print_info "Extracting VOSK model..."

    if ! command -v unzip >/dev/null 2>&1; then
        print_error "unzip command not found. Cannot extract VOSK model."
        rm -f "$TEMP_ZIP"
        return 1
    fi

    # Unpack beside the model rather than over it: any model already installed
    # stays usable until a complete, verified replacement exists, and staging in
    # MODELS_DIR keeps the swap a same-filesystem rename.
    local STAGING_DIR="$MODELS_DIR/.vosk-staging.$$"
    # A run killed mid-swap leaves its scratch directories behind; they are named
    # so they can be recognised and are of no use to a later run.
    find "$MODELS_DIR" -maxdepth 1 -type d \
        \( -name '.vosk-staging.*' -o -name '.*.replaced.*' \) \
        -exec rm -rf {} + 2>/dev/null || true
    if ! mkdir -p "$STAGING_DIR"; then
        print_error "Failed to create a staging directory under $MODELS_DIR"
        rm -f "$TEMP_ZIP"
        return 1
    fi

    if ! unzip -q "$TEMP_ZIP" -d "$STAGING_DIR"; then
        print_error "Failed to extract VOSK model"
        rm -f "$TEMP_ZIP"
        rm -rf "$STAGING_DIR"
        return 1
    fi
    rm -f "$TEMP_ZIP"

    if [ ! -d "$STAGING_DIR/$SMALL_MODEL_NAME" ]; then
        print_error "VOSK model extraction failed - $SMALL_MODEL_NAME not found in the archive"
        rm -rf "$STAGING_DIR"
        return 1
    fi

    chmod -R 755 "$STAGING_DIR/$SMALL_MODEL_NAME"
    echo "$(date)" > "$STAGING_DIR/$SMALL_MODEL_NAME/.voculinux_preinstalled"

    # Record the digest this tree was extracted from, so a later run can tell a
    # verified model from one that merely exists. Written before the swap: an
    # unstamped tree is one this function would download all over again.
    if [ -z "$EXPECTED_ZIP_DIGEST" ] ||
       ! printf '%s\n' "$EXPECTED_ZIP_DIGEST" > "$STAGING_DIR/$SMALL_MODEL_NAME/.voculinux_verified"; then
        print_error "Could not record the verified digest for $SMALL_MODEL_NAME"
        rm -rf "$STAGING_DIR"
        return 1
    fi

    # Swap. The old tree is moved aside rather than deleted, so a failed rename
    # can put it back.
    local REPLACED_DIR="$MODELS_DIR/.$SMALL_MODEL_NAME.replaced.$$"
    if [ -d "$SMALL_MODEL_PATH" ] && ! mv "$SMALL_MODEL_PATH" "$REPLACED_DIR"; then
        print_error "Could not move the existing VOSK model aside; keeping it"
        rm -rf "$STAGING_DIR"
        return 1
    fi
    if ! mv "$STAGING_DIR/$SMALL_MODEL_NAME" "$SMALL_MODEL_PATH"; then
        print_error "Failed to install the verified VOSK model"
        if [ -d "$REPLACED_DIR" ]; then
            mv "$REPLACED_DIR" "$SMALL_MODEL_PATH"
        fi
        rm -rf "$STAGING_DIR"
        return 1
    fi
    rm -rf "$REPLACED_DIR" "$STAGING_DIR"

    print_success "VOSK small model installed successfully at $SMALL_MODEL_PATH"
    return 0
}

# Function to download and install whisper.cpp tiny model
install_whispercpp_model() {
    print_info "Installing whisper.cpp tiny model (~39MB)..."

    # Create whisper.cpp models directory
    local WHISPERCPP_DIR="$DATA_DIR/models/whispercpp"
    mkdir -p "$WHISPERCPP_DIR"

    # whisper.cpp tiny model URL and path. The Hugging Face revision is pinned to
    # the one the digests in model_checksums.txt were taken at, so upstream
    # replacing ggml-tiny.bin cannot turn every install into a checksum failure.
    local WHISPERCPP_REVISION
    WHISPERCPP_REVISION=$(whispercpp_pinned_revision)
    if [ -z "$WHISPERCPP_REVISION" ]; then
        print_error "No whisper.cpp revision is pinned in $MODEL_CHECKSUMS_FILE."
        print_warning "whisper.cpp model will be downloaded and verified on first application run."
        return 1
    fi
    local TINY_MODEL_URL="https://huggingface.co/ggerganov/whisper.cpp/resolve/$WHISPERCPP_REVISION/ggml-tiny.bin"
    local TINY_MODEL_PATH="$WHISPERCPP_DIR/ggml-tiny.bin"

    # An existing file is not a verified file: it may predate checksum
    # verification, or have been replaced since. Hash it before trusting it,
    # and re-download rather than keep something that does not match.
    if [ -f "$TINY_MODEL_PATH" ]; then
        if verify_model_checksum "$TINY_MODEL_PATH" "ggml-tiny.bin"; then
            print_info "whisper.cpp tiny model already exists at $TINY_MODEL_PATH"
            return 0
        fi
        print_warning "The existing whisper.cpp tiny model does not match its pinned digest; replacing it."
        rm -f "$TINY_MODEL_PATH"
    fi

    # Check internet connectivity
    if ! command -v wget >/dev/null 2>&1 && ! command -v curl >/dev/null 2>&1; then
        print_warning "Neither wget nor curl found. Cannot download whisper.cpp model."
        print_warning "Model will be downloaded on first application run."
        return 1
    fi

    # Test internet connectivity (HTTP probes; ICMP ping is often blocked)
    if ! check_connectivity; then
        print_warning "No internet connection detected (HTTP probes to pypi.org/api.github.com/huggingface.co failed)."
        print_warning "whisper.cpp model will be downloaded on first application run."
        return 1
    fi

    print_info "Downloading whisper.cpp tiny model..."
    print_info "This may take a few minutes depending on your internet connection."

    local TEMP_FILE="$VOCULINUX_TMP_DIR/ggml-tiny.bin"

    if ! download_model_file "$TINY_MODEL_URL" "$TEMP_FILE" "whisper.cpp model"; then
        return 1
    fi

    # Verify before the file reaches its final location, so a failed download is
    # never installed. (The "already exists" path above hashes what it finds.)
    if ! verify_model_checksum "$TEMP_FILE" "ggml-tiny.bin"; then
        rm -f "$TEMP_FILE"
        print_warning "whisper.cpp model will be downloaded and verified on first application run."
        return 1
    fi

    # Move to final location
    mv "$TEMP_FILE" "$TINY_MODEL_PATH"

    # Verify the model file
    if [ -f "$TINY_MODEL_PATH" ]; then
        local MODEL_SIZE=$(du -h "$TINY_MODEL_PATH" | cut -f1)
        print_success "whisper.cpp tiny model installed successfully ($MODEL_SIZE)"

        # Create a marker file to indicate this model was pre-installed
        echo "$(date)" > "$WHISPERCPP_DIR/.voculinux_preinstalled"

        return 0
    else
        print_error "whisper.cpp model installation failed"
        return 1
    fi
}