#!/bin/bash -e

# Enable console access for Pi Zero
# Supports both:
# 1. USB gadget serial (g_serial) - connect via USB data port
# 2. Hardware UART (GPIO pins) - connect via CP210x/FTDI/PL2303 cable
# Date: 2025-11-27

on_chroot << EOF
echo "Configuring console access..."

# Enable USB serial console in systemd (for USB gadget mode)
systemctl enable serial-getty@ttyGS0.service

# Enable hardware UART console (for GPIO serial cables like CP210x/FTDI)
# On Pi Zero 2 W, serial0 maps to ttyS0 (mini UART) when Bluetooth is enabled
# or ttyAMA0 (PL011) when Bluetooth is disabled
systemctl enable serial-getty@serial0.service

# Add g_serial to /etc/modules for automatic loading (more reliable than cmdline.txt)
if ! grep -q "^g_serial" /etc/modules; then
    echo "g_serial" >> /etc/modules
    echo "✓ Added g_serial to /etc/modules"
else
    echo "✓ g_serial already in /etc/modules"
fi

echo "✓ USB console service enabled"
EOF

# Modify boot configuration (runs outside chroot, modifies boot partition)
echo "Enabling USB gadget mode in boot config..."

# Detect boot partition location (newer images use /boot/firmware)
if [ -f "${ROOTFS_DIR}/boot/firmware/config.txt" ]; then
    BOOT_DIR="${ROOTFS_DIR}/boot/firmware"
elif [ -f "${ROOTFS_DIR}/boot/config.txt" ]; then
    BOOT_DIR="${ROOTFS_DIR}/boot"
else
    echo "ERROR: Could not find boot partition"
    exit 1
fi

echo "Using boot directory: $BOOT_DIR"

# Add dwc2 overlay to config.txt (peripheral mode for USB gadget)
CONFIG_FILE="$BOOT_DIR/config.txt"
if ! grep -q "dtoverlay=dwc2" "$CONFIG_FILE"; then
    echo "" >> "$CONFIG_FILE"
    echo "# Enable USB gadget mode (OTG) for USB console access" >> "$CONFIG_FILE"
    echo "dtoverlay=dwc2,dr_mode=peripheral" >> "$CONFIG_FILE"
    echo "✓ Added dwc2 overlay to config.txt"
else
    echo "✓ dwc2 overlay already in config.txt"
fi

# Enable hardware UART for GPIO serial console (CP210x/FTDI/PL2303 cables)
if ! grep -q "^enable_uart=1" "$CONFIG_FILE"; then
    echo "" >> "$CONFIG_FILE"
    echo "# Enable hardware UART for GPIO serial console" >> "$CONFIG_FILE"
    echo "enable_uart=1" >> "$CONFIG_FILE"
    echo "✓ Added enable_uart=1 to config.txt"
else
    echo "✓ enable_uart already enabled in config.txt"
fi

# Add g_serial module to cmdline.txt
CMDLINE_FILE="$BOOT_DIR/cmdline.txt"
if [ -f "$CMDLINE_FILE" ]; then
    # Check if modules-load is already present
    if ! grep -q "modules-load=dwc2,g_serial" "$CMDLINE_FILE"; then
        # Add after rootwait
        sed -i 's/rootwait/rootwait modules-load=dwc2,g_serial/' "$CMDLINE_FILE"
        echo "✓ Added g_serial module to cmdline.txt"
    else
        echo "✓ g_serial module already in cmdline.txt"
    fi

    # Set default CPU cores for Pi Zero 2 W (2 of 4 cores for thermal management)
    # User can change this via Settings > System > CPU Cores
    if ! grep -q "maxcpus=" "$CMDLINE_FILE"; then
        sed -i 's/$/ maxcpus=2/' "$CMDLINE_FILE"
        echo "✓ Set maxcpus=2 for Pi Zero 2 W thermal management"
    else
        echo "✓ maxcpus already configured in cmdline.txt"
    fi
else
    echo "ERROR: cmdline.txt not found at $CMDLINE_FILE"
    exit 1
fi

echo "=== USB Console Configuration Complete ==="
echo "Connect USB cable to the Pi Zero's DATA port (left port when facing USB ports)"
echo "Access via: screen /dev/ttyACM0 115200 (Linux/Mac)"
echo "Login: dt / proxmark3"
