#!/bin/bash

# Ubuntu uBlue Installation Script
# Installs and configures the complete Ubuntu uBlue system

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'
NC='\033[0m'

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

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

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

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

log_step() {
    echo -e "${PURPLE}[STEP]${NC} $1"
}

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INSTALL_DIR="/usr/local"
CONFIG_DIR="/usr/local/etc/ubuntu-ublue"
BIN_DIR="/usr/local/bin"
LOG_DIR="/var/log/ubuntu-ublue"
ROOT_DIR="/var/lib/ubuntu-ublue"

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        log_info "Please run: sudo $0"
        exit 1
    fi
}

# Check Ubuntu version
check_ubuntu_version() {
    log_step "Checking Ubuntu version..."
    
    if ! command -v lsb_release >/dev/null 2>&1; then
        log_error "lsb_release not found. This script requires Ubuntu."
        exit 1
    fi
    
    local version
    version=$(lsb_release -rs)
    local codename
    codename=$(lsb_release -cs)
    
    log_info "Detected Ubuntu $version ($codename)"
    
    # Check if version is supported (20.04 or later)
    if [[ "$version" < "20.04" ]]; then
        log_error "Ubuntu $version is not supported. Please use Ubuntu 20.04 or later."
        exit 1
    fi
    
    log_success "Ubuntu version check passed"
}

# Install system dependencies
install_dependencies() {
    log_step "Installing system dependencies..."
    
    # Update package lists
    log_info "Updating package lists..."
    apt-get update
    
    # Install required packages
    local packages=(
        "squashfs-tools"
        "overlayroot"
        "chroot"
        "rsync"
        "tar"
        "gzip"
        "xz-utils"
        "zstd"
        "curl"
        "wget"
        "gnupg"
        "lsb-release"
        "util-linux"
        "mount"
        "umount"
        "findmnt"
        "lsof"
    )
    
    log_info "Installing packages: ${packages[*]}"
    apt-get install -y "${packages[@]}"
    
    # Install optional container runtime (podman preferred, fallback to docker)
    if command -v podman >/dev/null 2>&1; then
        log_info "Podman already installed"
    elif command -v docker >/dev/null 2>&1; then
        log_info "Docker found, will use as container runtime"
    else
        log_info "Installing Podman as container runtime..."
        apt-get install -y podman
    fi
    
    log_success "System dependencies installed"
}

# Create directory structure
create_directories() {
    log_step "Creating directory structure..."
    
    local directories=(
        "$ROOT_DIR"
        "$CONFIG_DIR"
        "$LOG_DIR"
        "/var/cache/ubuntu-ublue"
        "$ROOT_DIR/build"
        "$ROOT_DIR/workspace"
        "$ROOT_DIR/temp"
        "$ROOT_DIR/live-overlay"
        "$ROOT_DIR/live-overlay/upper"
        "$ROOT_DIR/live-overlay/work"
        "$ROOT_DIR/backups"
        "/var/lib/composefs-alternative"
        "/var/lib/composefs-alternative/images"
        "/var/lib/composefs-alternative/layers"
        "/var/lib/composefs-alternative/mounts"
    )
    
    for dir in "${directories[@]}"; do
        log_info "Creating directory: $dir"
        mkdir -p "$dir"
    done
    
    # Set proper permissions
    chmod 755 "$ROOT_DIR" "$CONFIG_DIR" "$LOG_DIR" "/var/cache/ubuntu-ublue"
    chmod 700 "$ROOT_DIR/backups"
    
    log_success "Directory structure created"
}

# Install scripts
install_scripts() {
    log_step "Installing Ubuntu uBlue scripts..."
    
    # Copy scripts to /usr/local/bin
    local scripts=(
        "apt-layer.sh"
        "bootloader-integration.sh"
        "oci-integration.sh"
        "test-integration.sh"
        "ublue-logrotate.sh"
    )
    
    for script in "${scripts[@]}"; do
        local source_path="$SCRIPT_DIR/$script"
        local target_path="$BIN_DIR/$script"
        
        if [[ -f "$source_path" ]]; then
            log_info "Installing $script..."
            cp "$source_path" "$target_path"
            chmod +x "$target_path"
            log_success "Installed $script"
        else
            log_warning "Script not found: $source_path"
        fi
    done
    
    # Install configuration file
    if [[ -f "$SCRIPT_DIR/ublue-config.sh" ]]; then
        log_info "Installing configuration file..."
        cp "$SCRIPT_DIR/ublue-config.sh" "$CONFIG_DIR/ublue-config.sh"
        chmod 644 "$CONFIG_DIR/ublue-config.sh"
        log_success "Configuration file installed"
    else
        log_warning "Configuration file not found: $SCRIPT_DIR/ublue-config.sh"
    fi
    
    log_success "Scripts installed"
}

# Create systemd services
create_systemd_services() {
    log_step "Creating systemd services..."
    
    # Create service directory
    mkdir -p /etc/systemd/system
    
    # Ubuntu uBlue transaction cleanup service
    cat > /etc/systemd/system/ubuntu-ublue-cleanup.service << 'EOF'
[Unit]
Description=Ubuntu uBlue Transaction Cleanup
After=multi-user.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/apt-layer.sh --cleanup-incomplete
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
    
    # Ubuntu uBlue log rotation service
    cat > /etc/systemd/system/ubuntu-ublue-logrotate.service << 'EOF'
[Unit]
Description=Ubuntu uBlue Log Rotation
After=multi-user.target

[Service]
Type=oneshot
ExecStart=/usr/sbin/logrotate /etc/logrotate.d/ubuntu-ublue
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
    
    # Enable services
    systemctl daemon-reload
    systemctl enable ubuntu-ublue-cleanup.service
    systemctl enable ubuntu-ublue-logrotate.service
    
    log_success "Systemd services created"
}

# Create logrotate configuration
create_logrotate_config() {
    log_step "Creating logrotate configuration..."
    
    cat > /etc/logrotate.d/ubuntu-ublue << 'EOF'
/var/log/ubuntu-ublue/*.log {
    daily
    missingok
    rotate 7
    compress
    delaycompress
    notifempty
    create 644 root root
    postrotate
        systemctl reload ubuntu-ublue-logrotate.service
    endscript
}
EOF
    
    log_success "Logrotate configuration created"
}

# Create man pages
create_man_pages() {
    log_step "Creating man pages..."
    
    mkdir -p /usr/local/share/man/man1
    
    # apt-layer man page
    cat > /usr/local/share/man/man1/apt-layer.1 << 'EOF'
.TH APT-LAYER 1 "December 2024" "Ubuntu uBlue" "User Commands"

.SH NAME
apt-layer \- Ubuntu uBlue layer management tool

.SH SYNOPSIS
.B apt-layer
[\fIOPTIONS\fR] \fIBASE_IMAGE\fR \fINEW_IMAGE\fR [\fIPACKAGES\fR...]

.SH DESCRIPTION
apt-layer is a tool for managing layers on Ubuntu uBlue systems using ComposeFS.
It provides functionality similar to rpm-ostree for Fedora Silverblue/Kinoite.

.SH OPTIONS
.TP
.B \-\-help
Show help message

.TP
.B \-\-list
List all available ComposeFS images/layers

.TP
.B \-\-info \fIIMAGE\fR
Show information about a specific image

.TP
.B \-\-live-install \fIPACKAGES\fR...
Install packages on live system with overlayfs

.TP
.B \-\-live-commit [\fIMESSAGE\fR]
Commit current live overlay changes as new layer

.TP
.B \-\-live-rollback
Rollback live overlay changes

.TP
.B \-\-container
Use container isolation for package installation

.TP
.B \-\-oci-export \fIIMAGE\fR \fIOCI_NAME\fR
Export ComposeFS image as OCI image

.TP
.B \-\-oci-import \fIOCI_NAME\fR \fITARGET_IMAGE\fR
Import OCI image as ComposeFS image

.SH EXAMPLES
.TP
Create a new layer:
.B apt-layer ubuntu-ublue/base/24.04 ubuntu-ublue/gaming/24.04 steam wine

.TP
Install packages on live system:
.B apt-layer \-\-live-install steam wine

.TP
Commit live changes:
.B apt-layer \-\-live-commit "Add gaming packages"

.SH FILES
.TP
.I /var/lib/ubuntu-ublue/
Ubuntu uBlue root directory

.TP
.I /usr/local/etc/ublue-config.sh
Configuration file

.TP
.I /var/log/ubuntu-ublue/
Log files

.SH SEE ALSO
.BR bootloader-integration (1),
.BR oci-integration (1)

.SH AUTHOR
Ubuntu uBlue Project
EOF
    
    # Update man database
    mandb -q
    
    log_success "Man pages created"
}

# Create completion scripts
create_completion_scripts() {
    log_step "Creating completion scripts..."
    
    mkdir -p /usr/local/share/bash-completion/completions
    
    # apt-layer completion
    cat > /usr/local/share/bash-completion/completions/apt-layer << 'EOF'
_apt_layer() {
    local cur prev opts
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"
    
    opts="--help --list --info --live-install --live-commit --live-rollback --container --oci-export --oci-import"
    
    case "${prev}" in
        --info|--rollback)
            # Complete with available images
            local images
            images=$(apt-layer --list 2>/dev/null | grep -v "Listing" | awk '{print $1}' | tr '\n' ' ')
            COMPREPLY=( $(compgen -W "${images}" -- "${cur}") )
            return 0
            ;;
        --live-install)
            # Complete with common packages
            local packages="steam wine firefox chromium-browser vlc"
            COMPREPLY=( $(compgen -W "${packages}" -- "${cur}") )
            return 0
            ;;
    esac
    
    if [[ ${cur} == * ]] ; then
        COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
        return 0
    fi
}

complete -F _apt_layer apt-layer
EOF
    
    log_success "Completion scripts created"
}

# Configure environment
configure_environment() {
    log_step "Configuring environment..."
    
    # Add to PATH if not already present
    if ! grep -q "/usr/local/bin" /etc/environment; then
        echo 'PATH="/usr/local/bin:$PATH"' >> /etc/environment
    fi
    
    # Create profile script
    cat > /etc/profile.d/ubuntu-ublue.sh << 'EOF'
# Ubuntu uBlue environment configuration
export UBLUE_ROOT="/var/lib/ubuntu-ublue"
export UBLUE_CONFIG_DIR="/usr/local/etc/ubuntu-ublue"

# Source configuration if available
if [[ -f "$UBLUE_CONFIG_DIR/ublue-config.sh" ]]; then
    source "$UBLUE_CONFIG_DIR/ublue-config.sh"
fi
EOF
    
    chmod +x /etc/profile.d/ubuntu-ublue.sh
    
    log_success "Environment configured"
}

# Run post-installation tests
run_tests() {
    log_step "Running post-installation tests..."
    
    if [[ -f "$BIN_DIR/test-integration.sh" ]]; then
        log_info "Running integration tests..."
        if "$BIN_DIR/test-integration.sh"; then
            log_success "All tests passed"
        else
            log_warning "Some tests failed - check logs for details"
        fi
    else
        log_warning "Test script not found - skipping tests"
    fi
}

# Show installation summary
show_summary() {
    log_step "Installation Summary"
    echo
    echo "Ubuntu uBlue has been successfully installed!"
    echo
    echo "Installed Components:"
    echo "  • apt-layer.sh - Main layer management tool"
    echo "  • bootloader-integration.sh - Boot entry management"
    echo "  • oci-integration.sh - OCI container integration"
    echo "  • test-integration.sh - Integration testing"
    echo "  • ublue-config.sh - Unified configuration"
    echo
    echo "Directories:"
    echo "  • Root: $ROOT_DIR"
    echo "  • Config: $CONFIG_DIR"
    echo "  • Logs: $LOG_DIR"
    echo "  • Scripts: $BIN_DIR"
    echo
    echo "Next Steps:"
    echo "  1. Create a base image: apt-layer --oci-import ubuntu:24.04 ubuntu-ublue/base/24.04"
    echo "  2. Create a layer: apt-layer ubuntu-ublue/base/24.04 ubuntu-ublue/desktop/24.04 gnome-shell"
    echo "  3. Install packages live: apt-layer --live-install steam"
    echo "  4. View help: apt-layer --help"
    echo
    echo "Documentation:"
    echo "  • Man pages: man apt-layer"
    echo "  • Logs: tail -f $LOG_DIR/ublue.log"
    echo "  • Tests: $BIN_DIR/test-integration.sh"
    echo
    log_success "Installation completed successfully!"
}

# Main installation function
main() {
    echo "Ubuntu uBlue Installation Script"
    echo "================================"
    echo
    
    # Check prerequisites
    check_root
    check_ubuntu_version
    
    # Install components
    install_dependencies
    create_directories
    install_scripts
    create_systemd_services
    create_logrotate_config
    create_man_pages
    create_completion_scripts
    configure_environment
    
    # Run tests
    run_tests
    
    # Show summary
    show_summary
}

# Handle command line arguments
case "${1:-}" in
    "help"|"-h"|"--help")
        cat << EOF
Ubuntu uBlue Installation Script

Usage: $0 [options]

Options:
  help, -h, --help    Show this help message

Examples:
  sudo $0              Install Ubuntu uBlue system
  sudo $0 help         Show this help message

This script installs and configures the complete Ubuntu uBlue system,
including all scripts, services, and configuration files.
EOF
        exit 0
        ;;
    "")
        main
        ;;
    *)
        log_error "Unknown option: $1"
        echo "Use '$0 help' for usage information"
        exit 1
        ;;
esac