#!/usr/bin/env bash
# Install as /var/lib/dietpi/postboot.d/10-pikit-firstboot and chmod +x.
# Runs once on first boot to finalize device-unique setup.

set -euo pipefail

FIRSTBOOT_DIR="/var/lib/pikit/firstboot"
STATE_FILE="$FIRSTBOOT_DIR/state.json"
LOG_FILE="$FIRSTBOOT_DIR/firstboot.log"
ERROR_FILE="$FIRSTBOOT_DIR/firstboot.error"
DONE_FILE="$FIRSTBOOT_DIR/firstboot.done"
LOCK_FILE="$FIRSTBOOT_DIR/firstboot.lock"
CERT_DIR="/etc/pikit/certs"
WEB_ASSETS="/var/www/pikit-web/assets"
PROFILE_FILE="/etc/pikit/profile.json"
MOTD_FILE="/etc/motd"
FIRSTBOOT_CONF="/etc/pikit/firstboot.conf"
APT_UA_OVERRIDE="/etc/apt/apt.conf.d/51pikit-unattended.conf"
SERVICE_JSON="/etc/pikit/services.json"

STEPS=(
  "Preparing system"
  "Generating security keys"
  "Securing the dashboard"
  "Updating software (this can take a while)"
  "Final checks"
  "Starting Pi-Kit"
)
STEP_STATUS=(pending pending pending pending pending pending)
CURRENT_STEP=""
CURRENT_INDEX=-1
PIKIT_FIRSTBOOT_UPDATES="${PIKIT_FIRSTBOOT_UPDATES:-1}"

log() {
  printf '[%s] %s\n' "$(date '+%Y-%m-%dT%H:%M:%S%z')" "$*"
}

load_config() {
  if [ -f "$FIRSTBOOT_CONF" ]; then
    # shellcheck disable=SC1090
    . "$FIRSTBOOT_CONF"
  fi
  PIKIT_FIRSTBOOT_UPDATES="${PIKIT_FIRSTBOOT_UPDATES:-1}"
}

skip_updates() {
  case "${PIKIT_FIRSTBOOT_UPDATES,,}" in
    0|false|no|off) return 0 ;;
  esac
  return 1
}

configure_unattended_defaults() {
  if [ -f "$APT_UA_OVERRIDE" ]; then
    log "Unattended-upgrades config already present; skipping defaults."
    return
  fi
  if ! command -v python3 >/dev/null 2>&1; then
    log "python3 missing; skipping unattended-upgrades defaults."
    return
  fi
  PROFILE_FILE="$PROFILE_FILE" SERVICE_JSON="$SERVICE_JSON" PYTHONPATH=/usr/local/bin python3 - <<'PY'
import sys
try:
    from pikit_api.auto_updates import set_updates_config
except Exception as e:
    print(f"pikit_api unavailable: {e}")
    sys.exit(0)

set_updates_config({"enable": True, "scope": "security"})
PY
  log "Unattended-upgrades defaults applied (security-only)."
}

apply_profile() {
  if [ ! -f "$PROFILE_FILE" ]; then
    log "Profile step skipped (no profile.json)."
    return
  fi
  PYTHONPATH=/usr/local/bin python3 - <<'PY'
import json
import os
import pathlib
import pwd
import grp
import shutil
import subprocess

profile_path = pathlib.Path(os.environ.get("PROFILE_FILE", "/etc/pikit/profile.json"))
services_path = pathlib.Path(os.environ.get("SERVICE_JSON", "/etc/pikit/services.json"))

def log(msg: str) -> None:
    print(msg)

def ipv6_enabled() -> bool:
    cfg = pathlib.Path("/etc/default/ufw")
    if not cfg.exists():
        return True
    for line in cfg.read_text().splitlines():
        if line.strip().startswith("IPV6="):
            return line.split("=", 1)[1].strip().lower() == "yes"
    return True

def get_ipv6_prefixes() -> list:
    if not ipv6_enabled():
        return []
    try:
        out = subprocess.check_output(["ip", "-6", "addr", "show", "scope", "global"], text=True)
    except Exception:
        return []
    prefixes = set()
    for line in out.splitlines():
        line = line.strip()
        if not line.startswith("inet6 "):
            continue
        if " temporary " in f" {line} ":
            continue
        parts = line.split()
        if len(parts) < 2:
            continue
        prefixes.add(parts[1])
    return sorted(prefixes)

try:
    profile = json.loads(profile_path.read_text())
except Exception as e:
    log(f"Profile load failed: {e}")
    profile = {}

ports = profile.get("firewall_ports") or []
firewall_enable = bool(profile.get("firewall_enable", False))
base_ports = profile.get("firewall_base_ports") or [22, 80, 443, 5252, 5253]
if ports:
    if shutil.which("ufw"):
        ipv6_prefixes = get_ipv6_prefixes()
        if ipv6_prefixes:
            log(f"IPv6 LAN prefixes: {', '.join(ipv6_prefixes)}")
        for raw in ports:
            try:
                port = int(raw)
            except Exception:
                continue
            for subnet in ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10", "169.254.0.0/16"):
                subprocess.run(["ufw", "allow", "from", subnet, "to", "any", "port", str(port)], check=False)
            for prefix in ipv6_prefixes:
                subprocess.run(["ufw", "allow", "from", prefix, "to", "any", "port", str(port)], check=False)
        if firewall_enable:
            for raw in base_ports:
                try:
                    port = int(raw)
                except Exception:
                    continue
                for subnet in ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10", "169.254.0.0/16"):
                    subprocess.run(["ufw", "allow", "from", subnet, "to", "any", "port", str(port)], check=False)
                for prefix in ipv6_prefixes:
                    subprocess.run(["ufw", "allow", "from", prefix, "to", "any", "port", str(port)], check=False)
            subprocess.run(["ufw", "--force", "enable"], check=False)
            log("UFW enabled.")
        log("Profile firewall rules applied.")
    else:
        log("Profile firewall step skipped (ufw missing).")
else:
    log("Profile firewall step skipped (no ports).")

def normalize_name(name: str) -> str:
    return name.strip().lower()

services = []
if services_path.exists():
    try:
        services = json.loads(services_path.read_text())
    except Exception:
        services = []
if not isinstance(services, list):
    services = []

profile_services = profile.get("services") or []
if isinstance(profile_services, list) and profile_services:
    for psvc in profile_services:
        if not isinstance(psvc, dict):
            continue
        p_name = normalize_name(str(psvc.get("name", "")))
        p_port = str(psvc.get("port", ""))
        p_path = str(psvc.get("path", "")) if psvc.get("path") is not None else ""
        replaced = False
        for svc in services:
            if not isinstance(svc, dict):
                continue
            s_name = normalize_name(str(svc.get("name", "")))
            s_port = str(svc.get("port", ""))
            s_path = str(svc.get("path", "")) if svc.get("path") is not None else ""
            if (p_name and s_name == p_name) or (p_port and s_port == p_port and s_path == p_path):
                svc.update(psvc)
                replaced = True
                break
        if not replaced:
            services.append(psvc)
    services_path.parent.mkdir(parents=True, exist_ok=True)
    services_path.write_text(json.dumps(services, indent=2))
    log("Profile services merged.")
else:
    log("Profile services step skipped (none).")

actions = profile.get("actions") or []
if isinstance(actions, list) and actions:
    for action in actions:
        if not isinstance(action, dict):
            continue
        action_type = action.get("type")
        if action_type == "tls_bundle":
            src_cert = pathlib.Path(action.get("source_cert", ""))
            src_key = pathlib.Path(action.get("source_key", ""))
            dest = pathlib.Path(action.get("dest", ""))
            if not src_cert.exists() or not src_key.exists():
                log(f"TLS bundle skipped (missing cert/key): {dest}")
                continue
            dest.parent.mkdir(parents=True, exist_ok=True)
            content = src_cert.read_bytes() + b"\n" + src_key.read_bytes() + b"\n"
            dest.write_bytes(content)
            owner = action.get("owner")
            if owner:
                user, _, group = str(owner).partition(":")
                try:
                    uid = pwd.getpwnam(user).pw_uid if user else -1
                except Exception:
                    uid = -1
                try:
                    gid = grp.getgrnam(group).gr_gid if group else -1
                except Exception:
                    gid = -1
                if uid != -1 or gid != -1:
                    os.chown(dest, uid if uid != -1 else -1, gid if gid != -1 else -1)
            mode = action.get("mode")
            if mode:
                try:
                    os.chmod(dest, int(str(mode), 8))
                except Exception:
                    pass
            restart = action.get("restart")
            if restart:
                subprocess.run(["systemctl", "restart", str(restart)], check=False)
            log(f"TLS bundle written: {dest}")
            continue

        if action_type == "replace_text":
            file_path = pathlib.Path(action.get("file", ""))
            match = str(action.get("match", ""))
            replacement = str(action.get("replace", ""))
            if not file_path.exists():
                log(f"Replace skipped (missing file): {file_path}")
                continue
            content = file_path.read_text()
            if match not in content:
                log(f"Replace skipped (pattern not found): {file_path}")
                continue
            file_path.write_text(content.replace(match, replacement, 1))
            restart = action.get("restart")
            if restart:
                subprocess.run(["systemctl", "restart", str(restart)], check=False)
            log(f"Replaced text in: {file_path}")
            continue
else:
    log("Profile actions step skipped (none).")
PY
}

write_state() {
  local state="$1"
  local current="$2"
  local steps_joined
  local status_joined
  steps_joined=$(IFS='|'; echo "${STEPS[*]}")
  status_joined=$(IFS='|'; echo "${STEP_STATUS[*]}")
  PIKIT_STATE_FILE="$STATE_FILE" PIKIT_STATE="$state" PIKIT_CURRENT_STEP="$current" PIKIT_STEPS="$steps_joined" PIKIT_STEP_STATUSES="$status_joined" \
    python3 - <<'PY'
import json
import os
import pathlib
from datetime import datetime, timezone

state_path = pathlib.Path(os.environ["PIKIT_STATE_FILE"]) if "PIKIT_STATE_FILE" in os.environ else pathlib.Path("/var/lib/pikit/firstboot/state.json")
state = os.environ.get("PIKIT_STATE", "running")
current = os.environ.get("PIKIT_CURRENT_STEP") or None
steps = (os.environ.get("PIKIT_STEPS") or "").split("|")
statuses = (os.environ.get("PIKIT_STEP_STATUSES") or "").split("|")
if len(statuses) < len(steps):
    statuses += ["pending"] * (len(steps) - len(statuses))
updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
state_path.parent.mkdir(parents=True, exist_ok=True)
state_path.write_text(json.dumps({
    "state": state,
    "current_step": current,
    "steps": [{"label": label, "status": status} for label, status in zip(steps, statuses)],
    "updated_at": updated_at,
}, indent=2))
PY
}

begin_step() {
  local idx="$1"
  if [ "$CURRENT_INDEX" -ge 0 ]; then
    STEP_STATUS[$CURRENT_INDEX]="done"
  fi
  CURRENT_INDEX="$idx"
  CURRENT_STEP="${STEPS[$idx]}"
  STEP_STATUS[$idx]="current"
  write_state "running" "$CURRENT_STEP"
  log "--- $CURRENT_STEP ---"
}

finish_step() {
  local idx="$1"
  local state="${2:-running}"
  local current="${3:-$CURRENT_STEP}"
  STEP_STATUS[$idx]="done"
  write_state "$state" "$current"
}

clear_motd_block() {
  if [ -f "$MOTD_FILE" ]; then
    sed -i '/^\[Pi-Kit firstboot\]/,/^\[\/Pi-Kit firstboot\]/d' "$MOTD_FILE" || true
  fi
}

write_motd_error() {
  clear_motd_block
  cat >> "$MOTD_FILE" <<'TXT'
[Pi-Kit firstboot]
Pi-Kit setup needs attention.
Error log: sudo cat /var/lib/pikit/firstboot/firstboot.error
Full log:  sudo cat /var/lib/pikit/firstboot/firstboot.log
If needed: sudo systemctl restart nginx pikit-api
[/Pi-Kit firstboot]
TXT
}

handle_error() {
  local line="$1"
  local msg="Firstboot failed at step: ${CURRENT_STEP:-unknown} (line $line)"
  log "$msg"
  STEP_STATUS[$CURRENT_INDEX]="error"
  write_state "error" "${CURRENT_STEP:-}" || true
  echo "$msg" > "$ERROR_FILE"
  if [ -f "$LOG_FILE" ]; then
    echo "--- recent log ---" >> "$ERROR_FILE"
    tail -n 120 "$LOG_FILE" >> "$ERROR_FILE"
  fi
  write_motd_error
  exit 1
}

mkdir -p "$FIRSTBOOT_DIR"
:> "$LOG_FILE"
exec >>"$LOG_FILE" 2>&1

log "Pi-Kit firstboot starting"
load_config

if [ -f "$DONE_FILE" ]; then
  log "Firstboot already completed; exiting."
  exit 0
fi

rm -f "$ERROR_FILE"
clear_motd_block

exec 9>"$LOCK_FILE"
if ! flock -n 9; then
  log "Another firstboot run is in progress; exiting."
  exit 0
fi

trap 'handle_error $LINENO' ERR

begin_step 0
mkdir -p "$CERT_DIR" "$WEB_ASSETS"
if getent group pikit-cert >/dev/null 2>&1; then
  chgrp pikit-cert "$CERT_DIR" || true
fi
chmod 750 "$CERT_DIR"
finish_step 0

begin_step 1
if [ -x /usr/local/bin/pikit-certgen.sh ]; then
  /usr/local/bin/pikit-certgen.sh
else
  if [ -s "$CERT_DIR/pikit-ca.crt" ] && [ -s "$CERT_DIR/pikit-ca.key" ] && [ -s "$CERT_DIR/pikit.local.crt" ] && [ -s "$CERT_DIR/pikit.local.key" ]; then
    log "TLS certs already present; skipping generation."
  else
    if ! command -v openssl >/dev/null 2>&1; then
      echo "openssl not installed" >&2
      exit 1
    fi
    rm -f "$CERT_DIR/pikit-ca.key" "$CERT_DIR/pikit-ca.crt" "$CERT_DIR/pikit-ca.srl" || true
    rm -f "$CERT_DIR/pikit.local.key" "$CERT_DIR/pikit.local.crt" "$CERT_DIR/pikit.local.csr" || true

    openssl genrsa -out "$CERT_DIR/pikit-ca.key" 2048
    openssl req -x509 -new -nodes -key "$CERT_DIR/pikit-ca.key" -sha256 -days 3650 \
      -out "$CERT_DIR/pikit-ca.crt" -subj "/CN=Pi-Kit CA"

    openssl genrsa -out "$CERT_DIR/pikit.local.key" 2048
    openssl req -new -key "$CERT_DIR/pikit.local.key" -out "$CERT_DIR/pikit.local.csr" -subj "/CN=pikit.local"

    SAN_CFG=$(mktemp)
    cat > "$SAN_CFG" <<'CFG'
authorityKeyIdentifier=keyid,issuer
basicConstraints=CA:FALSE
keyUsage = digitalSignature, nonRepudiation, keyEncipherment, dataEncipherment
subjectAltName = @alt_names

[alt_names]
DNS.1 = pikit.local
DNS.2 = pikit
CFG

    openssl x509 -req -in "$CERT_DIR/pikit.local.csr" -CA "$CERT_DIR/pikit-ca.crt" -CAkey "$CERT_DIR/pikit-ca.key" \
      -CAcreateserial -out "$CERT_DIR/pikit.local.crt" -days 825 -sha256 -extfile "$SAN_CFG"
    rm -f "$SAN_CFG" "$CERT_DIR/pikit.local.csr"

    chmod 600 "$CERT_DIR/pikit-ca.key" "$CERT_DIR/pikit.local.key"
    chmod 644 "$CERT_DIR/pikit-ca.crt" "$CERT_DIR/pikit.local.crt"
  fi
fi
finish_step 1

begin_step 2
cp "$CERT_DIR/pikit-ca.crt" "$WEB_ASSETS/pikit-ca.crt"
chmod 644 "$WEB_ASSETS/pikit-ca.crt"
if command -v sha256sum >/dev/null 2>&1; then
  sha256sum "$WEB_ASSETS/pikit-ca.crt" | awk '{print $1}' > "$WEB_ASSETS/pikit-ca.sha256"
elif command -v openssl >/dev/null 2>&1; then
  openssl dgst -sha256 "$WEB_ASSETS/pikit-ca.crt" | awk '{print $2}' > "$WEB_ASSETS/pikit-ca.sha256"
fi
if [ -s "$WEB_ASSETS/pikit-ca.sha256" ]; then
  chmod 644 "$WEB_ASSETS/pikit-ca.sha256"
fi
if command -v systemctl >/dev/null 2>&1; then
  systemctl reload nginx || systemctl restart nginx
fi
finish_step 2

begin_step 3
if skip_updates; then
  log "Skipping software updates (PIKIT_FIRSTBOOT_UPDATES=$PIKIT_FIRSTBOOT_UPDATES)."
else
  export DEBIAN_FRONTEND=noninteractive
  mkdir -p /var/cache/apt/archives/partial /var/lib/apt/lists/partial
  chmod 755 /var/cache/apt/archives /var/cache/apt/archives/partial /var/lib/apt/lists /var/lib/apt/lists/partial
  apt-get update
  apt-get -y -o Dpkg::Options::=--force-confold full-upgrade
fi
finish_step 3

begin_step 4
configure_unattended_defaults
apply_profile

if [ ! -f "$WEB_ASSETS/pikit-ca.crt" ]; then
  echo "CA bundle missing in web assets" >&2
  exit 1
fi
finish_step 4

begin_step 5
touch "$DONE_FILE"
touch /var/run/pikit-ready
finish_step 5 "done" "${STEPS[5]}"

log "Pi-Kit firstboot complete"
exit 0