#!/bin/bash
# Build a custom Fedora Asahi Remix rootfs with lab firstboot LVM setup.
#
# Downloads the upstream Fedora Asahi Remix Server package, injects our
# firstboot script + systemd service, and repackages it for the bastion.
#
# Requirements: root, curl, unzip, mount (loop), zip
# Output: bastion/asahi-repo/ directory with package + installer_data.json
#
# Usage: sudo ./scripts/build-asahi-rootfs.sh [--bastion-ip IP] [--http-port PORT]
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
ASAHI_DIR="$PROJECT_DIR/asahi-repo"
CACHE_DIR="$PROJECT_DIR/.asahi-cache"
WORK_DIR=""

# Defaults
BASTION_IP="${BASTION_IP:-192.168.8.23}"
HTTP_PORT="${HTTP_PORT:-8080}"
ROLE="${ROLE:-infra}"
HOSTNAME="${HOSTNAME:-mac-studio}"
MAC="${MAC:-00:00:00:00:00:00}"
ADMIN_USER="${ADMIN_USER:-michal}"

# Parse args
while [[ $# -gt 0 ]]; do
    case "$1" in
        --bastion-ip) BASTION_IP="$2"; shift 2 ;;
        --http-port) HTTP_PORT="$2"; shift 2 ;;
        --role) ROLE="$2"; shift 2 ;;
        --hostname) HOSTNAME="$2"; shift 2 ;;
        --mac) MAC="$2"; shift 2 ;;
        --admin-user) ADMIN_USER="$2"; shift 2 ;;
        *) echo "Unknown option: $1"; exit 1 ;;
    esac
done

# ── Resolve upstream package URL ─────────────────────────────────
echo "==> Fetching Asahi installer data..."
INSTALLER_DATA=$(curl -sfL "https://cdn.asahilinux.org/installer/installer_data.json")

# Find the Server variant package URL
SERVER_URL=$(echo "$INSTALLER_DATA" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for os in data.get('os_list', []):
    name = os.get('name', '').lower()
    if 'server' in name and 'uefi' not in name and not os.get('expert'):
        print(os['package'])
        break
" 2>/dev/null)

if [ -z "$SERVER_URL" ]; then
    echo "ERROR: Could not find Fedora Asahi Remix Server in installer data."
    echo "Available variants:"
    echo "$INSTALLER_DATA" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for os in data.get('os_list', []):
    print(f\"  - {os.get('name', '?')}\")" 2>/dev/null
    exit 1
fi

PACKAGE_NAME=$(basename "$SERVER_URL")
echo "    Variant: Fedora Asahi Remix Server"
echo "    Package: $PACKAGE_NAME"

# Also extract the partition layout and supported_fw from upstream
UPSTREAM_CONFIG=$(echo "$INSTALLER_DATA" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for os in data.get('os_list', []):
    name = os.get('name', '').lower()
    if 'server' in name and 'uefi' not in name and not os.get('expert'):
        json.dump(os, sys.stdout)
        break
")

# ── Download upstream package ────────────────────────────────────
mkdir -p "$CACHE_DIR" "$ASAHI_DIR"

CACHED_PKG="$CACHE_DIR/$PACKAGE_NAME"
if [ -f "$CACHED_PKG" ]; then
    echo "==> Using cached package: $CACHED_PKG"
else
    echo "==> Downloading $SERVER_URL..."
    curl -# -L -o "$CACHED_PKG" "$SERVER_URL"
fi

# ── Extract and modify rootfs ────────────────────────────────────
WORK_DIR=$(mktemp -d)
trap 'echo "==> Cleaning up..."; umount "$WORK_DIR/rootfs" 2>/dev/null || true; rm -rf "$WORK_DIR"' EXIT

echo "==> Extracting package..."
unzip -q -o "$CACHED_PKG" -d "$WORK_DIR/pkg"

# List contents
echo "    Package contents:"
ls -lh "$WORK_DIR/pkg/" | grep -v ^total | while read -r line; do echo "      $line"; done

# Find root.img
ROOT_IMG=$(find "$WORK_DIR/pkg" -name "root.img" -type f | head -1)
if [ -z "$ROOT_IMG" ]; then
    echo "ERROR: root.img not found in package."
    echo "Contents: $(ls "$WORK_DIR/pkg/")"
    exit 1
fi

echo "==> Mounting root.img..."
mkdir -p "$WORK_DIR/rootfs"
mount -o loop "$ROOT_IMG" "$WORK_DIR/rootfs"

# ── Read SSH keys from the system ────────────────────────────────
SSH_KEYS=""
REAL_USER="${SUDO_USER:-$USER}"
REAL_HOME=$(eval echo "~$REAL_USER")
for keyfile in "$REAL_HOME/.ssh/id_ed25519.pub" "$REAL_HOME/.ssh/id_ecdsa.pub" "$REAL_HOME/.ssh/id_rsa.pub"; do
    if [ -f "$keyfile" ]; then
        SSH_KEYS=$(cat "$keyfile")
        echo "    SSH key: $keyfile"
        break
    fi
done

if [ -z "$SSH_KEYS" ]; then
    echo "WARNING: No SSH public key found. You'll need to add keys manually."
fi

# ── Generate firstboot script from bastion ───────────────────────
echo "==> Generating firstboot script..."

# Try to get the script from a running bastion, fall back to local generation
FIRSTBOOT_SCRIPT=""
FIRSTBOOT_URL="http://$BASTION_IP:$HTTP_PORT/asahi/firstboot.sh?hostname=$HOSTNAME&role=$ROLE&mac=$MAC&user=$ADMIN_USER"
FIRSTBOOT_SCRIPT=$(curl -sf "$FIRSTBOOT_URL" 2>/dev/null || echo "")

if [ -z "$FIRSTBOOT_SCRIPT" ]; then
    echo "    Bastion not reachable, generating script locally..."
    # Generate a basic firstboot script inline
    FIRSTBOOT_SCRIPT=$(cd "$PROJECT_DIR" && node -e "
const { renderFirstbootScript } = require('./src/bastion/dist/templates/asahi-firstboot.sh.js');
process.stdout.write(renderFirstbootScript({
    hostname: '$HOSTNAME',
    role: '$ROLE',
    serverIp: '$BASTION_IP',
    httpPort: $HTTP_PORT,
    sshKeys: $([ -n "$SSH_KEYS" ] && echo "[\"$SSH_KEYS\"]" || echo "[]"),
    adminUser: '$ADMIN_USER',
    mac: '$MAC',
}));
" 2>/dev/null) || {
        echo "    ERROR: Could not generate firstboot script. Build the project first: npm run build"
        exit 1
    }
fi

# ── Inject files into rootfs ─────────────────────────────────────
echo "==> Injecting lab configuration into rootfs..."

# Firstboot script
echo "$FIRSTBOOT_SCRIPT" > "$WORK_DIR/rootfs/usr/local/bin/lab-firstboot.sh"
chmod 755 "$WORK_DIR/rootfs/usr/local/bin/lab-firstboot.sh"
echo "    Installed: /usr/local/bin/lab-firstboot.sh"

# Systemd service
cat > "$WORK_DIR/rootfs/etc/systemd/system/lab-firstboot.service" << 'UNIT'
[Unit]
Description=Lab first-boot LVM setup
After=local-fs.target network-online.target
Wants=network-online.target
ConditionPathExists=!/etc/lab-lvm-setup-done

[Service]
Type=oneshot
ExecStart=/usr/local/bin/lab-firstboot.sh
RemainAfterExit=yes
StandardOutput=journal+console
StandardError=journal+console

[Install]
WantedBy=multi-user.target
UNIT
echo "    Installed: /etc/systemd/system/lab-firstboot.service"

# Enable the service
mkdir -p "$WORK_DIR/rootfs/etc/systemd/system/multi-user.target.wants"
ln -sf /etc/systemd/system/lab-firstboot.service \
    "$WORK_DIR/rootfs/etc/systemd/system/multi-user.target.wants/lab-firstboot.service"
echo "    Enabled: lab-firstboot.service"

# SSH authorized keys for root (for initial access before firstboot runs user creation)
if [ -n "$SSH_KEYS" ]; then
    mkdir -p "$WORK_DIR/rootfs/root/.ssh"
    chmod 700 "$WORK_DIR/rootfs/root/.ssh"
    echo "$SSH_KEYS" > "$WORK_DIR/rootfs/root/.ssh/authorized_keys"
    chmod 600 "$WORK_DIR/rootfs/root/.ssh/authorized_keys"
    echo "    Installed: /root/.ssh/authorized_keys"
fi

# Ensure lvm2 and xfsprogs are installed (should be in server image already)
echo "    Checking required packages..."
if [ -f "$WORK_DIR/rootfs/usr/sbin/pvcreate" ]; then
    echo "      lvm2: present"
else
    echo "      WARNING: lvm2 not found in rootfs. LVM setup may fail."
fi
if [ -f "$WORK_DIR/rootfs/usr/sbin/mkfs.xfs" ]; then
    echo "      xfsprogs: present"
else
    echo "      WARNING: xfsprogs not found in rootfs. LVM setup may fail."
fi

# ── Unmount and repackage ────────────────────────────────────────
echo "==> Unmounting rootfs..."
umount "$WORK_DIR/rootfs"

echo "==> Repackaging..."
OUTPUT_PKG="$ASAHI_DIR/fedora-asahi-lab.zip"
rm -f "$OUTPUT_PKG"
(cd "$WORK_DIR/pkg" && zip -q "$OUTPUT_PKG" *)
echo "    Output: $OUTPUT_PKG ($(du -sh "$OUTPUT_PKG" | cut -f1))"

# ── Generate installer_data.json ─────────────────────────────────
echo "==> Generating installer_data.json..."

# Parse upstream config to get supported_fw, boot_object, next_object, and partition details
python3 << PYEOF > "$ASAHI_DIR/installer_data.json"
import json, sys

upstream = json.loads('''$UPSTREAM_CONFIG''')

# Build our custom installer data based on upstream
# Keep EFI and Boot partitions identical, modify Root to not expand,
# add Data partition that expands for LVM.
partitions = []
for p in upstream.get('partitions', []):
    if p.get('type') == 'EFI':
        partitions.append(p)
    elif p.get('name') == 'Boot':
        partitions.append(p)
    elif p.get('name') == 'Root':
        # Fixed size root, no expand
        root_p = dict(p)
        root_p['expand'] = False
        # Keep the original size (it's the minimum needed for the rootfs)
        partitions.append(root_p)

# Add Data partition for LVM
partitions.append({
    "name": "Data",
    "type": "Linux",
    "size": "1073741824B",  # 1GB minimum, will expand
    "expand": True
})

data = {
    "os_list": [{
        "name": "Fedora Asahi Lab (${ROLE})",
        "default_os_name": "Fedora Linux Lab",
        "boot_object": upstream.get("boot_object", "m1n1.bin"),
        "next_object": upstream.get("next_object", "m1n1/boot.bin"),
        "package": "fedora-asahi-lab.zip",
        "supported_fw": upstream.get("supported_fw", ["13.5"]),
        "partitions": partitions,
    }]
}

json.dump(data, sys.stdout, indent=2)
print()
PYEOF

echo "    Generated: $ASAHI_DIR/installer_data.json"

# Pretty-print the partition layout
echo ""
echo "    Partition layout:"
python3 -c "
import json
with open('$ASAHI_DIR/installer_data.json') as f:
    data = json.load(f)
for p in data['os_list'][0]['partitions']:
    size = p.get('size', '?')
    expand = ' (expand)' if p.get('expand') else ''
    image = f\" [{p['image']}]\" if 'image' in p else ''
    print(f\"      {p['name']:8s} {p['type']:8s} {size:>16s}{expand}{image}\")
"

echo ""
echo "==> Build complete!"
echo ""
echo "    Package: $ASAHI_DIR/fedora-asahi-lab.zip"
echo "    Config:  $ASAHI_DIR/installer_data.json"
echo ""
echo "    To serve from bastion, copy to the bastion's HTTP directory"
echo "    or configure REPO_BASE to point here."
echo ""
echo "    To install on Mac Studio:"
echo "      curl http://$BASTION_IP:$HTTP_PORT/asahi | sh"