#!/bin/bash
# inject-kubesolo.sh — Add KubeSolo binary, init system, and configs to rootfs
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
CACHE_DIR="${CACHE_DIR:-$PROJECT_ROOT/build/cache}"
ROOTFS_DIR="${ROOTFS_DIR:-$PROJECT_ROOT/build/rootfs-work}"
ROOTFS="$ROOTFS_DIR/rootfs"
VERSION="$(cat "$PROJECT_ROOT/VERSION")"

if [ ! -d "$ROOTFS" ]; then
    echo "ERROR: Rootfs not found: $ROOTFS"
    echo "Run extract-core.sh first."
    exit 1
fi

KUBESOLO_BIN="$CACHE_DIR/kubesolo"
if [ ! -f "$KUBESOLO_BIN" ]; then
    echo "ERROR: KubeSolo binary not found: $KUBESOLO_BIN"
    echo "See fetch-components.sh output for instructions."
    exit 1
fi

echo "==> Injecting KubeSolo into rootfs..."

# --- 1. KubeSolo binary ---
# Install to /usr/bin (NOT /usr/local/bin) because /usr/local is bind-mounted
# from the data partition at boot, which would hide the binary.
mkdir -p "$ROOTFS/usr/bin"
cp "$KUBESOLO_BIN" "$ROOTFS/usr/bin/kubesolo"
chmod +x "$ROOTFS/usr/bin/kubesolo"
echo "    Installed KubeSolo binary ($(du -h "$KUBESOLO_BIN" | cut -f1))"

# --- 2. Custom init system ---
echo "    Installing init system..."

# Main init — remove symlink first to avoid clobbering busybox
# (Tiny Core has /sbin/init -> ../bin/busybox; cp follows symlinks)
rm -f "$ROOTFS/sbin/init"
cp "$PROJECT_ROOT/init/init.sh" "$ROOTFS/sbin/init"
chmod +x "$ROOTFS/sbin/init"

# Init stages
mkdir -p "$ROOTFS/usr/lib/kubesolo-os/init.d"
for stage in "$PROJECT_ROOT"/init/lib/*.sh; do
    [ -f "$stage" ] || continue
    cp "$stage" "$ROOTFS/usr/lib/kubesolo-os/init.d/"
    chmod +x "$ROOTFS/usr/lib/kubesolo-os/init.d/$(basename "$stage")"
done
echo "    Installed $(ls "$ROOTFS/usr/lib/kubesolo-os/init.d/" | wc -l) init stages"

# Shared functions
if [ -f "$PROJECT_ROOT/init/lib/functions.sh" ]; then
    cp "$PROJECT_ROOT/init/lib/functions.sh" "$ROOTFS/usr/lib/kubesolo-os/functions.sh"
fi

# Emergency shell
if [ -f "$PROJECT_ROOT/init/emergency-shell.sh" ]; then
    cp "$PROJECT_ROOT/init/emergency-shell.sh" "$ROOTFS/usr/lib/kubesolo-os/emergency-shell.sh"
    chmod +x "$ROOTFS/usr/lib/kubesolo-os/emergency-shell.sh"
fi

# Shared library scripts (network, health)
for lib in network.sh health.sh; do
    src="$PROJECT_ROOT/build/rootfs/usr/lib/kubesolo-os/$lib"
    [ -f "$src" ] && cp "$src" "$ROOTFS/usr/lib/kubesolo-os/$lib"
done

# Cloud-init binary (Go, built separately)
CLOUDINIT_BIN="$CACHE_DIR/kubesolo-cloudinit"
if [ -f "$CLOUDINIT_BIN" ]; then
    cp "$CLOUDINIT_BIN" "$ROOTFS/usr/lib/kubesolo-os/kubesolo-cloudinit"
    chmod +x "$ROOTFS/usr/lib/kubesolo-os/kubesolo-cloudinit"
    echo "    Installed cloud-init binary ($(du -h "$CLOUDINIT_BIN" | cut -f1))"
else
    echo "    WARN: Cloud-init binary not found (run 'make build-cloudinit' to build)"
fi

# Update agent binary (Go, built separately)
UPDATE_BIN="$CACHE_DIR/kubesolo-update"
if [ -f "$UPDATE_BIN" ]; then
    cp "$UPDATE_BIN" "$ROOTFS/usr/lib/kubesolo-os/kubesolo-update"
    chmod +x "$ROOTFS/usr/lib/kubesolo-os/kubesolo-update"
    echo "    Installed update agent ($(du -h "$UPDATE_BIN" | cut -f1))"
else
    echo "    WARN: Update agent not found (run 'make build-update-agent' to build)"
fi

# --- 3. Custom kernel or TCZ kernel modules ---
# If a custom kernel was built (with CONFIG_CGROUP_BPF=y), use it.
# Otherwise fall back to TCZ-extracted modules with manual modules.dep.
CUSTOM_KERNEL_DIR="$CACHE_DIR/custom-kernel"
CUSTOM_VMLINUZ="$CUSTOM_KERNEL_DIR/vmlinuz"
CUSTOM_MODULES="$CUSTOM_KERNEL_DIR/modules"

# Detect kernel version from rootfs
KVER=""
for d in "$ROOTFS"/lib/modules/*/; do
    [ -d "$d" ] && KVER="$(basename "$d")" && break
done

if [ -z "$KVER" ]; then
    echo "    WARN: Could not detect kernel version from rootfs"
fi

echo "    Kernel version: $KVER"

if [ -f "$CUSTOM_VMLINUZ" ] && [ -d "$CUSTOM_MODULES/lib/modules/$KVER" ]; then
    # =========================================================================
    # Custom kernel path — selective module install (only what modules.list needs)
    # =========================================================================
    echo "    Using custom kernel (CONFIG_CGROUP_BPF=y)..."

    # Replace vmlinuz
    cp "$CUSTOM_VMLINUZ" "$ROOTFS_DIR/vmlinuz"
    echo "    Installed custom vmlinuz ($(du -h "$CUSTOM_VMLINUZ" | cut -f1))"

    # Selectively install ONLY modules from modules.list + their transitive deps.
    # This keeps the initramfs minimal — no sound, GPU, SCSI, etc. modules.
    echo "    Installing kernel modules (selective — modules.list + deps only)..."
    CUSTOM_MOD_DIR="$CUSTOM_MODULES/lib/modules/$KVER"

    rm -rf "$ROOTFS/lib/modules/$KVER"
    mkdir -p "$ROOTFS/lib/modules/$KVER/kernel"

    # Copy module metadata files (needed by modprobe)
    for f in modules.builtin modules.builtin.modinfo modules.order \
             modules.builtin.alias.bin modules.builtin.bin; do
        [ -f "$CUSTOM_MOD_DIR/$f" ] && cp "$CUSTOM_MOD_DIR/$f" "$ROOTFS/lib/modules/$KVER/"
    done

    # Use modprobe --show-depends to resolve each module + its transitive deps
    MODULES_LIST="$PROJECT_ROOT/build/config/modules.list"
    NEEDED_MODS=$(mktemp)
    while IFS= read -r mod; do
        # Skip comments and blank lines
        case "$mod" in \#*|"") continue ;; esac
        mod=$(echo "$mod" | xargs)  # trim whitespace
        [ -z "$mod" ] && continue

        # modprobe -S <ver> -d <root> --show-depends <module> lists all deps in load order
        # Output format: "insmod /path/to/module.ko" — extract path with awk
        modprobe -S "$KVER" -d "$CUSTOM_MODULES" --show-depends "$mod" 2>/dev/null \
            | awk '/^insmod/{print $2}' >> "$NEEDED_MODS" \
            || echo "    WARN: modprobe could not resolve: $mod"
    done < "$MODULES_LIST"

    # Deduplicate and copy each needed module
    sort -u "$NEEDED_MODS" | while IFS= read -r mod_path; do
        mod_path=$(echo "$mod_path" | xargs)  # trim whitespace
        [ -z "$mod_path" ] && continue
        # mod_path is absolute (e.g., /path/to/custom-kernel/modules/lib/modules/KVER/kernel/...)
        if [ ! -f "$mod_path" ]; then
            echo "    WARN: module not found: $mod_path"
            continue
        fi
        # Get the relative path under lib/modules/KVER/
        rel_path="${mod_path#$CUSTOM_MOD_DIR/}"
        dst="$ROOTFS/lib/modules/$KVER/$rel_path"
        mkdir -p "$(dirname "$dst")"
        cp "$mod_path" "$dst"
    done
    rm -f "$NEEDED_MODS"

    # Run depmod on the selective module set to generate correct metadata
    depmod -a -b "$ROOTFS" "$KVER" 2>/dev/null || true

    MOD_COUNT=$(find "$ROOTFS/lib/modules/$KVER" -name '*.ko*' | wc -l)
    MOD_SIZE=$(du -sh "$ROOTFS/lib/modules/$KVER" | cut -f1)
    echo "    Installed $MOD_COUNT kernel modules ($MOD_SIZE) — minimal set"

else
    # =========================================================================
    # Stock kernel path — extract TCZ modules + manual modules.dep
    # =========================================================================
    echo "    No custom kernel found, using stock kernel with TCZ modules..."

    if [ -n "$KVER" ]; then
        ROOTFS_MOD_DST="$ROOTFS/lib/modules/$KVER/kernel"

        NETFILTER_TCZ="$CACHE_DIR/ipv6-netfilter-${KVER}.tcz"
        if [ -f "$NETFILTER_TCZ" ]; then
            echo "    Extracting netfilter modules from $(basename "$NETFILTER_TCZ")..."
            TCZ_TMP=$(mktemp -d)
            if command -v unsquashfs >/dev/null 2>&1; then
                unsquashfs -d "$TCZ_TMP/content" "$NETFILTER_TCZ" >/dev/null 2>&1
            else
                echo "    ERROR: unsquashfs not found (install squashfs-tools)"
                rm -rf "$TCZ_TMP"
                exit 1
            fi
            TCZ_MOD_SRC="$TCZ_TMP/content/usr/local/lib/modules/$KVER/kernel"
            if [ -d "$TCZ_MOD_SRC" ]; then
                find "$TCZ_MOD_SRC" -name '*.ko.gz' | while IFS= read -r mod_file; do
                    rel_path="${mod_file#$TCZ_MOD_SRC/}"
                    dst_dir="$ROOTFS_MOD_DST/$(dirname "$rel_path")"
                    mkdir -p "$dst_dir"
                    cp "$mod_file" "$dst_dir/"
                done
                MOD_COUNT=$(find "$TCZ_MOD_SRC" -name '*.ko.gz' | wc -l)
                echo "    Installed $MOD_COUNT kernel modules from netfilter TCZ"
            fi
            rm -rf "$TCZ_TMP"
        else
            echo "    WARN: Netfilter TCZ not found. kube-proxy may not work."
        fi

        NET_BRIDGING_TCZ="$CACHE_DIR/net-bridging-${KVER}.tcz"
        if [ -f "$NET_BRIDGING_TCZ" ]; then
            echo "    Extracting bridge modules from $(basename "$NET_BRIDGING_TCZ")..."
            TCZ_TMP=$(mktemp -d)
            unsquashfs -d "$TCZ_TMP/content" "$NET_BRIDGING_TCZ" >/dev/null 2>&1
            TCZ_MOD_SRC="$TCZ_TMP/content/usr/local/lib/modules/$KVER/kernel"
            if [ -d "$TCZ_MOD_SRC" ]; then
                find "$TCZ_MOD_SRC" -name '*.ko.gz' | while IFS= read -r mod_file; do
                    rel_path="${mod_file#$TCZ_MOD_SRC/}"
                    dst_dir="$ROOTFS_MOD_DST/$(dirname "$rel_path")"
                    mkdir -p "$dst_dir"
                    cp "$mod_file" "$dst_dir/"
                done
                BR_COUNT=$(find "$TCZ_MOD_SRC" -name '*.ko.gz' | wc -l)
                echo "    Installed $BR_COUNT kernel modules from net-bridging TCZ"
            fi
            rm -rf "$TCZ_TMP"
        else
            echo "    WARN: Net-bridging TCZ not found. CNI bridge networking may not work."
        fi

        # Manual modules.dep for stock kernel (Ubuntu's depmod can't handle TC's kernel)
        MODULES_DEP="$ROOTFS/lib/modules/$KVER/modules.dep"
        if [ -f "$MODULES_DEP" ]; then
            echo "    Appending module entries to modules.dep..."
            cat >> "$MODULES_DEP" << 'MODDEP'
kernel/net/ipv6/ipv6.ko.gz:
kernel/net/ipv4/netfilter/nf_defrag_ipv4.ko.gz:
kernel/net/ipv6/netfilter/nf_defrag_ipv6.ko.gz: kernel/net/ipv6/ipv6.ko.gz
kernel/net/netfilter/nf_conntrack.ko.gz: kernel/net/ipv4/netfilter/nf_defrag_ipv4.ko.gz kernel/net/ipv6/netfilter/nf_defrag_ipv6.ko.gz
kernel/net/netfilter/nf_nat.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/nf_conntrack_netlink.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/nf_tables.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/nft_compat.ko.gz: kernel/net/netfilter/nf_tables.ko.gz
kernel/net/netfilter/nft_chain_nat.ko.gz: kernel/net/netfilter/nf_tables.ko.gz kernel/net/netfilter/nf_nat.ko.gz
kernel/net/netfilter/nft_ct.ko.gz: kernel/net/netfilter/nf_tables.ko.gz kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/nft_masq.ko.gz: kernel/net/netfilter/nf_tables.ko.gz kernel/net/netfilter/nf_nat.ko.gz kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/nft_nat.ko.gz: kernel/net/netfilter/nf_tables.ko.gz kernel/net/netfilter/nf_nat.ko.gz
kernel/net/netfilter/nft_redir.ko.gz: kernel/net/netfilter/nf_tables.ko.gz kernel/net/netfilter/nf_nat.ko.gz
kernel/net/netfilter/xt_conntrack.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/xt_MASQUERADE.ko.gz: kernel/net/netfilter/nf_nat.ko.gz kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/xt_mark.ko.gz:
kernel/net/netfilter/xt_comment.ko.gz:
kernel/net/netfilter/xt_multiport.ko.gz:
kernel/net/netfilter/xt_nat.ko.gz: kernel/net/netfilter/nf_nat.ko.gz
kernel/net/netfilter/xt_addrtype.ko.gz:
kernel/net/netfilter/xt_connmark.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz
kernel/net/netfilter/xt_REDIRECT.ko.gz: kernel/net/netfilter/nf_nat.ko.gz
kernel/net/netfilter/xt_recent.ko.gz:
kernel/net/netfilter/xt_statistic.ko.gz:
kernel/net/netfilter/xt_set.ko.gz: kernel/net/netfilter/ipset/ip_set.ko.gz
kernel/net/netfilter/ipset/ip_set.ko.gz:
kernel/net/ipv4/netfilter/nf_reject_ipv4.ko.gz:
kernel/net/ipv6/netfilter/nf_reject_ipv6.ko.gz:
kernel/net/ipv4/netfilter/ipt_REJECT.ko.gz: kernel/net/ipv4/netfilter/nf_reject_ipv4.ko.gz
kernel/net/ipv6/netfilter/ip6t_REJECT.ko.gz: kernel/net/ipv6/netfilter/nf_reject_ipv6.ko.gz
kernel/net/netfilter/nft_reject.ko.gz: kernel/net/netfilter/nf_tables.ko.gz
kernel/net/bridge/bridge.ko.gz: kernel/net/802/stp.ko.gz kernel/net/llc/llc.ko.gz
kernel/net/bridge/br_netfilter.ko.gz: kernel/net/bridge/bridge.ko.gz kernel/net/802/stp.ko.gz kernel/net/llc/llc.ko.gz
kernel/net/bridge/netfilter/nf_conntrack_bridge.ko.gz: kernel/net/netfilter/nf_conntrack.ko.gz kernel/net/bridge/bridge.ko.gz
MODDEP

            find "$ROOTFS_MOD_DST" -name '*.ko.gz' -path '*/net/*' | sort | while IFS= read -r mod_file; do
                rel_path="kernel/${mod_file#$ROOTFS_MOD_DST/}"
                if ! grep -q "^${rel_path}:" "$MODULES_DEP" 2>/dev/null; then
                    echo "${rel_path}:" >> "$MODULES_DEP"
                fi
            done
            echo "    Updated modules.dep with netfilter entries"
        fi
    fi
fi

# Install iptables-nft (nftables-based iptables) from the builder system
# Kernel 6.18 uses nf_tables, not legacy ip_tables, so we need xtables-nft-multi
echo "    Installing iptables-nft from builder..."
if [ -f /usr/sbin/xtables-nft-multi ]; then
    mkdir -p "$ROOTFS/usr/sbin"
    cp /usr/sbin/xtables-nft-multi "$ROOTFS/usr/sbin/"

    # Create standard symlinks
    for cmd in iptables iptables-save iptables-restore ip6tables ip6tables-save ip6tables-restore; do
        ln -sf xtables-nft-multi "$ROOTFS/usr/sbin/$cmd"
    done

    # Copy required shared libraries
    mkdir -p "$ROOTFS/usr/lib/x86_64-linux-gnu" "$ROOTFS/lib/x86_64-linux-gnu" "$ROOTFS/lib64"
    for lib in \
        /lib/x86_64-linux-gnu/libxtables.so.12* \
        /lib/x86_64-linux-gnu/libmnl.so.0* \
        /lib/x86_64-linux-gnu/libnftnl.so.11* \
        /lib/x86_64-linux-gnu/libc.so.6 \
        /lib64/ld-linux-x86-64.so.2; do
        [ -e "$lib" ] && cp -aL "$lib" "$ROOTFS${lib}" 2>/dev/null || true
    done

    # Copy xtables modules directory (match extensions)
    if [ -d /usr/lib/x86_64-linux-gnu/xtables ]; then
        mkdir -p "$ROOTFS/usr/lib/x86_64-linux-gnu/xtables"
        cp -a /usr/lib/x86_64-linux-gnu/xtables/*.so "$ROOTFS/usr/lib/x86_64-linux-gnu/xtables/" 2>/dev/null || true
    fi

    echo "    Installed iptables-nft (xtables-nft-multi) + shared libs"
else
    echo "    WARN: xtables-nft-multi not found in builder (install iptables package)"
fi

# Kernel modules list (for init to load at boot)
cp "$PROJECT_ROOT/build/config/modules.list" "$ROOTFS/usr/lib/kubesolo-os/modules.list"

# --- 4. Sysctl config ---
mkdir -p "$ROOTFS/etc/sysctl.d"
cp "$PROJECT_ROOT/build/rootfs/etc/sysctl.d/k8s.conf" "$ROOTFS/etc/sysctl.d/k8s.conf"

# --- 5. OS metadata ---
echo "$VERSION" > "$ROOTFS/etc/kubesolo-os-version"

cat > "$ROOTFS/etc/os-release" << EOF
NAME="KubeSolo OS"
VERSION="$VERSION"
ID=kubesolo-os
VERSION_ID=$VERSION
PRETTY_NAME="KubeSolo OS $VERSION"
HOME_URL="https://github.com/portainer/kubesolo"
BUG_REPORT_URL="https://github.com/portainer/kubesolo/issues"
EOF

# --- 6. Default KubeSolo config ---
mkdir -p "$ROOTFS/etc/kubesolo"
if [ -f "$PROJECT_ROOT/build/rootfs/etc/kubesolo/defaults.yaml" ]; then
    cp "$PROJECT_ROOT/build/rootfs/etc/kubesolo/defaults.yaml" "$ROOTFS/etc/kubesolo/defaults.yaml"
fi

# --- 7. Essential directories ---
mkdir -p "$ROOTFS/var/lib/kubesolo"
mkdir -p "$ROOTFS/var/lib/containerd"
mkdir -p "$ROOTFS/etc/kubesolo"
mkdir -p "$ROOTFS/etc/cni/net.d"
mkdir -p "$ROOTFS/opt/cni/bin"
mkdir -p "$ROOTFS/var/log"
mkdir -p "$ROOTFS/usr/local"
mkdir -p "$ROOTFS/mnt/data"
mkdir -p "$ROOTFS/run/containerd"

# --- 8. CA certificates (required for containerd to pull from registries) ---
mkdir -p "$ROOTFS/etc/ssl/certs"
if [ -f /etc/ssl/certs/ca-certificates.crt ]; then
    cp /etc/ssl/certs/ca-certificates.crt "$ROOTFS/etc/ssl/certs/ca-certificates.crt"
    echo "    Installed CA certificates bundle"
elif [ -f /etc/pki/tls/certs/ca-bundle.crt ]; then
    cp /etc/pki/tls/certs/ca-bundle.crt "$ROOTFS/etc/ssl/certs/ca-certificates.crt"
    echo "    Installed CA certificates bundle (from ca-bundle.crt)"
else
    echo "    WARN: No CA certificates found in builder — TLS verification will fail"
fi

# --- 9. Ensure /etc/hosts and /etc/resolv.conf exist ---
if [ ! -f "$ROOTFS/etc/hosts" ]; then
    cat > "$ROOTFS/etc/hosts" << EOF
127.0.0.1 localhost
::1       localhost
EOF
fi

if [ ! -f "$ROOTFS/etc/resolv.conf" ]; then
    cat > "$ROOTFS/etc/resolv.conf" << EOF
nameserver 8.8.8.8
nameserver 1.1.1.1
EOF
fi

# --- Summary ---
echo ""
echo "==> Injection complete. Rootfs contents:"
echo "    Total size: $(du -sh "$ROOTFS" | cut -f1)"
echo "    KubeSolo:   $(du -h "$ROOTFS/usr/bin/kubesolo" | cut -f1)"
echo "    Init stages: $(ls "$ROOTFS/usr/lib/kubesolo-os/init.d/" | wc -l)"
echo ""