#!/bin/bash
# Build Dangerous Pi image using pi-gen
#
# For faster builds, you can use pre-built PM3 binaries:
#   PM3_TARBALL=/path/to/pm3-*.tar.gz ./build-image.sh full
#
# For faster package downloads, run apt-cacher-ng locally:
#   docker run -d -p 3142:3142 --name apt-cache sameersbn/apt-cacher-ng
#   APT_PROXY=http://host.docker.internal:3142 ./build-image.sh full
#
# To skip base OS build (stages 0-2), use a pre-baked base image:
#   BASE_IMAGE=/path/to/base-image.tar.gz ./build-image.sh full

set -e

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Configuration
PI_GEN_DIR="/home/work/pi-gen-builder"  # Official pi-gen builder
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
DTPI_STAGE_DIR="$SCRIPT_DIR/pi-gen/stageDangerousPi"
CONFIG_FILE="$SCRIPT_DIR/pi-gen/config"
KEEP_CONTAINER=0

# Parse flags (after build type positional arg)
for arg in "$@"; do
    case $arg in
        --keep-container)
            KEEP_CONTAINER=1
            shift
            ;;
    esac
done

# Verify pi-gen directory exists
if [ ! -d "$PI_GEN_DIR" ]; then
    echo "Error: pi-gen directory not found at $PI_GEN_DIR"
    echo "Please run: cd /home/work && git clone https://github.com/RPi-Distro/pi-gen.git pi-gen-builder"
    exit 1
fi

echo -e "${GREEN}Using pi-gen at: $PI_GEN_DIR${NC}"

# Build frontend
build_frontend() {
    echo -e "${GREEN}Building frontend...${NC}"
    cd "$SCRIPT_DIR/app/frontend"

    # Check if npm is available
    if ! command -v npm &> /dev/null; then
        echo -e "${YELLOW}WARNING: npm not found, skipping frontend build${NC}"
        echo -e "${YELLOW}Install Node.js to build frontend: https://nodejs.org/${NC}"
        cd "$SCRIPT_DIR"
        return 1
    fi

    # Install dependencies if needed
    if [ ! -d "node_modules" ]; then
        echo -e "${GREEN}Installing frontend dependencies...${NC}"
        npm install
    fi

    # Build the frontend
    echo -e "${GREEN}Running npm build...${NC}"
    npm run build

    echo -e "${GREEN}Frontend build complete${NC}"
    cd "$SCRIPT_DIR"
    return 0
}

# Helper function to copy stages
copy_stages() {
    cd "$PI_GEN_DIR"

    # Copy config
    cp "$CONFIG_FILE" config

    # Copy DangerousPi stage
    rm -rf stageDangerousPi
    cp -r "$DTPI_STAGE_DIR" .

    # Remove intermediate image exports (we only want the final image)
    rm -f stage2/EXPORT_IMAGE

    # Build frontend before copying
    if build_frontend; then
        echo -e "${GREEN}Frontend built successfully${NC}"
    else
        echo -e "${YELLOW}Continuing without frontend build${NC}"
    fi

    # Return to pi-gen directory
    cd "$PI_GEN_DIR"

    # Copy Dangerous Pi source files into the stage
    echo -e "${GREEN}Copying Dangerous Pi source files...${NC}"
    mkdir -p stageDangerousPi/04-dangerous-pi/files
    cp -r "$SCRIPT_DIR/app" stageDangerousPi/04-dangerous-pi/files/
    cp -r "$SCRIPT_DIR/systemd" stageDangerousPi/04-dangerous-pi/files/
    cp -r "$SCRIPT_DIR/scripts" stageDangerousPi/04-dangerous-pi/files/
    cp "$SCRIPT_DIR/requirements.txt" stageDangerousPi/04-dangerous-pi/files/

    # Copy dt-design-system packages (needed by frontend's file: dependency)
    DT_DS_DIR="$(dirname "$SCRIPT_DIR")/dt-design-system"
    if [ -d "$DT_DS_DIR/packages/web" ]; then
        echo -e "${GREEN}Copying dt-design-system packages...${NC}"
        mkdir -p stageDangerousPi/04-dangerous-pi/files/dt-design-system/packages
        cp -r "$DT_DS_DIR/packages/web" stageDangerousPi/04-dangerous-pi/files/dt-design-system/packages/
        cp -r "$DT_DS_DIR/packages/tokens" stageDangerousPi/04-dangerous-pi/files/dt-design-system/packages/
        cp "$DT_DS_DIR/package.json" stageDangerousPi/04-dangerous-pi/files/dt-design-system/
    else
        echo -e "${YELLOW}WARNING: dt-design-system not found at $DT_DS_DIR${NC}"
        echo -e "${YELLOW}Frontend npm ci will fail without it${NC}"
    fi

    # Pass PM3_TARBALL into the stage if set
    if [ -n "${PM3_TARBALL:-}" ] && [ -f "$PM3_TARBALL" ]; then
        echo -e "${GREEN}Staging PM3 tarball: $PM3_TARBALL${NC}"
        mkdir -p stageDangerousPi/02-pm3-install/files
        cp "$PM3_TARBALL" stageDangerousPi/02-pm3-install/files/
    fi
}

# Install pre-baked base image if BASE_IMAGE is set
install_base_image() {
    if [ -z "${BASE_IMAGE:-}" ]; then
        return 1
    fi

    if [ ! -f "$BASE_IMAGE" ]; then
        echo -e "${YELLOW}WARNING: BASE_IMAGE file not found: $BASE_IMAGE${NC}"
        return 1
    fi

    echo -e "${GREEN}Installing pre-baked base image: $BASE_IMAGE${NC}"
    cd "$PI_GEN_DIR"

    # Extract base image into work directory as stage2 rootfs
    local WORK_DIR="work/${IMG_NAME:-dangerous-pi}"
    mkdir -p "${WORK_DIR}/stage2"
    tar -xzf "$BASE_IMAGE" -C "${WORK_DIR}/stage2/"
    echo -e "${GREEN}Base image extracted, skipping stages 0-2${NC}"

    # Mark stages 0-2 as SKIP
    touch stage0/SKIP stage1/SKIP stage2/SKIP
    return 0
}

# Build type
BUILD_TYPE="${1:-full}"

case "$BUILD_TYPE" in
    full)
        echo -e "${GREEN}Starting full build (all stages)...${NC}"

        # Remove container but KEEP the volume for caching
        if docker ps -a --filter name=pigen_work --format "{{.Names}}" | grep -q pigen_work; then
            echo -e "${YELLOW}Removing existing container (keeping cache)...${NC}"
            docker rm pigen_work > /dev/null 2>&1 || true
        fi

        copy_stages

        # Use pre-baked base image if available, otherwise build all stages
        if install_base_image; then
            echo -e "${GREEN}Using pre-baked base image (stages 0-2 skipped)${NC}"
        else
            # Remove any SKIP files to ensure full build
            rm -f stage0/SKIP stage1/SKIP stage2/SKIP
        fi
        rm -f stageDangerousPi/SKIP

        # Full build (uses docker, no root needed)
        PIGEN_DOCKER_MODE=1 ./build-docker.sh
        ;;

    from-dtpi)
        echo -e "${GREEN}Building from DangerousPi stage (customizations only)...${NC}"
        copy_stages

        # Check for cached stage2
        if ls work/*/stage2/*.qcow2 2>/dev/null | grep -q . || \
           [ -d "work/dangerous-pi/stage2/rootfs" ]; then
            echo -e "${GREEN}Found cached base stages, creating SKIP markers...${NC}"
            touch stage0/SKIP stage1/SKIP stage2/SKIP
            echo -e "${GREEN}Will build: stageDangerousPi only (~5 min)${NC}"
        elif install_base_image; then
            echo -e "${GREEN}Using pre-baked base image${NC}"
        else
            echo -e "${YELLOW}WARNING: No cached stages found!${NC}"
            echo -e "${YELLOW}You need to run './build-image.sh full' first${NC}"
            echo -e "${YELLOW}Falling back to full build...${NC}"
            rm -f stage0/SKIP stage1/SKIP stage2/SKIP
        fi

        # Ensure DangerousPi stage builds
        rm -f stageDangerousPi/SKIP

        # Build with continue flag
        CONTINUE=1 PIGEN_DOCKER_MODE=1 ./build-docker.sh
        ;;

    clean)
        echo -e "${YELLOW}Removing ALL build cache and container...${NC}"
        if docker ps -a --filter name=pigen_work --format "{{.Names}}" | grep -q pigen_work; then
            docker rm -v pigen_work > /dev/null 2>&1 || true
            echo -e "${GREEN}Removed container and volume${NC}"
        else
            docker volume rm pigen_work > /dev/null 2>&1 || true
            echo -e "${GREEN}Removed volume${NC}"
        fi
        echo ""
        echo "Cache wiped. Next build will be from scratch."
        exit 0
        ;;

    test)
        echo -e "${GREEN}Validating stage scripts...${NC}"

        # Syntax check - PM3 install
        bash -n "$DTPI_STAGE_DIR/02-pm3-install/00-run.sh"
        bash -n "$DTPI_STAGE_DIR/02-pm3-install/00-run-chroot.sh"
        bash -n "$DTPI_STAGE_DIR/02-pm3-install/00-run-chroot-compile.sh"

        # Syntax check - WiFi AP
        bash -n "$DTPI_STAGE_DIR/01-Wireless-AP/00-run-chroot.sh"

        # Syntax check - ttyd
        bash -n "$DTPI_STAGE_DIR/03-ttyd/00-run-chroot.sh"

        # Syntax check - Dangerous Pi
        bash -n "$DTPI_STAGE_DIR/04-dangerous-pi/00-run.sh"
        bash -n "$DTPI_STAGE_DIR/04-dangerous-pi/00-run-chroot.sh"
        bash -n "$DTPI_STAGE_DIR/04-dangerous-pi/01-run-chroot.sh"

        # Syntax check - OS updates
        if [ -f "$DTPI_STAGE_DIR/07-os-updates/00-run-chroot.sh" ]; then
            bash -n "$DTPI_STAGE_DIR/07-os-updates/00-run-chroot.sh"
            bash -n "$DTPI_STAGE_DIR/07-os-updates/01-run-chroot.sh"
        fi

        echo -e "${GREEN}All scripts pass syntax validation${NC}"

        # Check required files exist
        [ -d "$(dirname "$0")/app" ] || { echo "Error: app/ directory not found"; exit 1; }
        [ -f "$(dirname "$0")/requirements.txt" ] || { echo "Error: requirements.txt not found"; exit 1; }
        [ -d "$(dirname "$0")/systemd" ] || { echo "Error: systemd/ directory not found"; exit 1; }

        echo -e "${GREEN}All required files exist${NC}"
        ;;

    *)
        echo "Usage: $0 {full|from-dtpi|test|clean} [--keep-container]"
        echo ""
        echo "Build Commands:"
        echo "  full      - Build all stages (~15-20 min with pre-built PM3, cached base)"
        echo "              Use for: First build or when base system needs updates"
        echo "              Preserves cache for fast rebuilds"
        echo ""
        echo "  from-dtpi - Rebuild only customizations (~5 min)"
        echo "              Use for: Daily development, app changes"
        echo "              Skips: stages 0-2, PM3 (uses cache)"
        echo "              Requires: Previous 'full' build"
        echo ""
        echo "Utility Commands:"
        echo "  test      - Validate all scripts without building"
        echo "  clean     - Wipe ALL cache (next build will be from scratch)"
        echo ""
        echo "Flags:"
        echo "  --keep-container  - Keep Docker container after build (for debugging)"
        echo ""
        echo "Environment Variables:"
        echo "  PM3_TARBALL       Path to pre-built PM3 tarball (aarch64)"
        echo "                    Skips PM3 compilation (~45 min savings)"
        echo "                    Example: PM3_TARBALL=./pm3-v4.20728-aarch64-linux-cp312.tar.gz"
        echo ""
        echo "  BASE_IMAGE        Path to pre-baked base image tarball (stages 0-2)"
        echo "                    Skips base OS build (~30 min savings on cold build)"
        echo "                    Example: BASE_IMAGE=./base-image-trixie-arm64.tar.gz"
        echo ""
        echo "  APT_PROXY         URL of apt-cacher-ng proxy for faster package downloads"
        echo "                    Example: APT_PROXY=http://host.docker.internal:3142"
        echo ""
        echo "Stage Structure:"
        echo "  stage0, stage1, stage2  -> Base Raspberry Pi OS (cached)"
        echo "  stageDangerousPi        -> PM3 + WiFi AP + ttyd + app + UPS + HTTPS + updates"
        exit 1
        ;;
esac

echo -e "${GREEN}Build complete!${NC}"
echo "Image location: $PI_GEN_DIR/deploy/"

# Clean up container but preserve work volume for future builds
if [ $KEEP_CONTAINER -eq 0 ]; then
    if docker ps -a --filter name=pigen_work --format "{{.Names}}" | grep -q pigen_work; then
        echo -e "${GREEN}Cleaning up build container (preserving cache)...${NC}"
        docker rm pigen_work > /dev/null 2>&1
        echo -e "${GREEN}Container removed, cache preserved for fast rebuilds${NC}"
    fi
else
    echo -e "${YELLOW}Container kept for debugging (--keep-container flag)${NC}"
    echo "Inspect with: docker exec -it pigen_work bash"
    echo "Remove with: docker rm -v pigen_work"
fi