#!/bin/bash
#
# EcoOS Installer
# Installs EcoOS from live USB to disk
#

set -e

# Configuration
TIMEOUT=10
HOSTNAME="ecoos"
USERNAME="ecouser"
SQUASHFS_PATH="/run/live/medium/live/filesystem.squashfs"

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

log() {
    echo -e "${GREEN}[EcoOS]${NC} $1" >&2
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1" >&2
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
    exit 1
}

# Get the device the live system is running from
get_live_device() {
    local live_dev=""
    # Find the device backing /run/live/medium
    if mountpoint -q /run/live/medium 2>/dev/null; then
        live_dev=$(findmnt -n -o SOURCE /run/live/medium | sed 's/[0-9]*$//')
    fi
    # Also check /cdrom for older casper
    if [ -z "$live_dev" ] && mountpoint -q /cdrom 2>/dev/null; then
        live_dev=$(findmnt -n -o SOURCE /cdrom | sed 's/[0-9]*$//')
    fi
    echo "$live_dev"
}

# List available disks (excluding live media and loop devices)
list_disks() {
    local live_dev=$(get_live_device)
    local disks=()

    for disk in /sys/block/sd* /sys/block/nvme* /sys/block/vd*; do
        [ -e "$disk" ] || continue
        local name=$(basename "$disk")
        local dev="/dev/$name"

        # Skip if this is the live media
        if [ "$dev" = "$live_dev" ]; then
            continue
        fi

        # Skip removable devices (USB sticks) - but allow if it's the only option
        local removable=$(cat "$disk/removable" 2>/dev/null || echo "0")

        # Get size in GB
        local size_bytes=$(cat "$disk/size" 2>/dev/null || echo "0")
        local size_gb=$((size_bytes * 512 / 1024 / 1024 / 1024))

        # Skip disks smaller than 10GB
        if [ "$size_gb" -lt 10 ]; then
            continue
        fi

        # Get model
        local model=$(cat "$disk/device/model" 2>/dev/null | tr -d '\n' || echo "Unknown")

        disks+=("$dev|$size_gb|$model|$removable")
    done

    printf '%s\n' "${disks[@]}"
}

# Select disk with timeout
# All UI output goes to stderr so stdout only returns the device path
select_disk() {
    local disks
    mapfile -t disks < <(list_disks)

    if [ ${#disks[@]} -eq 0 ]; then
        error "No suitable disks found for installation"
    fi

    echo "" >&2
    echo -e "${BLUE}╔════════════════════════════════════════════════════════════╗${NC}" >&2
    echo -e "${BLUE}║${NC}              ${GREEN}EcoOS Disk Installation${NC}                       ${BLUE}║${NC}" >&2
    echo -e "${BLUE}╚════════════════════════════════════════════════════════════╝${NC}" >&2
    echo "" >&2
    echo "Available disks:" >&2
    echo "" >&2

    local i=1
    local default_disk=""
    local default_idx=1
    local max_size=0

    for disk_info in "${disks[@]}"; do
        IFS='|' read -r dev size model removable <<< "$disk_info"

        local marker=""
        if [ "$size" -gt "$max_size" ]; then
            max_size=$size
            default_disk=$dev
            default_idx=$i
        fi

        printf "  ${YELLOW}%d)${NC} %-12s %4d GB  %s\n" "$i" "$dev" "$size" "$model" >&2
        ((i++))
    done

    echo "" >&2
    echo -e "Default: ${GREEN}$default_disk${NC} (largest disk)" >&2
    echo "" >&2
    echo -e "${YELLOW}WARNING: Selected disk will be COMPLETELY ERASED!${NC}" >&2
    echo "" >&2

    # Countdown with input check
    local selected=""
    local remaining=$TIMEOUT

    while [ $remaining -gt 0 ]; do
        printf "\rSelect disk [1-%d] or press Enter for default (%ds remaining): " "${#disks[@]}" "$remaining" >&2

        if read -t 1 -n 1 input; then
            if [ -z "$input" ]; then
                # Enter pressed - use default
                selected=$default_idx
                break
            elif [[ "$input" =~ ^[0-9]$ ]] && [ "$input" -ge 1 ] && [ "$input" -le ${#disks[@]} ]; then
                selected=$input
                echo "" >&2
                break
            else
                echo "" >&2
                warn "Invalid selection. Please enter 1-${#disks[@]}"
                remaining=$TIMEOUT
            fi
        fi
        ((remaining--))
    done

    if [ -z "$selected" ]; then
        selected=$default_idx
        echo "" >&2
        log "Timeout - auto-selecting default disk"
    fi

    # Get selected disk
    local idx=$((selected - 1))
    IFS='|' read -r TARGET_DISK size model removable <<< "${disks[$idx]}"

    echo "" >&2
    log "Selected: $TARGET_DISK ($size GB - $model)"
    echo "" >&2

    # Final confirmation with shorter timeout
    echo -e "${RED}╔════════════════════════════════════════════════════════════╗${NC}" >&2
    echo -e "${RED}║  ALL DATA ON $TARGET_DISK WILL BE PERMANENTLY DESTROYED!  ║${NC}" >&2
    echo -e "${RED}╚════════════════════════════════════════════════════════════╝${NC}" >&2
    echo "" >&2

    local confirm_timeout=10
    printf "Press 'y' to confirm, any other key to cancel (%ds): " "$confirm_timeout" >&2

    if read -t $confirm_timeout -n 1 confirm; then
        echo "" >&2
        if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
            error "Installation cancelled by user"
        fi
    else
        echo "" >&2
        log "Auto-confirming installation..."
    fi

    # Only this goes to stdout - the actual device path
    echo "$TARGET_DISK"
}

# Partition the disk
partition_disk() {
    local disk=$1

    log "Partitioning $disk..."

    # Wipe existing partition table
    wipefs -a "$disk" >/dev/null 2>&1 || true

    # Create GPT partition table
    parted -s "$disk" mklabel gpt

    # Create EFI partition (512MB)
    parted -s "$disk" mkpart ESP fat32 1MiB 513MiB
    parted -s "$disk" set 1 esp on

    # Create root partition (rest of disk)
    parted -s "$disk" mkpart root ext4 513MiB 100%

    # Wait for partitions to appear
    sleep 2
    partprobe "$disk"
    sleep 1

    # Determine partition names (nvme vs sd)
    if [[ "$disk" == *"nvme"* ]]; then
        EFI_PART="${disk}p1"
        ROOT_PART="${disk}p2"
    else
        EFI_PART="${disk}1"
        ROOT_PART="${disk}2"
    fi

    log "Created partitions: EFI=$EFI_PART, Root=$ROOT_PART"
}

# Format partitions
format_partitions() {
    log "Formatting partitions..."

    mkfs.fat -F 32 -n "EFI" "$EFI_PART"
    mkfs.ext4 -F -L "EcoOS" "$ROOT_PART"

    log "Partitions formatted"
}

# Mount partitions
mount_partitions() {
    log "Mounting partitions..."

    mkdir -p /mnt/target
    mount "$ROOT_PART" /mnt/target

    mkdir -p /mnt/target/boot/efi
    mount "$EFI_PART" /mnt/target/boot/efi

    log "Partitions mounted at /mnt/target"
}

# Copy system files
copy_system() {
    log "Copying system files (this may take several minutes)..."

    # Check if squashfs exists
    if [ ! -f "$SQUASHFS_PATH" ]; then
        # Try alternative paths (including casper paths for Ubuntu)
        for path in /run/live/medium/live/filesystem.squashfs \
                    /cdrom/casper/filesystem.squashfs \
                    /cdrom/live/filesystem.squashfs \
                    /isodevice/casper/filesystem.squashfs \
                    /lib/live/mount/medium/live/filesystem.squashfs \
                    /rofs/../casper/filesystem.squashfs; do
            if [ -f "$path" ]; then
                SQUASHFS_PATH="$path"
                break
            fi
        done
    fi

    # If still not found, try to find it
    if [ ! -f "$SQUASHFS_PATH" ]; then
        local found=$(find /cdrom /run /media -name "filesystem.squashfs" 2>/dev/null | head -1)
        if [ -n "$found" ]; then
            SQUASHFS_PATH="$found"
        fi
    fi

    if [ ! -f "$SQUASHFS_PATH" ]; then
        error "Cannot find filesystem.squashfs"
    fi

    log "Extracting from $SQUASHFS_PATH..."

    # Extract squashfs
    unsquashfs -f -d /mnt/target "$SQUASHFS_PATH"

    log "System files copied"
}

# Configure the installed system
configure_system() {
    log "Configuring system..."

    # Get UUIDs
    local root_uuid=$(blkid -s UUID -o value "$ROOT_PART")
    local efi_uuid=$(blkid -s UUID -o value "$EFI_PART")

    # Create fstab
    cat > /mnt/target/etc/fstab << EOF
# EcoOS fstab
UUID=$root_uuid  /          ext4  defaults,noatime  0  1
UUID=$efi_uuid   /boot/efi  vfat  umask=0077        0  1
EOF

    # Set hostname
    echo "$HOSTNAME" > /mnt/target/etc/hostname
    cat > /mnt/target/etc/hosts << EOF
127.0.0.1   localhost
127.0.1.1   $HOSTNAME

::1         localhost ip6-localhost ip6-loopback
ff02::1     ip6-allnodes
ff02::2     ip6-allrouters
EOF

    # Ensure seat group exists (for seatd)
    if ! grep -q "^seat:" /mnt/target/etc/group; then
        chroot /mnt/target groupadd seat 2>/dev/null || true
    fi

    # Ensure render group exists (for GPU access)
    if ! grep -q "^render:" /mnt/target/etc/group; then
        chroot /mnt/target groupadd render 2>/dev/null || true
    fi

    # Ensure ecouser exists and is configured
    # Groups: video,render (GPU), audio, input (devices), sudo, seat (seatd)
    if ! grep -q "^$USERNAME:" /mnt/target/etc/passwd; then
        chroot /mnt/target useradd -m -s /bin/bash -c "EcoOS User" -G video,render,audio,input,sudo,adm,cdrom,plugdev,seat "$USERNAME"
    else
        # Add to required groups if user already exists
        chroot /mnt/target usermod -a -G video,render,seat "$USERNAME" 2>/dev/null || true
        # Set the GECOS/full name field if missing
        chroot /mnt/target chfn -f "EcoOS User" "$USERNAME" 2>/dev/null || true
    fi

    # Create Sway config directory for ecouser
    mkdir -p /mnt/target/home/$USERNAME/.config/sway
    if [ -f /mnt/target/etc/sway/config ]; then
        cp /mnt/target/etc/sway/config /mnt/target/home/$USERNAME/.config/sway/config
    fi
    chroot /mnt/target chown -R $USERNAME:$USERNAME /home/$USERNAME/.config 2>/dev/null || true

    # Set a default password (ecouser:ecouser) - should be changed on first boot
    echo "$USERNAME:ecouser" | chroot /mnt/target chpasswd

    # Enable sudo for ecouser
    echo "$USERNAME ALL=(ALL) NOPASSWD:ALL" > /mnt/target/etc/sudoers.d/ecouser
    chmod 440 /mnt/target/etc/sudoers.d/ecouser

    # Remove live-boot packages marker if present
    rm -f /mnt/target/etc/live 2>/dev/null || true

    # Enable systemd-networkd for static IP (more reliable than NetworkManager)
    chroot /mnt/target systemctl enable systemd-networkd.service 2>/dev/null || true
    chroot /mnt/target systemctl enable systemd-networkd-wait-online.service 2>/dev/null || true

    # Disable NetworkManager to avoid conflicts
    chroot /mnt/target systemctl disable NetworkManager.service 2>/dev/null || true
    chroot /mnt/target systemctl mask NetworkManager.service 2>/dev/null || true

    # Create network config for systemd-networkd (DHCP for QEMU/VMs)
    mkdir -p /mnt/target/etc/systemd/network
    cat > /mnt/target/etc/systemd/network/10-wired.network << 'NETEOF'
[Match]
Name=ens* enp* eth*

[Network]
DHCP=yes
NETEOF

    # Fix critical directory permissions - /etc must be world-readable
    # for systemd-networkd and other services to read their config files
    # This is CRITICAL - squashfs may have wrong permissions from Docker build
    log "Fixing /etc permissions..."

    # Fix /etc and all subdirectories recursively
    find /mnt/target/etc -type d -exec chmod 755 {} \;

    # Fix critical files that must be world-readable
    for file in passwd group hosts hostname profile bash.bashrc environment \
                shells nsswitch.conf resolv.conf machine-id ld.so.conf; do
        if [ -f "/mnt/target/etc/$file" ]; then
            chmod 644 "/mnt/target/etc/$file"
        fi
    done

    # Shadow files should be root-only readable
    chmod 640 /mnt/target/etc/shadow 2>/dev/null || true
    chmod 640 /mnt/target/etc/gshadow 2>/dev/null || true

    # Sudoers files need specific permissions
    chmod 440 /mnt/target/etc/sudoers 2>/dev/null || true
    find /mnt/target/etc/sudoers.d -type f -exec chmod 440 {} \; 2>/dev/null || true

    # Network config
    chmod 644 /mnt/target/etc/systemd/network/10-wired.network

    log "systemd-networkd enabled for networking"

    # Enable other services for installed system
    chroot /mnt/target systemctl enable eco-daemon.service 2>/dev/null || true
    chroot /mnt/target systemctl enable seatd.service 2>/dev/null || true
    chroot /mnt/target systemctl enable ssh.service 2>/dev/null || true
    chroot /mnt/target systemctl enable debug-network.service 2>/dev/null || true

    log "System configured"
}

# Install bootloader
install_bootloader() {
    log "Installing GRUB bootloader..."

    # Mount necessary filesystems for chroot
    mount --bind /dev /mnt/target/dev
    mount --bind /dev/pts /mnt/target/dev/pts
    mount --bind /proc /mnt/target/proc
    mount --bind /sys /mnt/target/sys
    mount --bind /run /mnt/target/run

    # Fix GRUB default config - remove casper/live boot parameters and add serial console
    if [ -f /mnt/target/etc/default/grub ]; then
        # Remove any boot=casper or live-related parameters
        sed -i 's/boot=casper//g' /mnt/target/etc/default/grub
        # Update GRUB_CMDLINE_LINUX_DEFAULT with serial console for debugging
        sed -i 's/^GRUB_CMDLINE_LINUX_DEFAULT=.*/GRUB_CMDLINE_LINUX_DEFAULT="console=tty1 console=ttyS0,115200n8"/' /mnt/target/etc/default/grub
        # If line doesn't exist, add it
        if ! grep -q "GRUB_CMDLINE_LINUX_DEFAULT" /mnt/target/etc/default/grub; then
            echo 'GRUB_CMDLINE_LINUX_DEFAULT="console=tty1 console=ttyS0,115200n8"' >> /mnt/target/etc/default/grub
        fi
        # Enable serial terminal in GRUB
        echo 'GRUB_TERMINAL="console serial"' >> /mnt/target/etc/default/grub
        echo 'GRUB_SERIAL_COMMAND="serial --speed=115200 --unit=0 --word=8 --parity=no --stop=1"' >> /mnt/target/etc/default/grub
    fi

    # Disable casper-related services
    log "Disabling live boot services..."
    chroot /mnt/target systemctl disable casper.service 2>/dev/null || true
    chroot /mnt/target systemctl disable casper-md5check.service 2>/dev/null || true
    chroot /mnt/target systemctl mask casper.service 2>/dev/null || true
    chroot /mnt/target systemctl mask casper-md5check.service 2>/dev/null || true

    # Remove casper initramfs hooks to prevent live-boot behavior
    log "Removing casper initramfs hooks..."
    rm -rf /mnt/target/usr/share/initramfs-tools/scripts/casper 2>/dev/null || true
    rm -rf /mnt/target/usr/share/initramfs-tools/scripts/casper-premount 2>/dev/null || true
    rm -rf /mnt/target/usr/share/initramfs-tools/scripts/casper-bottom 2>/dev/null || true
    rm -f /mnt/target/usr/share/initramfs-tools/hooks/casper 2>/dev/null || true
    rm -f /mnt/target/etc/initramfs-tools/conf.d/casper.conf 2>/dev/null || true

    # Regenerate initramfs without casper hooks
    log "Regenerating initramfs..."
    chroot /mnt/target update-initramfs -u -k all

    # Ensure proper boot target
    chroot /mnt/target systemctl set-default multi-user.target 2>/dev/null || true

    # Install GRUB
    chroot /mnt/target grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=EcoOS --recheck

    # Generate GRUB config
    chroot /mnt/target update-grub

    # Cleanup mounts (use lazy unmount for stubborn mounts, reverse order)
    sync
    umount -l /mnt/target/run 2>/dev/null || true
    umount -l /mnt/target/sys 2>/dev/null || true
    umount -l /mnt/target/proc 2>/dev/null || true
    umount -l /mnt/target/dev/pts 2>/dev/null || true
    umount -l /mnt/target/dev 2>/dev/null || true

    log "Bootloader installed"
}

# Cleanup and reboot
cleanup_and_reboot() {
    log "Cleaning up..."

    # Sync disks
    sync

    # Unmount
    umount /mnt/target/boot/efi
    umount /mnt/target

    echo ""
    echo -e "${GREEN}╔════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}║          EcoOS Installation Complete!                      ║${NC}"
    echo -e "${GREEN}╚════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo "The system will reboot in 10 seconds..."
    echo "Remove the USB drive when the screen goes blank."
    echo ""

    sleep 10
    reboot
}

# Main installation flow
main() {
    clear

    echo ""
    echo -e "${GREEN}"
    echo "  ███████╗ ██████╗ ██████╗  ██████╗ ███████╗"
    echo "  ██╔════╝██╔════╝██╔═══██╗██╔═══██╗██╔════╝"
    echo "  █████╗  ██║     ██║   ██║██║   ██║███████╗"
    echo "  ██╔══╝  ██║     ██║   ██║██║   ██║╚════██║"
    echo "  ███████╗╚██████╗╚██████╔╝╚██████╔╝███████║"
    echo "  ╚══════╝ ╚═════╝ ╚═════╝  ╚═════╝ ╚══════╝"
    echo -e "${NC}"
    echo "              System Installer v1.0"
    echo ""

    # Check if running as root
    if [ "$(id -u)" -ne 0 ]; then
        error "This script must be run as root"
    fi

    # Select target disk
    TARGET_DISK=$(select_disk)

    # Partition disk
    partition_disk "$TARGET_DISK"

    # Format partitions
    format_partitions

    # Mount partitions
    mount_partitions

    # Copy system
    copy_system

    # Configure system
    configure_system

    # Install bootloader
    install_bootloader

    # Cleanup and reboot
    cleanup_and_reboot
}

# Run main function
main "$@"