#!/bin/bash

# Shader compilation script for Vulkan SPIR-V shaders
# Requires glslc (from Vulkan SDK) or glslangValidator

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUTPUT_DIR="$SCRIPT_DIR/compiled"

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

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check for shader compiler
COMPILER=""
if command -v glslc >/dev/null 2>&1; then
    COMPILER="glslc"
    log_info "Using glslc compiler"
elif command -v glslangValidator >/dev/null 2>&1; then
    COMPILER="glslangValidator"
    log_info "Using glslangValidator compiler"
else
    log_error "No shader compiler found!"
    log_error "Please install Vulkan SDK or make sure glslc/glslangValidator is in PATH"
    exit 1
fi

# Create output directory
mkdir -p "$OUTPUT_DIR"

compile_shader() {
    local input_file="$1"
    local output_file="$2"
    
    log_info "Compiling $input_file -> $output_file"
    
    if [ "$COMPILER" = "glslc" ]; then
        glslc "$input_file" -o "$output_file"
    else
        glslangValidator -V "$input_file" -o "$output_file"
    fi
    
    if [ $? -eq 0 ]; then
        log_success "Successfully compiled $input_file"
        
        # Show file size
        local size=$(stat -f%z "$output_file" 2>/dev/null || stat -c%s "$output_file" 2>/dev/null)
        log_info "Output size: $size bytes"
    else
        log_error "Failed to compile $input_file"
        return 1
    fi
}

# Compile all shaders
log_info "Starting shader compilation..."

# Triangle shaders
if [ -f "$SCRIPT_DIR/triangle.vert" ]; then
    compile_shader "$SCRIPT_DIR/triangle.vert" "$OUTPUT_DIR/triangle_vert.spv"
else
    log_warning "triangle.vert not found"
fi

if [ -f "$SCRIPT_DIR/triangle.frag" ]; then
    compile_shader "$SCRIPT_DIR/triangle.frag" "$OUTPUT_DIR/triangle_frag.spv"
else
    log_warning "triangle.frag not found"
fi

# Compute shader (if exists)
if [ -f "$SCRIPT_DIR/compute.comp" ]; then
    compile_shader "$SCRIPT_DIR/compute.comp" "$OUTPUT_DIR/compute.spv"
fi

# Create header file with embedded SPIR-V for Go
create_go_header() {
    local spv_file="$1"
    local var_name="$2"
    local go_file="$OUTPUT_DIR/shaders.go"
    
    if [ ! -f "$spv_file" ]; then
        return
    fi
    
    log_info "Creating Go header for $spv_file"
    
    if [ ! -f "$go_file" ]; then
        cat > "$go_file" << 'EOF'
// Code generated by compile.sh. DO NOT EDIT.

package shaders

// Compiled SPIR-V shader bytecode
var (
EOF
    fi
    
    echo -n "	$var_name = []byte{" >> "$go_file"
    
    # Convert binary to Go byte array
    hexdump -v -e '16/1 "%d," "\n"' "$spv_file" | sed 's/,$//' | while read line; do
        echo "" >> "$go_file"
        echo -n "		$line," >> "$go_file"
    done
    
    echo "" >> "$go_file"
    echo "	}" >> "$go_file"
}

# Create Go header file
if [ -d "$OUTPUT_DIR" ] && [ -n "$(ls -A $OUTPUT_DIR/*.spv 2>/dev/null)" ]; then
    log_info "Creating Go header with embedded shaders..."
    
    rm -f "$OUTPUT_DIR/shaders.go"
    
    # Add package header
    cat > "$OUTPUT_DIR/shaders.go" << 'EOF'
// Code generated by compile.sh. DO NOT EDIT.

package shaders

// Compiled SPIR-V shader bytecode
var (
EOF
    
    # Add each shader
    for spv_file in "$OUTPUT_DIR"/*.spv; do
        if [ -f "$spv_file" ]; then
            basename=$(basename "$spv_file" .spv)
            var_name=$(echo "$basename" | sed 's/[^a-zA-Z0-9_]/_/g')
            var_name="${var_name^}" # Capitalize first letter
            
            echo "	${var_name}SPIRV = []byte{" >> "$OUTPUT_DIR/shaders.go"
            
            # Convert binary to Go byte array
            hexdump -v -e '16/1 "%d," "\n"' "$spv_file" | sed 's/,$//' | while read line; do
                echo "		$line," >> "$OUTPUT_DIR/shaders.go"
            done
            
            echo "	}" >> "$OUTPUT_DIR/shaders.go"
            echo "" >> "$OUTPUT_DIR/shaders.go"
        fi
    done
    
    echo ")" >> "$OUTPUT_DIR/shaders.go"
    
    log_success "Created Go header at $OUTPUT_DIR/shaders.go"
fi

log_success "Shader compilation complete!"
log_info "Compiled shaders are in: $OUTPUT_DIR"

# List compiled files
if [ -d "$OUTPUT_DIR" ]; then
    log_info "Generated files:"
    ls -la "$OUTPUT_DIR"
fi