#!/bin/bash

# Companion Dashboard Raspberry Pi Kiosk Install Script
# This script installs Docker, sets up the Companion Dashboard app, and configures kiosk mode

set -e  # Exit on any error

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

# Configuration
REPO_URL="https://github.com/tomhillmeyer/companion-dashboard"
CONTAINER_NAME="companion-dashboard"
IMAGE_NAME="companion-dashboard:latest"
APP_PORT=${APP_PORT:-3000}  # Internal app port

# 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 running as root
check_root() {
    if [[ $EUID -eq 0 ]]; then
        print_error "This script should not be run as root. Please run as a regular user with sudo privileges."
        exit 1
    fi
}

# Function to detect if this is an update vs fresh install
detect_installation_type() {
    if [ -d ~/companion-dashboard-kiosk ] && [ -f ~/companion-dashboard-kiosk/Dockerfile ]; then
        print_status "Existing installation detected - running update"
        return 0  # Update mode
    else
        print_status "No existing installation found - running fresh install"
        return 1  # Fresh install mode
    fi
}

# Function to check if we're on a Raspberry Pi or compatible ARM system
check_system() {
    print_status "Checking system compatibility..."
    
    # Check if we're on ARM architecture
    ARCH=$(uname -m)
    print_status "Detected architecture: $ARCH"
    
    if [[ "$ARCH" != "armv7l" && "$ARCH" != "aarch64" && "$ARCH" != "arm64" ]]; then
        print_warning "This script is designed for ARM-based systems (Raspberry Pi). Detected: $ARCH"
        print_warning "Continuing anyway, but some steps may need adjustment."
    fi
    
    # Check if we have apt (Debian/Ubuntu based)
    if ! command -v apt >/dev/null 2>&1; then
        print_error "This script requires apt package manager (Debian/Ubuntu/Raspbian)"
        exit 1
    fi
    
    # Test network connectivity
    print_status "Testing network connectivity..."
    if ! ping -c 1 -W 5 8.8.8.8 >/dev/null 2>&1; then
        print_error "No internet connectivity detected. Please check your network connection."
        exit 1
    fi
    
    print_success "System check passed"
}

# Function to update system packages with timeout
update_system() {
    print_status "Updating system packages (this may take a few minutes)..."
    
    # Update package list with timeout
    print_status "Updating package lists..."
    if ! timeout 300 sudo apt update; then
        print_error "apt update failed or timed out (5 minutes). Check your internet connection and repository configuration."
        exit 1
    fi
    
    # Upgrade packages with timeout
    print_status "Upgrading packages..."
    if ! timeout 600 sudo apt upgrade -y; then
        print_warning "apt upgrade failed or timed out (10 minutes). Continuing anyway..."
    fi
    
    print_success "System updated"
}

# Function to install required system packages
install_system_packages() {
    print_status "Installing required system packages..."
    
    # Essential packages for GUI and Docker
    PACKAGES=(
        "curl"
        "wget"
        "git"
        "unzip"
        "ca-certificates"
        "gnupg"
        "lsb-release"
        "xorg"
        "openbox"
        "chromium-browser"
        "x11-xserver-utils"
        "unclutter"
        "sed"
    )
    
    # Install packages with timeout
    if ! timeout 600 sudo apt install -y "${PACKAGES[@]}"; then
        print_error "Package installation failed or timed out. Please check your internet connection."
        exit 1
    fi
    
    print_success "System packages installed"
}

# Function to check and fix Raspberry Pi cgroup configuration
fix_rpi_cgroups() {
    print_status "Checking Raspberry Pi cgroup configuration..."
    
    # Check if we're likely on a Raspberry Pi
    if [ ! -f /boot/cmdline.txt ] && [ ! -f /boot/firmware/cmdline.txt ]; then
        print_status "Not a Raspberry Pi system, skipping cgroup fix"
        return
    fi
    
    # Determine the correct cmdline.txt location
    CMDLINE_FILE=""
    if [ -f /boot/firmware/cmdline.txt ]; then
        CMDLINE_FILE="/boot/firmware/cmdline.txt"
    elif [ -f /boot/cmdline.txt ]; then
        CMDLINE_FILE="/boot/cmdline.txt"
    fi
    
    if [ -z "$CMDLINE_FILE" ]; then
        print_warning "Could not find cmdline.txt, skipping cgroup fix"
        return
    fi
    
    # Check if cgroup parameters are already present
    if grep -q "cgroup_enable=cpuset cgroup_enable=memory cgroup_memory=1" "$CMDLINE_FILE"; then
        print_success "Raspberry Pi cgroup configuration already correct"
        return
    fi
    
    print_status "Adding cgroup configuration to $CMDLINE_FILE..."
    
    # Backup the original file
    sudo cp "$CMDLINE_FILE" "${CMDLINE_FILE}.backup"
    
    # Add cgroup parameters to the end of the line (not a new line)
    sudo sed -i '$ s/$/ cgroup_enable=cpuset cgroup_enable=memory cgroup_memory=1/' "$CMDLINE_FILE"
    
    print_success "Raspberry Pi cgroup configuration updated"
    print_warning "A reboot will be required for cgroup changes to take effect"
    
    # Set flag that reboot is needed
    REBOOT_REQUIRED=true
}

# Function to install Docker with better error handling
install_docker() {
    if command -v docker >/dev/null 2>&1; then
        print_status "Docker is already installed, checking if it's working..."
        
        # Test if Docker is actually working
        if timeout 30 sudo docker run --rm hello-world >/dev/null 2>&1; then
            print_success "Docker is installed and working correctly"
            return
        else
            print_warning "Docker is installed but not working correctly, attempting to fix..."
        fi
    fi
    
    print_status "Installing Docker..."
    
    # Method 1: Try official Docker repository installation
    if install_docker_official; then
        print_success "Docker installed successfully via official repository"
        return
    fi
    
    print_warning "Official Docker installation failed, trying convenience script..."
    
    # Method 2: Try Docker convenience script as fallback
    if install_docker_convenience; then
        print_success "Docker installed successfully via convenience script"
        return
    fi
    
    print_error "All Docker installation methods failed"
    print_error "Please install Docker manually and re-run this script"
    exit 1
}

# Function to install Docker via official repository
install_docker_official() {
    print_status "Attempting official Docker repository installation..."
    
    # Clean up any failed installation
    sudo apt remove --purge -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin 2>/dev/null || true
    
    # Add Docker's official GPG key with timeout
    sudo mkdir -p /etc/apt/keyrings
    if ! timeout 60 curl -fsSL https://download.docker.com/linux/debian/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg; then
        print_warning "Failed to add Docker GPG key"
        return 1
    fi
    
    # Add Docker repository
    echo \
      "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/debian \
      $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
    
    # Install Docker
    if ! timeout 300 sudo apt update; then
        print_warning "Failed to update package list for Docker repository"
        return 1
    fi
    
    if ! timeout 600 sudo apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin; then
        print_warning "Failed to install Docker packages"
        return 1
    fi
    
    # Add current user to docker group
    sudo usermod -aG docker $USER
    
    # Try to start Docker service
    if ! sudo systemctl enable docker; then
        print_warning "Failed to enable Docker service"
        return 1
    fi
    
    if ! sudo systemctl start docker; then
        print_warning "Docker service failed to start, checking for common issues..."
        
        # Check if it's a cgroup issue and we haven't rebooted yet
        if [ "$REBOOT_REQUIRED" = true ]; then
            print_warning "Cgroup configuration was just updated. Docker may work after reboot."
            print_warning "Continuing with installation, but you may need to reboot and re-run the script."
            return 0
        fi
        
        return 1
    fi
    
    # Test Docker with timeout
    if timeout 60 sudo docker run --rm hello-world >/dev/null 2>&1; then
        return 0
    else
        print_warning "Docker test failed"
        return 1
    fi
}

# Function to install Docker via convenience script
install_docker_convenience() {
    print_status "Downloading Docker convenience script..."
    
    if ! timeout 60 curl -fsSL https://get.docker.com -o /tmp/get-docker.sh; then
        print_warning "Failed to download Docker convenience script"
        return 1
    fi
    
    if ! timeout 600 sudo sh /tmp/get-docker.sh; then
        print_warning "Docker convenience script failed"
        return 1
    fi
    
    # Clean up
    rm -f /tmp/get-docker.sh
    
    # Add current user to docker group
    sudo usermod -aG docker $USER
    
    # Try to start Docker service
    if ! sudo systemctl enable docker; then
        return 1
    fi
    
    if ! sudo systemctl start docker; then
        return 1
    fi
    
    # Test Docker
    if timeout 60 sudo docker run --rm hello-world >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to create Dockerfile
create_dockerfile() {
    print_status "Creating/updating Dockerfile..."
    
    mkdir -p ~/companion-dashboard-kiosk
    cd ~/companion-dashboard-kiosk
    
    # Always recreate Dockerfile to ensure it's up to date
    cat > Dockerfile << 'EOF'
# Use Node.js LTS on Debian
FROM node:18-bullseye

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    python3 \
    python3-pip \
    build-essential \
    libnss3-dev \
    libatk-bridge2.0-dev \
    libdrm2 \
    libxcomposite1 \
    libxdamage1 \
    libxrandr2 \
    libgbm1 \
    libxss1 \
    libasound2 \
    libatspi2.0-0 \
    libgtk-3-0 \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Clone the repository (always get latest)
RUN git clone https://github.com/tomhillmeyer/companion-dashboard.git . || (rm -rf .git && git clone https://github.com/tomhillmeyer/companion-dashboard.git .)

# Install dependencies
RUN npm install

# Build the application
RUN npm run build || npm run electron:build || echo "Build step may vary"

# Note: No EXPOSE needed since we're only accessing locally

# Create a startup script
RUN echo '#!/bin/bash\n\
if [ -f "dist/main.js" ]; then\n\
    node dist/main.js\n\
elif [ -f "build/main.js" ]; then\n\
    node build/main.js\n\
elif [ -f "main.js" ]; then\n\
    node main.js\n\
else\n\
    npm start\n\
fi' > /app/start.sh && chmod +x /app/start.sh

# Start the application
CMD ["/app/start.sh"]
EOF
    
    print_success "Dockerfile created/updated"
}

# Function to build Docker image
build_docker_image() {
    print_status "Building Docker image (this may take several minutes)..."
    
    cd ~/companion-dashboard-kiosk
    
    # Stop and remove existing container if it exists
    docker stop $CONTAINER_NAME 2>/dev/null || true
    docker rm $CONTAINER_NAME 2>/dev/null || true
    
    # Remove old image to force fresh build
    docker rmi $IMAGE_NAME 2>/dev/null || true
    
    # Build the image with no cache to ensure latest code
    # Use timeout to prevent hanging
    if ! timeout 1800 docker build --no-cache -t $IMAGE_NAME .; then
        print_error "Docker build failed or timed out (30 minutes)"
        exit 1
    fi
    
    print_success "Docker image built successfully"
}

# Function to create systemd service for auto-start
create_systemd_service() {
    print_status "Creating systemd service for auto-start..."
    
    # Create the service file
    sudo tee /etc/systemd/system/companion-dashboard-kiosk.service > /dev/null << EOF
[Unit]
Description=Companion Dashboard Kiosk
After=docker.service
Requires=docker.service

[Service]
Type=simple
User=$USER
Environment=DISPLAY=:0
Environment=XAUTHORITY=/home/$USER/.Xauthority
ExecStartPre=/usr/bin/docker stop $CONTAINER_NAME || true
ExecStartPre=/usr/bin/docker rm $CONTAINER_NAME || true
ExecStart=/usr/bin/docker run --rm --name $CONTAINER_NAME --network host -e DISPLAY=:0 $IMAGE_NAME
ExecStop=/usr/bin/docker stop $CONTAINER_NAME
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF
    
    sudo systemctl daemon-reload
    sudo systemctl enable companion-dashboard-kiosk.service
    
    print_success "Systemd service created and enabled"
}

# Function to configure X11 and kiosk mode
setup_kiosk_mode() {
    print_status "Setting up kiosk mode..."
    
    # Create .xinitrc for kiosk mode
    cat > ~/.xinitrc << EOF
#!/bin/bash

# Disable screen saver and power management
xset s off
xset -dpms
xset s noblank

# Hide cursor after inactivity
unclutter -idle 0.5 -root &

# Start window manager
openbox-session &

# Wait for Docker container to be ready
sleep 10

# Start Chromium in kiosk mode pointing to the app
chromium-browser --noerrdialogs --disable-infobars --kiosk --app=http://localhost:$APP_PORT &

# Keep X session alive
while true; do
    sleep 60
done
EOF
    
    chmod +x ~/.xinitrc
    
    # Configure auto-login and auto-start X (for Raspberry Pi OS)
    if [ -f /etc/systemd/system/getty@tty1.service.d/override.conf ]; then
        print_warning "Auto-login may already be configured"
    else
        sudo mkdir -p /etc/systemd/system/getty@tty1.service.d/
        sudo tee /etc/systemd/system/getty@tty1.service.d/override.conf > /dev/null << EOF
[Service]
ExecStart=
ExecStart=-/sbin/agetty --noissue --autologin $USER %I \$TERM
Type=idle
EOF
    fi
    
    # Add startx to .bash_profile for auto-start X
    if ! grep -q "startx" ~/.bash_profile 2>/dev/null; then
        echo "
# Auto-start X session on login
if [ -z \"\$DISPLAY\" ] && [ \"\$XDG_VTNR\" = 1 ]; then
    exec startx
fi" >> ~/.bash_profile
    fi
    
    print_success "Kiosk mode configured"
}

# Function to create management scripts
create_management_scripts() {
    print_status "Creating management scripts..."
    
    mkdir -p ~/companion-dashboard-kiosk/scripts
    
    # Start script
    cat > ~/companion-dashboard-kiosk/scripts/start.sh << 'EOF'
#!/bin/bash
echo "Starting Companion Dashboard..."
sudo systemctl start companion-dashboard-kiosk.service
sudo systemctl status companion-dashboard-kiosk.service
EOF
    
    # Stop script
    cat > ~/companion-dashboard-kiosk/scripts/stop.sh << 'EOF'
#!/bin/bash
echo "Stopping Companion Dashboard..."
sudo systemctl stop companion-dashboard-kiosk.service
EOF
    
    # Restart script
    cat > ~/companion-dashboard-kiosk/scripts/restart.sh << 'EOF'
#!/bin/bash
echo "Restarting Companion Dashboard..."
sudo systemctl restart companion-dashboard-kiosk.service
sudo systemctl status companion-dashboard-kiosk.service
EOF
    
    # Update script
    cat > ~/companion-dashboard-kiosk/scripts/update.sh << 'EOF'
#!/bin/bash
echo "Updating Companion Dashboard..."
cd ~/companion-dashboard-kiosk
sudo systemctl stop companion-dashboard-kiosk.service
docker rmi companion-dashboard:latest || true
docker build -t companion-dashboard:latest .
sudo systemctl start companion-dashboard-kiosk.service
echo "Update complete!"
EOF
    
    # Logs script
    cat > ~/companion-dashboard-kiosk/scripts/logs.sh << 'EOF'
#!/bin/bash
echo "Showing Companion Dashboard logs (Ctrl+C to exit)..."
sudo journalctl -u companion-dashboard-kiosk.service -f
EOF
    
    chmod +x ~/companion-dashboard-kiosk/scripts/*.sh
    
    print_success "Management scripts created in ~/companion-dashboard-kiosk/scripts/"
}

# Main installation function
main() {
    print_status "Starting Companion Dashboard Raspberry Pi Kiosk installation..."
    
    check_root
    detect_installation_type
    IS_UPDATE=$?
    
    check_system
    
    if [ $IS_UPDATE -eq 1 ]; then
        # Fresh install - do everything
        update_system
        install_system_packages
        fix_rpi_cgroups  # Fix Raspberry Pi cgroup issues before Docker
        install_docker
        create_dockerfile
        build_docker_image
        create_systemd_service
        setup_kiosk_mode
        create_management_scripts
        
        print_success "Fresh installation completed successfully!"
        echo
        print_status "Next steps:"
        if [ "$REBOOT_REQUIRED" = true ]; then
            echo -e "  ${YELLOW}IMPORTANT: A reboot is required for cgroup changes to take effect${NC}"
            echo -e "  1. Reboot your Raspberry Pi: ${YELLOW}sudo reboot${NC}"
            echo -e "  2. After reboot, the dashboard should start automatically in kiosk mode"
            echo -e "  3. If the service doesn't start automatically, run: ${YELLOW}sudo systemctl start companion-dashboard-kiosk.service${NC}"
        else
            echo -e "  1. Reboot your Raspberry Pi: ${YELLOW}sudo reboot${NC}"
            echo -e "  2. After reboot, the dashboard should start automatically in kiosk mode"
        fi
    else
        # Update existing install
        print_status "Updating existing installation..."
        sudo systemctl stop companion-dashboard-kiosk.service 2>/dev/null || true
        
        create_dockerfile
        build_docker_image
        create_management_scripts  # Refresh management scripts
        
        sudo systemctl start companion-dashboard-kiosk.service
        
        print_success "Update completed successfully!"
        echo
        print_status "The service has been restarted with the latest version"
    fi
    
    echo -e "  3. Use the management scripts in ~/companion-dashboard-kiosk/scripts/:"
    echo -e "     - ${YELLOW}./start.sh${NC}   - Start the service"
    echo -e "     - ${YELLOW}./stop.sh${NC}    - Stop the service"
    echo -e "     - ${YELLOW}./restart.sh${NC} - Restart the service"
    echo -e "     - ${YELLOW}./update.sh${NC}  - Update the app"
    echo -e "     - ${YELLOW}./logs.sh${NC}    - View logs"
    echo
    print_status "The dashboard will be available locally at: http://localhost:$APP_PORT"
    echo -e "  Configure your Companion connection in the dashboard settings"
    echo -e "  The app has full network access to connect to your Companion instance"
    echo
    if [ $IS_UPDATE -eq 1 ]; then
        print_warning "If you encounter Docker permission issues, log out and log back in"
    fi
}

# Run main function
main "$@"