#!/bin/bash
# ISO Download Script with Checksum Verification
# Downloads distribution ISOs and verifies integrity

set -euo pipefail

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

# Configuration
ISO_DIR="${ISO_DIR:-/var/lib/vz/template/iso}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$SCRIPT_DIR")"

# Logging
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

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

# Parse YAML (basic parser for distributions.yml)
parse_yaml() {
    local yaml_file=$1
    local distro=$2
    local version=$3
    local field=$4

    # This is a simplified parser - for production use yq or python
    grep -A 20 "\"$version\":" "$yaml_file" | grep "$field:" | head -1 | sed 's/.*: "\(.*\)"/\1/' | tr -d '"'
}

# Download file with progress
download_file() {
    local url=$1
    local output=$2

    log_info "Downloading from: $url"

    if command -v wget &> /dev/null; then
        wget -c -O "$output" "$url" || return 1
    elif command -v curl &> /dev/null; then
        curl -L -C - -o "$output" "$url" || return 1
    else
        log_error "Neither wget nor curl found. Please install one."
        return 1
    fi
}

# Extract checksum from checksum file
extract_checksum() {
    local checksum_file=$1
    local iso_filename=$2

    if [[ -f "$checksum_file" ]]; then
        # Try different checksum file formats
        # Format 1: SHA256 (filename) = checksum
        local checksum=$(grep "$iso_filename" "$checksum_file" | grep -oP '= \K[a-f0-9]{64}' | head -1)
        if [[ -n "$checksum" ]]; then
            echo "$checksum"
            return
        fi
        # Format 2: checksum  filename
        checksum=$(grep "$iso_filename" "$checksum_file" | awk '{print $1}' | grep -E '^[a-f0-9]{64}$' | head -1)
        if [[ -n "$checksum" ]]; then
            echo "$checksum"
            return
        fi
    fi
}

# Verify checksum
verify_checksum() {
    local file=$1
    local expected_checksum=$2
    local checksum_type=${3:-sha256}

    log_info "Verifying ${checksum_type} checksum..."

    local actual_checksum
    case "$checksum_type" in
        sha256)
            actual_checksum=$(sha256sum "$file" | awk '{print $1}')
            ;;
        sha512)
            actual_checksum=$(sha512sum "$file" | awk '{print $1}')
            ;;
        md5)
            actual_checksum=$(md5sum "$file" | awk '{print $1}')
            ;;
        *)
            log_error "Unknown checksum type: $checksum_type"
            return 1
            ;;
    esac

    if [[ "$actual_checksum" == "$expected_checksum" ]]; then
        log_info "Checksum verification PASSED ✓"
        return 0
    else
        log_error "Checksum verification FAILED ✗"
        log_error "Expected: $expected_checksum"
        log_error "Got:      $actual_checksum"
        return 1
    fi
}

# Main download function
download_iso() {
    local distro=$1
    local version=$2

    local config_file="$ROOT_DIR/config/distributions.yml"

    if [[ ! -f "$config_file" ]]; then
        log_error "Configuration file not found: $config_file"
        return 1
    fi

    # Parse configuration
    local iso_url=$(parse_yaml "$config_file" "$distro" "$version" "iso_url")
    local checksum_url=$(parse_yaml "$config_file" "$distro" "$version" "iso_checksum_url")
    local checksum_type=$(parse_yaml "$config_file" "$distro" "$version" "iso_checksum_type")

    if [[ -z "$iso_url" ]]; then
        log_error "Could not find ISO URL for $distro $version"
        return 1
    fi

    local iso_filename=$(basename "$iso_url")
    local iso_path="$ISO_DIR/$iso_filename"
    local checksum_path="$ISO_DIR/${iso_filename}.checksum"

    # Create ISO directory if it doesn't exist
    if [[ ! -d "$ISO_DIR" ]]; then
        log_info "Creating ISO directory: $ISO_DIR"
        sudo mkdir -p "$ISO_DIR"
    fi

    # Check if ISO already exists and is not empty
    if [[ -f "$iso_path" ]]; then
        local file_size=$(stat -c%s "$iso_path" 2>/dev/null || stat -f%z "$iso_path" 2>/dev/null)
        if [[ "$file_size" -gt 1000000 ]]; then
            # File exists and is larger than 1MB (likely valid) - just use it
            log_info "ISO already exists: $iso_path ($(numfmt --to=iec-i --suffix=B $file_size 2>/dev/null || echo "${file_size} bytes"))"
            return 0
        else
            # File exists but is empty or too small, remove it
            log_warn "ISO file exists but is empty or incomplete, removing..."
            rm -f "$iso_path"
        fi
    fi

    # Download checksum file first
    if [[ -n "$checksum_url" ]]; then
        log_info "Downloading checksum file..."
        download_file "$checksum_url" "$checksum_path" || {
            log_warn "Failed to download checksum file"
        }
    fi

    # Download ISO
    log_info "Downloading ISO: $iso_filename"
    download_file "$iso_url" "$iso_path" || {
        log_error "Failed to download ISO"
        return 1
    }

    # Verify checksum if available
    if [[ -f "$checksum_path" ]]; then
        local expected_checksum=$(extract_checksum "$checksum_path" "$iso_filename")

        if [[ -n "$expected_checksum" ]]; then
            verify_checksum "$iso_path" "$expected_checksum" "$checksum_type" || {
                log_error "ISO verification failed. Removing downloaded file."
                rm -f "$iso_path"
                return 1
            }
        else
            log_warn "Could not extract checksum from checksum file"
        fi
    else
        log_warn "No checksum file available - skipping verification"
    fi

    log_info "ISO downloaded successfully: $iso_path"

    # Set proper permissions
    sudo chmod 644 "$iso_path"

    return 0
}

# Usage
usage() {
    cat <<EOF
Usage: $0 DISTRO VERSION

Downloads and verifies distribution ISO files.

Arguments:
  DISTRO    Distribution name (ubuntu, debian, rocky, almalinux, fedora)
  VERSION   Version number (e.g., 22.04, 24.04, 12, 9, 40)

Environment Variables:
  ISO_DIR   Directory to store ISOs (default: /var/lib/vz/template/iso)

Examples:
  $0 ubuntu 24.04
  $0 debian 12
  $0 rocky 9

EOF
}

# Main
main() {
    if [[ $# -lt 2 ]]; then
        usage
        exit 1
    fi

    local distro=$1
    local version=$2

    download_iso "$distro" "$version"
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi
