#!/bin/bash
# ==============================================================================
# Wasot Simple PGO Release Build Script
# ==============================================================================
#
# A simpler, more robust PGO build script that builds iteratively.
#
# Usage:
#   ./scripts/build/build_release_simple.sh                     # Default features
#   ./scripts/build/build_release_simple.sh --features sqlite   # With sqlite
#   ./scripts/build/build_release_simple.sh --all-features      # All non-python features
#   ./scripts/build/build_release_simple.sh --no-default-features # Minimal
#   ./scripts/build/build_release_simple.sh --clean             # Clean first
#   ./scripts/build/build_release_simple.sh --simple            # Skip PGO entirely
#
# ==============================================================================

set -euo pipefail

readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
readonly BINARY_NAME="wasot"
readonly RELEASE_DIR="${PROJECT_ROOT}/target/release"
readonly TARGET_TRIPLE="x86_64-unknown-linux-gnu"

# Colors
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly CYAN='\033[0;36m'
readonly RED='\033[0;31m'
readonly NC='\033[0m'

print_phase() { echo -e "\n${CYAN}>>> $1${NC}\n"; }
print_step() { echo -e "${BLUE}[*] $1${NC}"; }
print_success() { echo -e "${GREEN}[+] $1${NC}"; }
print_warning() { echo -e "${YELLOW}[!] $1${NC}"; }
print_error() { echo -e "${RED}[x] $1${NC}"; }

# Globals
FEATURE_FLAGS=""
DO_CLEAN=false
SIMPLE_BUILD=false

usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Options:
  --features FEATURES   Comma-separated features (e.g., sqlite,grpc)
  --all-features        Enable all non-python features (sqlite,grpc,optimized-alloc)
  --no-default-features Build without default features
  --clean               Clean before building
  --simple              Simple release build (no PGO)
  --help, -h            Show this help

Examples:
  $0                              # Build with default features (sqlite)
  $0 --no-default-features        # Minimal build
  $0 --features sqlite,grpc       # Specific features
  $0 --all-features               # All features except python
  $0 --simple --all-features      # Simple build with all features
EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --features)
                FEATURE_FLAGS="--features $2"
                shift 2
                ;;
            --all-features)
                # Note: python feature requires maturin, not cargo build
                FEATURE_FLAGS="--no-default-features --features sqlite,grpc,optimized-alloc"
                shift
                ;;
            --no-default-features)
                FEATURE_FLAGS="--no-default-features"
                shift
                ;;
            --clean)
                DO_CLEAN=true
                shift
                ;;
            --simple)
                SIMPLE_BUILD=true
                shift
                ;;
            --help|-h)
                usage
                ;;
            *)
                print_error "Unknown option: $1"
                usage
                ;;
        esac
    done
}

check_prerequisites() {
    print_phase "Checking Prerequisites"

    if ! command -v cargo &>/dev/null; then
        print_error "cargo not found"
        exit 1
    fi
    print_success "cargo available"

    if [[ "$SIMPLE_BUILD" == false ]]; then
        if ! command -v cargo-pgo &>/dev/null; then
            print_warning "cargo-pgo not installed. Installing..."
            cargo install cargo-pgo || {
                print_error "Failed to install cargo-pgo"
                exit 1
            }
        fi
        print_success "cargo-pgo available"

        if ! rustup component list --installed | grep -q "llvm-tools"; then
            print_step "Installing llvm-tools component..."
            rustup component add llvm-tools
        fi
        print_success "llvm-tools available"
    fi
}

clean_artifacts() {
    print_phase "Cleaning Build Artifacts"
    cargo clean 2>/dev/null || true
    rm -rf "$RELEASE_DIR"
    print_success "Clean complete"
}

# Simple build without PGO
build_simple() {
    print_phase "Simple Release Build"
    print_step "Features: ${FEATURE_FLAGS:-default}"

    # shellcheck disable=SC2086
    RUSTFLAGS="-C target-cpu=native" cargo build --release $FEATURE_FLAGS --bin "$BINARY_NAME"

    mkdir -p "$RELEASE_DIR"
    local src_bin="./target/release/$BINARY_NAME"

    if [[ ! -f "$src_bin" ]]; then
        print_error "Binary not found at $src_bin"
        exit 1
    fi

    cp "$src_bin" "$RELEASE_DIR/"
    print_success "Binary built: $RELEASE_DIR/$BINARY_NAME"
}

# Find the PGO binary (more robust than original script)
find_pgo_binary() {
    local profile="${1:-release}"
    local bin

    # Try target triple path first (cargo-pgo uses this)
    bin="${PROJECT_ROOT}/target/${TARGET_TRIPLE}/${profile}/${BINARY_NAME}"
    if [[ -f "$bin" && -x "$bin" ]]; then
        echo "$bin"
        return 0
    fi

    # Fallback: search for the binary
    bin=$(find "${PROJECT_ROOT}/target" -name "$BINARY_NAME" \
        -path "*/${profile}/*" -type f -executable \
        ! -path "*/deps/*" \
        2>/dev/null | head -1)

    if [[ -n "$bin" && -f "$bin" ]]; then
        echo "$bin"
        return 0
    fi

    return 1
}

# PGO Step 1: Build instrumented binary
pgo_build_instrumented() {
    print_phase "PGO Step 1: Building Instrumented Binary"
    print_step "Features: ${FEATURE_FLAGS:-default}"

    # Clear old profile data
    rm -rf "${PROJECT_ROOT}/target/pgo-profiles"

    # Build with PGO instrumentation
    # shellcheck disable=SC2086
    cargo pgo build -- $FEATURE_FLAGS --bin "$BINARY_NAME"

    print_success "Instrumented binary built"
}

# PGO Step 2: Train the binary with workloads
pgo_train() {
    print_phase "PGO Step 2: Training"

    local bin
    bin=$(find_pgo_binary "release") || {
        print_error "Could not find instrumented binary"
        exit 1
    }
    print_step "Using binary: $bin"

    # Set profraw output location
    export LLVM_PROFILE_FILE="${PROJECT_ROOT}/target/pgo-profiles/wasot_%m_%p.profraw"
    mkdir -p "${PROJECT_ROOT}/target/pgo-profiles"

    # Training workload 1: XML conversion
    print_step "Training: XML conversion..."
    if [[ -f "${PROJECT_ROOT}/examples/demo.wasot" ]]; then
        "$bin" xml-listener "${PROJECT_ROOT}/examples/demo.wasot" --pretty >/dev/null 2>&1 || true
    fi

    # Training workload 2: Selector queries
    print_step "Training: Selector queries..."
    if [[ -f "${PROJECT_ROOT}/examples/selector-showcase.wasot" ]]; then
        "$bin" select "${PROJECT_ROOT}/examples/selector-showcase.wasot" -s "paragraph" >/dev/null 2>&1 || true
        "$bin" select "${PROJECT_ROOT}/examples/selector-showcase.wasot" -s "section > heading" >/dev/null 2>&1 || true
    fi

    # Training workload 3: SQLite conversion (if examples exist)
    print_step "Training: SQLite conversion..."
    if [[ -f "${PROJECT_ROOT}/examples/demo.wasot" ]]; then
        local tmp_db="/tmp/wasot-pgo-train-$$.db"
        "$bin" to-sqlite "${PROJECT_ROOT}/examples/demo.wasot" -d "$tmp_db" 2>/dev/null || true
        rm -f "$tmp_db"
    fi

    # Training workload 4: JSONL conversion
    print_step "Training: JSONL conversion..."
    if [[ -f "${PROJECT_ROOT}/examples/demo.wasot" ]]; then
        "$bin" to-jsonl "${PROJECT_ROOT}/examples/demo.wasot" >/dev/null 2>&1 || true
    fi

    # Training workload 5: Help and version (for CLI parsing paths)
    print_step "Training: CLI parsing..."
    "$bin" --help >/dev/null 2>&1 || true
    "$bin" --version >/dev/null 2>&1 || true

    unset LLVM_PROFILE_FILE

    # Verify profile data was generated
    local profile_count
    profile_count=$(find "${PROJECT_ROOT}/target/pgo-profiles" -name "*.profraw" 2>/dev/null | wc -l)
    if [[ "$profile_count" -eq 0 ]]; then
        print_error "No profile data generated during training"
        exit 1
    fi
    print_success "Training complete ($profile_count profile files)"
}

# PGO Step 3: Build optimized binary
pgo_build_optimized() {
    print_phase "PGO Step 3: Building Optimized Binary"

    # shellcheck disable=SC2086
    cargo pgo optimize -- $FEATURE_FLAGS --bin "$BINARY_NAME"

    print_success "PGO-optimized binary built"
}

# Finalize: Copy to release dir and strip
finalize() {
    print_phase "Finalizing"

    mkdir -p "$RELEASE_DIR"

    local src_bin
    if [[ "$SIMPLE_BUILD" == true ]]; then
        src_bin="./target/release/$BINARY_NAME"
    else
        src_bin=$(find_pgo_binary "release") || {
            print_error "Could not find final binary"
            exit 1
        }
    fi

    if [[ ! -f "$src_bin" ]]; then
        print_error "Binary not found at $src_bin"
        exit 1
    fi

    local final_bin="$RELEASE_DIR/$BINARY_NAME"
    cp "$src_bin" "$final_bin"

    local size_before
    size_before=$(du -h "$final_bin" | cut -f1)
    strip --strip-all "$final_bin"
    local size_after
    size_after=$(du -h "$final_bin" | cut -f1)
    print_success "Stripped: $size_before -> $size_after"

    # Verify the binary works
    if "$final_bin" --version &>/dev/null; then
        print_success "Binary works"
        "$final_bin" --version

        echo 'sudo cp target/release/wasot /usr/local/bin/wasot'
    else
        print_error "Binary execution failed!"
        exit 1
    fi

    echo -e "\n${GREEN}Build complete: ${CYAN}$final_bin${NC}"
    echo -e "Features: ${FEATURE_FLAGS:-default}"
}

main() {
    cd "$PROJECT_ROOT"

    parse_args "$@"
    check_prerequisites

    [[ "$DO_CLEAN" == true ]] && clean_artifacts

    if [[ "$SIMPLE_BUILD" == true ]]; then
        build_simple
    else
        pgo_build_instrumented
        pgo_train
        pgo_build_optimized
    fi

    finalize
}

trap 'print_error "Build interrupted"; exit 1' INT TERM
main "$@"
