#!/bin/bash

# =============================================================================
# recomprimir — Recomprime una ISO híbrida live con xz + bcj x86
# Autor: Charlie Martinez, GPLv2
# Uso: sudo ./recomprimir <fichero.iso>
# Pensado para ISOs generadas con penguin-eggs (Quirinux GNU/Linux)
# =============================================================================

# -----------------------------------------------------------------------------
# Directorio de trabajo
# Cámbialo si lo prefieres fijo, por ejemplo:
#   WORKBASE="/home/charlie/recomprimir-work"
# Por defecto elige automáticamente la partición con más espacio libre.
# -----------------------------------------------------------------------------
WORKBASE=""

detectar_workbase() {
    WORKBASE=$(df --output=avail,target -x tmpfs -x devtmpfs -x squashfs -x overlay \
        | sort -rn | awk 'NR==2{print $2}')
    [[ -z "$WORKBASE" || "$WORKBASE" == "/" ]] && WORKBASE="/var/tmp"
    WORKBASE="${WORKBASE}/recomprimir-work"
}

# -----------------------------------------------------------------------------
# Colores y funciones
# -----------------------------------------------------------------------------
RED='\033[0;31m'   GREEN='\033[0;32m'
YELLOW='\033[1;33m' CYAN='\033[0;36m'
BOLD='\033[1m'      NC='\033[0m'

info()  { echo -e "${CYAN}[INFO]${NC}  $*"; }
ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
die()   { error "$*"; exit 1; }

# -----------------------------------------------------------------------------
# Comprobaciones previas
# -----------------------------------------------------------------------------
[[ $EUID -ne 0 ]] && die "Ejecuta el script como root: sudo $0 <fichero.iso>"

[[ $# -lt 1 ]] && {
    echo -e "${BOLD}Uso:${NC} sudo $0 <fichero.iso>"
    exit 1
}

ISO_ORIGINAL="$(realpath "$1")"
[[ ! -f "$ISO_ORIGINAL" ]] && die "No se encuentra el fichero: $ISO_ORIGINAL"
[[ "${ISO_ORIGINAL,,}" != *.iso ]] && die "El fichero no tiene extensión .iso"

for cmd in xorriso unsquashfs mksquashfs mount umount dd chroot; do
    command -v "$cmd" &>/dev/null \
        || die "Falta: '$cmd'  →  apt install $cmd"
done

# -----------------------------------------------------------------------------
# Variables
# -----------------------------------------------------------------------------
ISO_BASENAME="$(basename "$ISO_ORIGINAL" .iso)"
ISO_DIR="$(dirname "$ISO_ORIGINAL")"
ISO_SALIDA="${ISO_DIR}/${ISO_BASENAME}-xz.iso"

[[ -z "$WORKBASE" ]] && detectar_workbase

WORKDIR="${WORKBASE}"
DIR_ISO="${WORKDIR}/iso"
DIR_SQUASHFS="${WORKDIR}/squashfs"
DIR_NEWISO="${WORKDIR}/newiso"
MBR_BIN="${WORKDIR}/mbr.bin"

# Comprobar espacio (necesitamos ~3× el tamaño de la ISO)
ISO_SIZE_KB=$(du -sk "${ISO_ORIGINAL}" | cut -f1)
ESPACIO_DISPONIBLE_KB=$(df -k "$(dirname "${WORKBASE}")" | awk 'NR==2{print $4}')
if [[ $ESPACIO_DISPONIBLE_KB -lt $(( ISO_SIZE_KB * 3 )) ]]; then
    die "Espacio insuficiente en $(dirname "${WORKBASE}").
  Disponible : $(( ESPACIO_DISPONIBLE_KB / 1024 )) MB
  Necesario  : ~$(( ISO_SIZE_KB * 3 / 1024 )) MB
  Ajusta la variable WORKBASE al inicio del script."
fi

# Limpieza solo si hay error
cleanup() {
    desmontar_chroot 2>/dev/null || true
    mountpoint -q "${DIR_ISO}" 2>/dev/null && umount -lf "${DIR_ISO}" 2>/dev/null || true
    [[ -f "${ISO_SALIDA}.tmp" ]] && rm -f "${ISO_SALIDA}.tmp" || true
    [[ -d "${WORKDIR}" ]] && rm -rf "${WORKDIR}" || true
}
trap cleanup ERR

# -----------------------------------------------------------------------------
# Funciones para montar/desmontar el entorno chroot
# -----------------------------------------------------------------------------
montar_chroot() {
    info "Montando entorno chroot en: ${DIR_SQUASHFS}"

    # Verificar que el directorio raíz del squashfs existe y tiene contenido
    [[ ! -d "${DIR_SQUASHFS}" ]] \
        && die "El directorio squashfs no existe: ${DIR_SQUASHFS}"
    [[ ! -d "${DIR_SQUASHFS}/usr" ]] \
        && die "El squashfs parece vacío o incompleto: ${DIR_SQUASHFS}"

    # Detectar el shell disponible dentro del squashfs
    CHROOT_SHELL=""
    for sh in /bin/bash /usr/bin/bash /bin/sh /usr/bin/sh; do
        [[ -f "${DIR_SQUASHFS}${sh}" ]] && { CHROOT_SHELL="$sh"; break; }
    done
    [[ -z "$CHROOT_SHELL" ]] \
        && die "No se encontró ningún shell (/bin/bash, /bin/sh) dentro del squashfs."
    info "Shell encontrado en la ISO: ${CHROOT_SHELL}"

    # Montar solo si no están ya montados (reentrada segura)
    for punto in dev dev/pts proc sys run; do
        mountpoint -q "${DIR_SQUASHFS}/${punto}" 2>/dev/null \
            || mount --bind "/${punto}" "${DIR_SQUASHFS}/${punto}"
    done

    # Soporte para aplicaciones gráficas (X11)
    if [[ -n "${DISPLAY:-}" && -d /tmp/.X11-unix ]]; then
        mkdir -p "${DIR_SQUASHFS}/tmp/.X11-unix"
        mountpoint -q "${DIR_SQUASHFS}/tmp/.X11-unix" 2>/dev/null \
            || mount --bind /tmp/.X11-unix "${DIR_SQUASHFS}/tmp/.X11-unix"
        xhost +local: &>/dev/null || true
        info "Soporte gráfico X11 activado (DISPLAY=${DISPLAY})"
    else
        warn "No se detectó sesión X11 — las aplicaciones gráficas no estarán disponibles"
    fi

    ok "Entorno chroot montado"
}

desmontar_chroot() {
    # Desmontar en orden inverso los puntos conocidos
    for punto in run sys proc dev/pts dev; do
        mountpoint -q "${DIR_SQUASHFS}/${punto}" 2>/dev/null \
            && umount -lf "${DIR_SQUASHFS}/${punto}" 2>/dev/null || true
    done
    # Desmontar socket X11 si estaba montado
    mountpoint -q "${DIR_SQUASHFS}/tmp/.X11-unix" 2>/dev/null \
        && umount -lf "${DIR_SQUASHFS}/tmp/.X11-unix" 2>/dev/null || true
    # Desmontar cualquier montaje residual que haya quedado dentro del squashfs
    # (gvfs, fuse, doc/by-app, etc.) — orden inverso por profundidad de ruta
    grep "${DIR_SQUASHFS}" /proc/mounts 2>/dev/null \
        | awk '{print $2}' \
        | sort -r \
        | while read -r punto_residual; do
            umount -lf "${punto_residual}" 2>/dev/null || true
        done
}

# -----------------------------------------------------------------------------
# Menú interactivo
# -----------------------------------------------------------------------------
menu_interactivo() {
    while true; do
        echo
        echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
        echo -e "${BOLD}  MENÚ DE EDICIÓN                                       ${NC}"
        echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
        echo -e "${YELLOW}  ⚠  No cierres esta ventana hasta que se haya${NC}"
        echo -e "${YELLOW}     finalizado el procedimiento.${NC}"
        echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
        echo
        echo -e "  ${BOLD}1)${NC} Editar la ISO  ${CYAN}(chroot)${NC}"
        echo -e "  ${BOLD}2)${NC} Recomprimir"
        echo -e "  ${BOLD}3)${NC} Salir"
        echo
        read -r -p "  Elige una opción [1-3]: " OPCION </dev/tty

        case "$OPCION" in

            1)
                echo
                echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
                echo -e "  Ahora estás dentro de tu ISO, puedes modificarla,"
                echo -e "  actualizarla, instalar y desinstalar programas, etc."
                echo -e "  Escribe ${CYAN}'salir'${NC} para regresar al menú."
                echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
                echo
                montar_chroot
                echo
                echo -e "${YELLOW}  Estás dentro de la ISO. El prompt cambiará a:${NC}"
                echo -e "${BOLD}  root@ISO-QUIRINUX:/#${NC}"
                echo -e "${YELLOW}  Si ves el hostname de tu máquina real, algo falló.${NC}"                echo
                # Lanzar chroot con hostname visual distintivo y shell detectado.
                # El usuario escribe 'salir' para volver al menú (en lugar de exit,
                # que propagaría la señal y mataría el script).
                RCFILE_TMP="${DIR_SQUASHFS}/tmp/.chroot_rc_$$"
                SIGNAL_FILE="${DIR_SQUASHFS}/tmp/.chroot_done_$$"
                rm -f "${SIGNAL_FILE}"
                cat > "${RCFILE_TMP}" << RCEOF
export PS1='\[\033[1;31m\]root@ISO-QUIRINUX\[\033[0m\]:\[\033[1;34m\]\w\[\033[0m\]# '
[ -f /etc/bash.bashrc ] && source /etc/bash.bashrc 2>/dev/null || true
[ -f ~/.bashrc ]        && source ~/.bashrc         2>/dev/null || true
export PS1='\[\033[1;31m\]root@ISO-QUIRINUX\[\033[0m\]:\[\033[1;34m\]\w\[\033[0m\]# '
alias salir='touch /tmp/.chroot_done_$$ && exit'
echo ""
echo "  Escribe 'salir' para volver al menú."
echo ""
RCEOF
                set +e
                (
                    chroot "${DIR_SQUASHFS}" \
                        /usr/bin/env -i \
                        HOME=/root \
                        TERM="${TERM}" \
                        DISPLAY="${DISPLAY:-:0}" \
                        XAUTHORITY="${XAUTHORITY:-/root/.Xauthority}" \
                        PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \
                        "${CHROOT_SHELL}" --rcfile /tmp/.chroot_rc_$$
                )
                set -e
                rm -f "${RCFILE_TMP}"
                # Al volver, desmontamos
                desmontar_chroot
                echo
                ok "Has salido del chroot. De vuelta al menú."
                echo -e "  ${CYAN}Hostname actual: $(hostname)${NC}  ← debe ser tu máquina host"
                ;;

            2)
                # Asegurarse de que chroot esté desmontado antes de recomprimir
                desmontar_chroot
                echo
                ok "Continuando con la recompresión..."
                return 0   # Salir del menú y seguir el script
                ;;

            3)
                echo
                warn "Saliendo y eliminando los archivos descomprimidos..."
                desmontar_chroot
                mountpoint -q "${DIR_ISO}" 2>/dev/null \
                    && umount -lf "${DIR_ISO}" 2>/dev/null || true
                rm -rf "${WORKDIR}"
                ok "Archivos eliminados: ${WORKDIR}"
                echo
                info "Operación cancelada por el usuario. ¡Hasta luego!"
                echo
                exit 0
                ;;

            *)
                warn "Opción no válida. Elige 1, 2 o 3."
                ;;
        esac
    done
}

# -----------------------------------------------------------------------------
# Encabezado
# -----------------------------------------------------------------------------
echo
echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
echo -e "${BOLD}  recomprimir — xz + bcj x86 · Quirinux                ${NC}"
echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
echo
info "ISO original   : $ISO_ORIGINAL"
info "ISO resultado  : $ISO_SALIDA"
info "Dir. de trabajo: $WORKDIR"
info "Espacio libre  : $(( ESPACIO_DISPONIBLE_KB / 1024 )) MB"
echo

# -----------------------------------------------------------------------------
# PASO 1 — Preservar el MBR híbrido (primeros 432 bytes de la ISO original)
# -----------------------------------------------------------------------------
info "Paso 1/6 — Preservando MBR híbrido (432 bytes)..."
mkdir -p "${WORKDIR}"
dd if="${ISO_ORIGINAL}" of="${MBR_BIN}" bs=1 count=432 status=none
ok "MBR guardado"

# -----------------------------------------------------------------------------
# PASO 2 — Montar la ISO y copiar su estructura
# -----------------------------------------------------------------------------
info "Paso 2/6 — Copiando estructura de la ISO..."
mkdir -p "${DIR_ISO}" "${DIR_SQUASHFS}" "${DIR_NEWISO}"
mount -o loop,ro "${ISO_ORIGINAL}" "${DIR_ISO}"
cp -a "${DIR_ISO}/." "${DIR_NEWISO}/"

SQUASHFS_RELPATH=$(find "${DIR_NEWISO}" -name "filesystem.squashfs" -printf '%P\n' | head -1)
[[ -z "$SQUASHFS_RELPATH" ]] && die "No se encontró filesystem.squashfs en la ISO."
info "squashfs en: ${SQUASHFS_RELPATH}"

chmod -R u+w "${DIR_NEWISO}"
rm -f "${DIR_NEWISO}/${SQUASHFS_RELPATH}"
ok "Estructura copiada"

# -----------------------------------------------------------------------------
# PASO 3 — Descomprimir el squashfs
# -----------------------------------------------------------------------------
info "Paso 3/6 — Descomprimiendo filesystem.squashfs..."
unsquashfs -xattrs -d "${DIR_SQUASHFS}" "${DIR_ISO}/${SQUASHFS_RELPATH}"
umount -lf "${DIR_ISO}"
ok "Squashfs descomprimido"

# -----------------------------------------------------------------------------
# MENÚ INTERACTIVO — Editar, Recomprimir o Salir
# -----------------------------------------------------------------------------
menu_interactivo

# -----------------------------------------------------------------------------
# PASO 4 — Recomprimir con xz + bcj x86
# -----------------------------------------------------------------------------
SQUASHFS_NUEVO="${DIR_NEWISO}/${SQUASHFS_RELPATH}"
mkdir -p "$(dirname "${SQUASHFS_NUEVO}")"

info "Paso 4/6 — Recomprimiendo con xz + bcj x86..."
warn "Este paso puede tardar entre 30 y 90 minutos."
echo

SECONDS=0
mksquashfs "${DIR_SQUASHFS}" "${SQUASHFS_NUEVO}" \
    -comp xz -Xbcj x86 -b 1M -Xdict-size 1M \
    -xattrs \
    -no-progress -noappend

TIEMPO_COMP=$SECONDS
echo
ok "Recompresión en $(( TIEMPO_COMP / 60 )) min $(( TIEMPO_COMP % 60 )) seg"
info "Tamaño squashfs nuevo: $(du -sh "${SQUASHFS_NUEVO}" | cut -f1)"

# -----------------------------------------------------------------------------
# PASO 5 — Actualizar metadatos
# -----------------------------------------------------------------------------
info "Paso 5/6 — Actualizando metadatos..."

SIZE_FILE="$(dirname "${SQUASHFS_NUEVO}")/filesystem.size"
printf '%s' "$(du -sx --block-size=1 "${DIR_SQUASHFS}" | cut -f1)" > "${SIZE_FILE}"

if [[ -f "${DIR_NEWISO}/md5sum.txt" ]]; then
    pushd "${DIR_NEWISO}" > /dev/null
    find . -type f ! -name 'md5sum.txt' | sort | xargs md5sum > md5sum.txt
    popd > /dev/null
    ok "md5sum.txt recalculado"
fi
ok "Metadatos actualizados"

# -----------------------------------------------------------------------------
# PASO 6 — Reconstruir la ISO híbrida
# -----------------------------------------------------------------------------
info "Paso 6/6 — Reconstruyendo la ISO híbrida..."

ISO_TMP="${ISO_SALIDA}.tmp"

VOLID=$(xorriso -indev "${ISO_ORIGINAL}" -report_system_area as_mkisofs 2>/dev/null \
    | grep -- "-volid" | sed "s/.*-volid '\([^']*\)'.*/\1/" | head -1) || true
[[ -z "$VOLID" ]] && VOLID="QUIRINUX"
info "Volume ID: ${VOLID}"

BOOT_PARAMS=$(xorriso -indev "${ISO_ORIGINAL}" \
    -report_system_area as_mkisofs 2>/dev/null \
    | grep -v "^$" | grep -v -- "-volid" \
    | tr '\n' ' ')

cat > "${WORKDIR}/build_iso.sh" << BUILDSCRIPT
#!/usr/bin/env bash
xorriso -as mkisofs \\
  -iso-level 3 \\
  -full-iso9660-filenames \\
  -volid '${VOLID}' \\
  ${BOOT_PARAMS} \\
  -output '${ISO_TMP}' \\
  '${DIR_NEWISO}'
BUILDSCRIPT
chmod +x "${WORKDIR}/build_iso.sh"

bash "${WORKDIR}/build_iso.sh"

[[ ! -f "${ISO_TMP}" ]] && die "xorriso no generó la ISO."

info "Restaurando MBR híbrido original..."
dd if="${MBR_BIN}" of="${ISO_TMP}" bs=1 count=432 conv=notrunc status=none
ok "MBR restaurado"

mv "${ISO_TMP}" "${ISO_SALIDA}"

# Limpieza normal
trap - ERR
info "Limpiando directorio de trabajo..."
rm -rf "${WORKDIR}"

# -----------------------------------------------------------------------------
# Resumen
# -----------------------------------------------------------------------------
echo
echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
echo -e "${BOLD}  RESUMEN FINAL                                         ${NC}"
echo -e "${BOLD}════════════════════════════════════════════════════════${NC}"
echo -e "  Original  : ${YELLOW}$(du -sh "${ISO_ORIGINAL}" | cut -f1)${NC}  →  $(basename "${ISO_ORIGINAL}")"
echo -e "  Nueva     : ${GREEN}$(du -sh "${ISO_SALIDA}"   | cut -f1)${NC}  →  $(basename "${ISO_SALIDA}")"
echo -e "  Compresión: ${BOLD}xz + bcj x86${NC}"
echo -e "  Tiempo    : ${BOLD}$(( SECONDS / 60 )) min $(( SECONDS % 60 )) seg${NC}"
echo
ok "¡Listo!"
echo
