#!/bin/bash

# ParticleOS apt-layer Tool
# Enhanced version with container support and multiple package managers
# Inspired by Vanilla OS Apx approach
# Usage: apt-layer <base-branch> <new-branch> <packages...>
#        apt-layer --list | --info <branch> | --rollback <branch> | --help

set -euo pipefail

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

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

log_debug() {
    echo -e "${YELLOW}[DEBUG]${NC} $1"
}

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

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

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

log_layer() {
    echo -e "${PURPLE}[LAYER]${NC} $1"
}

log_container() {
    echo -e "${CYAN}[CONTAINER]${NC} $1"
}

# Configuration
WORKSPACE="/workspace"
OSTREE_REPO="$WORKSPACE/cache/ostree-repo"
BUILD_DIR="$WORKSPACE/cache/build"
CONTAINER_RUNTIME="podman"  # or "docker"

# Show usage
show_usage() {
    cat << EOF
ParticleOS apt-layer Tool - Enhanced with Container Support
Like rpm-ostree + Vanilla OS Apx for Ubuntu/Debian

Usage:
  apt-layer <base-branch> <new-branch> [packages...]
    # Add a new layer to an existing OSTree image (build or user)

  apt-layer --container <base-branch> <new-branch> [packages...]
    # Create layer using container isolation (like Apx)

  apt-layer --oci-export <branch> <image-name>
    # Export OSTree branch as OCI image

  apt-layer --list
    # List all available branches/layers

  apt-layer --info <branch>
    # Show information about a specific branch/layer

  apt-layer --rollback <branch>
    # Rollback a branch/layer to the previous commit

  apt-layer --help
    # Show this help message

Examples:
  # Traditional layer creation
  apt-layer particleos/base/trixie particleos/gaming/trixie steam wine

  # Container-based layer creation (Apx-style)
  apt-layer --container particleos/base/trixie particleos/gaming/trixie steam wine

  # Export as OCI image
  apt-layer --oci-export particleos/gaming/trixie particleos/gaming:latest

  # User custom layer with container isolation
  apt-layer --container particleos/desktop/trixie particleos/mycustom/trixie my-favorite-app

Description:
  apt-layer lets you add, inspect, and rollback layers on OSTree-based systems using apt packages.
  It can be used by system builders and end users, similar to rpm-ostree for Fedora Silverblue/Bazzite.
  Enhanced with container support inspired by Vanilla OS Apx.

EOF
}

# Check dependencies
check_dependencies() {
    log_info "Checking dependencies..."
    
    local missing_deps=()
    
    for dep in ostree chroot apt-get; do
        if ! command -v "$dep" >/dev/null 2>&1; then
            missing_deps+=("$dep")
        fi
    done
    
    # Check for container runtime
    if [[ "${1:-}" == "--container" ]] && ! command -v "$CONTAINER_RUNTIME" >/dev/null 2>&1; then
        missing_deps+=("$CONTAINER_RUNTIME")
    fi
    
    if [ ${#missing_deps[@]} -ne 0 ]; then
        log_error "Missing dependencies: ${missing_deps[*]}"
        exit 1
    fi
    
    log_success "All dependencies found"
}

# Create layer using container isolation (Apx-style)
create_container_layer() {
    local base_branch="$1"
    local new_branch="$2"
    shift 2
    local packages=("$@")
    
    log_container "Creating container-based layer: $new_branch"
    log_info "Base branch: $base_branch"
    log_info "Packages to install: ${packages[*]}"
    
    # Check if base branch exists
    if ! ostree --repo="$OSTREE_REPO" refs | grep -q "^$base_branch$"; then
        log_error "Base branch '$base_branch' not found"
        log_info "Available branches:"
        ostree --repo="$OSTREE_REPO" refs
        exit 1
    fi
    
    # Create workspace for container layer
    local layer_dir="$BUILD_DIR/container-layer-$(basename "$new_branch")"
    log_debug "Creating container layer workspace: $layer_dir"
    
    # Clean up any existing layer directory
    rm -rf "$layer_dir" 2>/dev/null || true
    mkdir -p "$layer_dir"
    
    # Checkout base branch to layer directory
    log_info "Checking out base branch..."
    ostree --repo="$OSTREE_REPO" checkout "$base_branch" "$layer_dir"
    
    # Create containerfile for the layer
    local containerfile="$layer_dir/Containerfile.layer"
    cat > "$containerfile" << EOF
FROM scratch
COPY . /
RUN apt-get update && apt-get install -y ${packages[*]} && apt-get clean && apt-get autoremove -y
EOF
    
    # Build container image
    log_info "Building container image..."
    local image_name="particleos-layer-$(basename "$new_branch"):latest"
    
    if "$CONTAINER_RUNTIME" build -f "$containerfile" -t "$image_name" "$layer_dir"; then
        log_success "Container image built: $image_name"
    else
        log_error "Failed to build container image"
        exit 1
    fi
    
    # Extract the modified filesystem
    log_info "Extracting modified filesystem..."
    local container_id
    container_id=$("$CONTAINER_RUNTIME" create "$image_name")
    
    # Copy files from container
    "$CONTAINER_RUNTIME" cp "$container_id:/" "$layer_dir/"
    "$CONTAINER_RUNTIME" rm "$container_id"
    
    # Clean up device files that can't be in OSTree
    log_debug "Cleaning up device files..."
    rm -rf "$layer_dir"/{dev,proc,sys}/* 2>/dev/null || true
    
    # Create commit message
    local commit_message="Add container layer: $new_branch"
    if [ ${#packages[@]} -gt 0 ]; then
        commit_message+=" with packages: ${packages[*]}"
    fi
    
    # Create OSTree commit
    log_info "Creating OSTree commit..."
    local commit_hash
    if commit_hash=$(ostree --repo="$OSTREE_REPO" commit --branch="$new_branch" --tree=dir="$layer_dir" --subject="$commit_message" 2>&1); then
        log_success "Container layer created successfully: $new_branch"
        echo "Commit: $commit_hash"
        echo "Branch: $new_branch"
        echo "Container image: $image_name"
    else
        log_error "Failed to create commit: $commit_hash"
        exit 1
    fi
    
    # Clean up layer directory
    rm -rf "$layer_dir"
    
    log_success "Container layer build completed successfully!"
}

# Export OSTree branch as OCI image
export_oci_image() {
    local branch="$1"
    local image_name="$2"
    
    log_container "Exporting OSTree branch as OCI image: $branch -> $image_name"
    
    if ! ostree --repo="$OSTREE_REPO" refs | grep -q "^$branch$"; then
        log_error "Branch '$branch' not found"
        exit 1
    fi
    
    # Create temporary directory for export
    local export_dir="$BUILD_DIR/oci-export-$(basename "$branch")"
    rm -rf "$export_dir" 2>/dev/null || true
    mkdir -p "$export_dir"
    
    # Checkout branch
    log_info "Checking out branch for OCI export..."
    ostree --repo="$OSTREE_REPO" checkout "$branch" "$export_dir"
    
    # Create containerfile for OCI image
    local containerfile="$export_dir/Containerfile"
    cat > "$containerfile" << EOF
FROM scratch
COPY . /
CMD ["/bin/bash"]
EOF
    
    # Build OCI image
    log_info "Building OCI image..."
    if "$CONTAINER_RUNTIME" build -f "$containerfile" -t "$image_name" "$export_dir"; then
        log_success "OCI image exported successfully: $image_name"
        echo "Image: $image_name"
        echo "Branch: $branch"
    else
        log_error "Failed to export OCI image"
        exit 1
    fi
    
    # Clean up
    rm -rf "$export_dir"
}

# Main execution
main() {
    # Check dependencies
    check_dependencies
    
    # Parse command line arguments
    case "${1:-}" in
        --help|-h)
            show_usage
            exit 0
            ;;
        --list)
            list_branches
            exit 0
            ;;
        --info)
            if [ -z "${2:-}" ]; then
                log_error "Branch name required for --info"
                show_usage
                exit 1
            fi
            show_branch_info "$2"
            exit 0
            ;;
        --rollback)
            if [ -z "${2:-}" ]; then
                log_error "Branch name required for --rollback"
                show_usage
                exit 1
            fi
            rollback_branch "$2"
            exit 0
            ;;
        --container)
            # Container-based layer creation
            if [ $# -lt 4 ]; then
                log_error "Insufficient arguments for --container"
                show_usage
                exit 1
            fi
            
            local base_branch="$2"
            local new_branch="$3"
            shift 3
            local packages=("$@")
            
            create_container_layer "$base_branch" "$new_branch" "${packages[@]}"
            ;;
        --oci-export)
            # Export OSTree branch as OCI image
            if [ $# -lt 3 ]; then
                log_error "Insufficient arguments for --oci-export"
                show_usage
                exit 1
            fi
            
            local branch="$2"
            local image_name="$3"
            
            export_oci_image "$branch" "$image_name"
            ;;
        "")
            log_error "No arguments provided"
            show_usage
            exit 1
            ;;
        *)
            # Regular layer creation
            if [ $# -lt 2 ]; then
                log_error "Insufficient arguments"
                show_usage
                exit 1
            fi
            
            local base_branch="$1"
            local new_branch="$2"
            shift 2
            local packages=("$@")
            
            create_layer "$base_branch" "$new_branch" "${packages[@]}"
            ;;
    esac
}

# Run main function
main "$@"