#!/bin/bash

# Golang AVX512 FFT Build and Test Script
# This script uses a Go container to build and test the FFT implementation

set -e  # Exit on any error

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

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

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

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

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

# Function to check if Docker is available
check_docker() {
    if ! command -v docker &> /dev/null; then
        print_error "Docker is not installed or not in PATH"
        print_error "Please install Docker and try again"
        exit 1
    fi

    if ! docker info &> /dev/null; then
        print_error "Docker daemon is not running"
        print_error "Please start Docker and try again"
        exit 1
    fi

    print_success "Docker is available and running"
}

# Function to check if required files exist
check_files() {
    local required_files=(
        "go.mod"
        "fft.go"
        "fft_avx512_working.s"
        "fft_test.go"
        "README.md"
    )

    local missing_files=()

    for file in "${required_files[@]}"; do
        if [[ ! -f "$file" ]]; then
            missing_files+=("$file")
        fi
    done

    if [[ ${#missing_files[@]} -gt 0 ]]; then
        print_error "Missing required files:"
        for file in "${missing_files[@]}"; do
            echo "  - $file"
        done
        exit 1
    fi

    print_success "All required files are present"
}

# Function to create Dockerfile
create_dockerfile() {
    print_status "Creating Dockerfile for Go environment"

    cat > Dockerfile << 'EOF'
FROM golang:1.21-bullseye

# Install required packages
RUN apt-get update && apt-get install -y \
    gcc \
    g++ \
    make \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Copy go mod files first for better caching
COPY go.mod go.sum* ./

# Download dependencies
RUN go mod download

# Copy source code
COPY . .

# Build the application
RUN go build -o fft .

# Run tests
RUN go test -v .

# Run benchmarks
RUN go test -bench=. -benchmem .

# Show binary info
RUN ls -la fft
RUN file fft

# Show Go version and environment
RUN go version
RUN go env GOOS GOARCH GOAMD64

# Check if AVX512 is supported (this will show in container)
RUN echo "Container CPU info:" && cat /proc/cpuinfo | grep -i avx512 | head -5 || echo "No AVX512 info available in container"

# Keep container running for interactive use
CMD ["/bin/bash"]
EOF

    print_success "Dockerfile created"
}

# Function to build and run container
build_and_run_container() {
    print_status "Building Go container image"

    # Build the image
    docker build -t golang-fft:latest .

    if [[ $? -eq 0 ]]; then
        print_success "Container image built successfully"
    else
        print_error "Failed to build container image"
        exit 1
    fi

    print_status "Running container for interactive testing"

    # Run the container interactively
    docker run -it --rm \
        --name golang-fft-test \
        golang-fft:latest
}

# Function to run quick test without interactive mode
run_quick_test() {
    print_status "Running quick build and test in container"

    # Run container, execute tests, and exit
    docker run --rm \
        --name golang-fft-quick \
        golang-fft:latest \
        bash -c "
            echo '=== Building application ==='
            go build -o fft .

            echo '=== Running tests ==='
            go test -v .

            echo '=== Running benchmarks ==='
            go test -bench=. -benchmem .

            echo '=== Application info ==='
            ls -la fft
            file fft

            echo '=== Go environment ==='
            go version
            go env GOOS GOARCH GOAMD64

            echo '=== CPU info ==='
            cat /proc/cpuinfo | grep -i avx512 | head -5 || echo 'No AVX512 info available'
        "
}

# Function to clean up
cleanup() {
    print_status "Cleaning up Docker resources"

    # Stop and remove containers
    docker stop golang-fft-test golang-fft-quick 2>/dev/null || true
    docker rm golang-fft-test golang-fft-quick 2>/dev/null || true

    # Remove image
    docker rmi golang-fft:latest 2>/dev/null || true

    # Remove Dockerfile
    rm -f Dockerfile

    print_success "Cleanup completed"
}

# Function to show help
show_help() {
    echo "Golang AVX512 FFT Build and Test Script"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -h, --help          Show this help message"
    echo "  -q, --quick         Run quick test without interactive mode"
    echo "  -c, --cleanup       Clean up Docker resources and exit"
    echo "  -i, --interactive   Run interactive container (default)"
    echo ""
    echo "Examples:"
    echo "  $0                  # Run interactive container"
    echo "  $0 --quick          # Run quick test and exit"
    echo "  $0 --cleanup        # Clean up and exit"
    echo ""
}

# Main script logic
main() {
    local mode="interactive"

    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -q|--quick)
                mode="quick"
                shift
                ;;
            -c|--cleanup)
                cleanup
                exit 0
                ;;
            -i|--interactive)
                mode="interactive"
                shift
                ;;
            *)
                print_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done

    print_status "Starting Golang AVX512 FFT build and test process"

    # Check prerequisites
    check_docker
    check_files

    # Create Dockerfile
    create_dockerfile

    # Handle different modes
    case $mode in
        "quick")
            run_quick_test
            ;;
        "interactive")
            build_and_run_container
            ;;
    esac

    print_success "Process completed successfully"
}

# Trap to ensure cleanup on script exit
trap cleanup EXIT

# Run main function with all arguments
main "$@"