#!/bin/bash
set -e

# Train 500KB zstd dictionary from parakeet_historical posts

DICT_SIZE=524288  # 500KB (512 * 1024)
DICT_VERSION=1
OUTPUT_DIR="./parakeet-db/src/dicts"
TEMP_DIR="/tmp/parakeet_dict_training"

echo "=========================================="
echo "Post Content Dictionary Training"
echo "=========================================="
echo "Dictionary size: 500KB"
echo "Database: parakeet_historical"
echo ""

# Create directories
mkdir -p "$OUTPUT_DIR"
mkdir -p "$TEMP_DIR"

# Extract all posts from historical database
echo "[1/4] Extracting posts from parakeet_historical..."
psql parakeet_historical -c "
    COPY (
        SELECT content
        FROM posts
        WHERE content IS NOT NULL
            AND LENGTH(content) > 0
    ) TO STDOUT
" > "$TEMP_DIR/all_posts.txt"

# Show stats
TOTAL_SIZE=$(stat -f%z "$TEMP_DIR/all_posts.txt" 2>/dev/null || stat -c%s "$TEMP_DIR/all_posts.txt")
TOTAL_LINES=$(wc -l < "$TEMP_DIR/all_posts.txt")
TOTAL_SIZE_MB=$(awk "BEGIN {printf \"%.2f\", $TOTAL_SIZE / 1048576}")

echo "  Extracted: $TOTAL_LINES posts"
echo "  Total size: ${TOTAL_SIZE_MB} MB"
echo ""

# Train dictionary
# Use -B to split file into blocks (posts are ~75 bytes avg, use 256 byte blocks)
echo "[2/4] Training 500KB dictionary..."
zstd --train "$TEMP_DIR/all_posts.txt" \
    -o "$OUTPUT_DIR/post_content_v${DICT_VERSION}.dict" \
    --maxdict=$DICT_SIZE \
    -B256

DICT_ACTUAL_SIZE=$(stat -f%z "$OUTPUT_DIR/post_content_v${DICT_VERSION}.dict" 2>/dev/null || stat -c%s "$OUTPUT_DIR/post_content_v${DICT_VERSION}.dict")
DICT_SIZE_KB=$(awk "BEGIN {printf \"%.2f\", $DICT_ACTUAL_SIZE / 1024}")

echo "  Dictionary created: ${DICT_SIZE_KB} KB"
echo ""

# Test compression ratio
echo "[3/4] Testing compression ratio..."

# Create test sample (10K posts)
head -n 10000 "$TEMP_DIR/all_posts.txt" > "$TEMP_DIR/test_sample.txt"
TEST_SIZE=$(stat -f%z "$TEMP_DIR/test_sample.txt" 2>/dev/null || stat -c%s "$TEMP_DIR/test_sample.txt")

# Compress without dictionary
zstd -q -c "$TEMP_DIR/test_sample.txt" > "$TEMP_DIR/test_no_dict.zst"
NO_DICT_SIZE=$(stat -f%z "$TEMP_DIR/test_no_dict.zst" 2>/dev/null || stat -c%s "$TEMP_DIR/test_no_dict.zst")

# Compress with dictionary
zstd -q -D "$OUTPUT_DIR/post_content_v${DICT_VERSION}.dict" \
    -c "$TEMP_DIR/test_sample.txt" > "$TEMP_DIR/test_with_dict.zst"
WITH_DICT_SIZE=$(stat -f%z "$TEMP_DIR/test_with_dict.zst" 2>/dev/null || stat -c%s "$TEMP_DIR/test_with_dict.zst")

# Calculate ratios
NO_DICT_RATIO=$(awk "BEGIN {printf \"%.2f\", $TEST_SIZE / $NO_DICT_SIZE}")
WITH_DICT_RATIO=$(awk "BEGIN {printf \"%.2f\", $TEST_SIZE / $WITH_DICT_SIZE}")
IMPROVEMENT=$(awk "BEGIN {printf \"%.2f\", $WITH_DICT_RATIO - $NO_DICT_RATIO}")

echo "  Test sample size:    $TEST_SIZE bytes"
echo "  Without dictionary:  $NO_DICT_SIZE bytes (${NO_DICT_RATIO}x compression)"
echo "  With dictionary:     $WITH_DICT_SIZE bytes (${WITH_DICT_RATIO}x compression)"
echo "  Improvement:         +${IMPROVEMENT}x"
echo ""

# Cleanup
echo "[4/4] Cleaning up..."
rm -rf "$TEMP_DIR"
echo ""

echo "=========================================="
echo "Dictionary training complete!"
echo "=========================================="
echo "Location: $OUTPUT_DIR/post_content_v${DICT_VERSION}.dict"
echo "Size: ${DICT_SIZE_KB} KB"
echo "Compression ratio: ${WITH_DICT_RATIO}x"
echo ""
echo "Next steps:"
echo "1. Review the compression ratio above"
echo "2. Run diesel migration to update schema"
echo "3. Update consumer and parakeet code"