#!/usr/bin/env bash

set -e
# set -u

export LANG=C

script_path=$(readlink -f ${0%/*})

cache_dir="${script_path}/cache"
channels_dir="${script_path}/channels"
isolinux_dir="${script_path}/isolinux"

codename="focal"
mirror="http://ftp.jaist.ac.jp/pub/Linux/ubuntu/"
os_name="Ubuntu"
iso_name="ubuntu"

arch=amd64
work_dir="${script_path}/work"
out_dir="${script_path}/out"
iso_label="${os_name}_${codename}_${arch}"
iso_publisher='Fascode Network <https://fascode.net>'
iso_application="${os_name} Live/Rescue CD"
iso_version="${codename}-$(date +%Y.%m.%d)"
iso_filename="${iso_name}-${iso_version}-${arch}.iso"

channel_name="serene"

username="liveuser"
password="liveuser"
usershell="/bin/bash"

debug=false
cache_only=false


start_time="$(date +%s)"

_msg_common() {
    if [[ "${debug}" = true ]]; then
        local _current_time="$(date +%s)"
        local _time

        _time="$(("${_current_time}"-"${start_time}"))"

        if [[ "${_time}" -ge 3600 ]]; then
            echo "[$(date -d @${_time} +%H:%M.%S)]$("${script_path}/echo_color" -t 6 "[LUBS Core]")"
        elif [[ "${_time}" -ge 60 ]]; then
            echo "[00:$(date -d @${_time} +%M.%S)]$("${script_path}/echo_color" -t 6 "[LUBS Core]")"
        else
            echo "[00:00.$(date -d @${_time} +%S)] $("${script_path}/echo_color" -t 6 "[LUBS Core]")"
        fi
    else
        "${script_path}/echo_color" -t 6 "[LUBS Core]"
    fi
}

# Show an INFO message
# _msg_info <message>
_msg_info() {
    local _msg="${@}"
    echo "$(_msg_common)  $("${script_path}/echo_color" -t 2 "Info:") ${_msg}"
}

# Show an debug message
# _msg_debug <message>
_msg_debug() {
    if [[ "${debug}" = true ]]; then
        local _msg="${@}"
        echo "$(_msg_common)  $("${script_path}/echo_color" -t 3 "Debug:") ${_msg}"
    fi
}

# Show an ERROR message then exit with status
# _msg_error <message> <exit code>
_msg_error() {
    local _msg="${1}"
    local _error=${2}
    echo "$(_msg_common)  $("${script_path}/echo_color" -t 1 "Error:") ${_msg}"
    if [[ ! ${_error} = 0 ]]; then
        exit ${_error}
    fi
}

# Unmount chroot dir
umount_chroot () {
    local mount
    for mount in $(mount | awk '{print $3}' | grep $(realpath ${work_dir}) | sort -r); do
        _msg_info "Unmounting ${mount}"
        umount -fl "${mount}"
    done
}

# Helper function to run make_*() only one time.
run_once() {
    local name
    umount_chroot

    if [[ "run_bootfiles" == "$1" ]]; then
        name="$2"
    else
        name="$1"
    fi

    if [[ ! -e "${work_dir}/build.${name}" ]]; then
        _msg_info "$(echo $name | sed "s@_@ @g") is starting."

        if [[ "run_bootfiles" == "$1" ]]; then
            "$1" "$2"
        else
            "$1"
        fi

        _msg_info "$(echo $name | sed "s@_@ @g") was done!"
        touch "${work_dir}/build.${name}"
    fi
}

run_cmd() {
    local mount
    for mount in "dev" "dev/pts" "proc" "sys" "run/systemd/resolve/stub-resolv.conf"; do
        [[ "${mount}" == "run/systemd/resolve/stub-resolv.conf" ]] && mount --bind /etc/resolv.conf "${work_dir}/airootfs/${mount}" || mount --bind /${mount} "${work_dir}/airootfs/${mount}"
    done
    chroot "${work_dir}/airootfs" "${@}"
    for mount in $(mount | awk '{print $3}' | grep $(realpath ${work_dir}) | sort -r); do
        umount -fl "${mount}"
    done
}

run_bootfiles() {
    cd "${bootfiles_dir}"
    "$1"
    cd - > /dev/null
}

_apt_install() {
    run_cmd apt-get --no-install-recommends --yes install ${@}
}

# rm helper
# Delete the file if it exists.
# For directories, rm -rf is used.
# If the file does not exist, skip it.
# remove <file> <file> ...
remove() {
    local _list
    local _file
    _list=($(echo "$@"))
    for _file in "${_list[@]}"; do
        _msg_debug "Removeing ${_file}"
        if [[ -f ${_file} ]]; then
            rm -f "${_file}"
        elif [[ -d ${_file} ]]; then
            rm -rf "${_file}"
        fi
    done
}

# Show help
_usage () {
    echo "usage ${0} [options] [channel]"
    echo
    echo " General options:"
    echo
    echo "    -a | --arch <str>      Set architecture"
    echo "                           Default: ${arch}"
    echo "    -c | --codename <str>  Set ubuntu codename"
    echo "                           Default: ${codename}"
    echo "    -m | --mirror <url>    Set apt mirror server."
    echo "                           Default: ${mirror}"
    echo "    -o | --out <out_dir>   Set the output directory"
    echo "                           Default: ${out_dir}"
    echo "    -w | --work <work_dir> Set the working directory"
    echo "                           Default: ${work_dir}"
    echo
    echo "    -d | --debug           "
    echo "    -h | --help            This help message and exit"
    echo
    echo "You can switch between installed packages, files included in images, etc. by channel."
    echo
    echo " Channel:"
    local _channel
    local channel_list
    local description
    for _channel in $(ls -l "${channels_dir}" | awk '$1 ~ /d/ {print $9 }'); do
        if [[ -n $(ls "${channels_dir}/${_channel}") ]] && [[ ! "${_channel}" = "share" ]]; then
            channel_list="${channel_list[@]} ${_channel}"
        fi
    done
    for _channel in ${channel_list[@]}; do
        if [[ -f "${channels_dir}/${_channel}/description.txt" ]]; then
            description=$(cat "${channels_dir}/${_channel}/description.txt")
        else
            description="This channel does not have a description.txt."
        fi
        echo -ne "    ${_channel}"
        for i in $( seq 1 $(( 23 - ${#_channel} )) ); do
            echo -ne " "
        done
        echo -ne "${description}\n"
    done
}


prepare_build() {
    if [[ ${EUID} -ne 0 ]]; then
        _msg_error "This script must be run as root." 1
    fi
    
    [[ ! -d "${work_dir}" ]] && mkdir -p "${work_dir}"
    [[ ! -d "${out_dir}" ]] && mkdir -p "${out_dir}"

    umount_chroot 

    # Check codename
    if [[ -z $(grep -h -v ^'#' ${channels_dir}/${channel_name}/codename.${arch} | grep -x ${codename}) ]]; then
        _msg_error "This codename (${channel_name}) is not supported on this channel (${codename})."
    fi

}


make_basefs() {
    local debootstrap_status
    statusfile="${cache_dir}/${codename}/status"

    debootstrap_status() {
        if [[ ! -d "$(dirname ${statusfile})" ]]; then
            mkdir -p "$(dirname ${statusfile})"
        fi
        echo "${1}" > "${statusfile}"
    }

    if [[ -f "${statusfile}" ]] && [[ $(cat "${statusfile}" 2> /dev/null) = "Done" ]]; then
        _msg_info "${codename} cache is found."
    else
        remove "${cache_dir}/${codename}"
        debootstrap_status "Running"
        _msg_info "Installing Ubuntu to '${cache_dir}/${codename}/airootfs'..."
        mkdir -p "${cache_dir}/${codename}/airootfs"
        debootstrap --arch=${arch} --include=linux-image-generic  --verbose --merged-usr "${codename}" "${cache_dir}/${codename}/airootfs" ${mirror}
        _msg_info "${codename} installed successfully!"
        debootstrap_status "Done"
    fi
    if [[ "${cache_only}" = true ]]; then
        exit 0
    fi
    
    rm -rf "${work_dir}/airootfs" && mkdir -p "${work_dir}/airootfs"
    _msg_info "copy base files from '${cache_dir}/${codename}/airootfs' to '${work_dir}/airootfs'..."
    rsync  -au "${cache_dir}/${codename}/airootfs/" "${work_dir}/airootfs"
    echo 'PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:${PATH}' >> "${work_dir}/airootfs/etc/bash.bashrc"
    run_cmd apt-get update
    # run_cmd apt-get upgrade
}

make_sourcelist() {
    cp ${script_path}/source.list.d/${codename}/* ${work_dir}/airootfs/etc/apt
    run_cmd apt-get update

    if [[ -n $(find ${channels_dir}/*/repo -type f) ]]; then
        _apt_install gnupg

        for repo in $(find ${channels_dir}/*/repo -name *.list); do
            key="$(dirname ${repo})/$(basename ${repo} | sed "s/list/key/")"
            
            if [[ -f "${key}" ]]; then
                if $(file ${key} | grep -q "PGP/GPG key"); then
                    cp ${key} ${work_dir}/airootfs/$(basename ${key})
                else
                    wget -q -O ${work_dir}/airootfs/$(basename ${key}) $(cat ${key})
                fi

                run_cmd apt-key add $(basename ${key})
                rm ${work_dir}/airootfs/$(basename ${key})
                cp ${repo} ${work_dir}/airootfs/etc/apt/sources.list.d
            fi
        done

        run_cmd apt-get update
    fi
    
    # PPA
    local PPA_FILELIST
    local _PPA_FILE
    local _ppa

    PPA_FILELIST=("${channels_dir}/share/ppa_list.${arch}" "${channels_dir}/${channel_name}/ppa_list.${arch}")
    _apt_install software-properties-common
    for _PPA_FILE in ${PPA_FILELIST[@]}; do
        if [[ -f "${_PPA_FILE}" ]]; then
            for _ppa in $(grep -h -v ^'#' ${_PPA_FILE}); do
                run_cmd add-apt-repository --yes "${_ppa}"
            done
        fi
    done
}

make_systemd() {
    _apt_install systemd-sysv
    run_cmd dbus-uuidgen > /etc/machine-id
    run_cmd ln -fs /etc/machine-id /var/lib/dbus/machine-id
}

make_apt_packages() {
    remove "${work_dir}/airootfs/aptpkglist"

    if [[ -f "${channels_dir}/share/packages.${arch}" ]]; then
        grep -h -v ^'#' "${channels_dir}/share/packages.${arch}" | grep -v "^$" >> "${work_dir}/airootfs/aptpkglist"
    fi
    if [[ -f "${channels_dir}/${channel_name}/packages.${arch}" ]]; then
        grep -h -v ^'#' "${channels_dir}/${channel_name}/packages.${arch}" | grep -v "^$" >> "${work_dir}/airootfs/aptpkglist"
    fi
    if [[ -n "$(echo $(<"${work_dir}/airootfs/aptpkglist"))" ]]; then
        run_cmd env -i bash -c 'DEBIAN_FRONTEND=noninteractive apt-get --no-install-recommends --yes install $(echo $(<aptpkglist))'
    fi
    remove "${work_dir}/airootfs/aptpkglist"
}

make_snap_packages() {
    remove "${work_dir}/airootfs/snappkglist"

    if [[ -f "${channels_dir}/share/snap-packages.${arch}" ]]; then
        grep -h -v ^'#' "${channels_dir}/share/snap-packages.${arch}" | grep -v "^$" >> "${work_dir}/airootfs/snappkglist"
    fi
    if [[ -f "${channels_dir}/${channel_name}/snap-packages.${arch}" ]]; then
        grep -h -v ^'#' "${channels_dir}/${channel_name}/snap-packages.${arch}" | grep -v "^$" >> "${work_dir}/airootfs/snappkglist"
    fi
    if [[ -n "$(echo $(<"${work_dir}/airootfs/snappkglist"))" ]]; then
        _apt_install snapd
        run_cmd env -i bash -c 'snap install $(echo $(<snappkglist))'
    fi
    remove "${work_dir}/airootfs/snappkglist"
}

make_config() {
    # Locales
    run_cmd env -i bash -c 'DEBIAN_FRONTEND=noninteractive LANG=en_US.UTF-8 LC_ALL=C LANGUAGE=en_US.UTF-8 dpkg-reconfigure locales'

    # resolvconf
    run_cmd env -i bash -c 'DEBIAN_FRONTEND=noninteractive dpkg-reconfigure resolvconf'

    # NetworkManager
    cp ${channels_dir}/share/airootfs/etc/NetworkManager/NetworkManager.conf ${work_dir}/airootfs/etc/NetworkManager/NetworkManager.conf

    # Timezone
    #run_cmd echo -ne "UTC" > '/etc/timezone'
    #run_cmd dpkg-reconfigure -f noninteractive tzdata

    run_cmd env -i bash -c 'DEBIAN_FRONTEND=noninteractive dpkg-reconfigure network-manager'
    run_cmd truncate -s 0 /etc/machine-id
}

make_customize_airootfs() {
    #-- Overwrite airootfs files --#
    local copy_airootfs
    copy_airootfs() {
        local _dir="${1%/}"
        if [[ -d "${_dir}" ]]; then
            cp -af "${_dir}"/* "${work_dir}/airootfs"
        fi
    }
    copy_airootfs "${channels_dir}/share/airootfs"
    copy_airootfs "${channels_dir}/${channel_name}/airootfs"


    #-- Run scripts --#
    local _customize_siript_options="-a ${arch} -u '${username}' -p '${password}' -s '${usershell}'"
    if ${bash_debug}; then
        _customize_siript_options="$_customize_siript_options -x"
    elif ${debug}; then
        _customize_siript_options="$_customize_siript_options -d"
    fi

    # run_customize_script <script path from lubs> <script path in chroot>
    local run_customize_script
    run_customize_script() {
        if [[ -f "${1}" ]]; then
            chmod 755 "${1}"
            run_cmd "${2} ${_customize_siript_options}"
        fi
    }
    run_customize_script "${work_dir}/airootfs/root/customize_airootfs.sh" "/root/customize_airootfs.sh"
    run_customize_script "${work_dir}/airootfs/root/customize_airootfs_${channel_name}.sh"  "/root/customize_airootfs_${channel_name}.sh"
}

make_clean() {
    run_cmd apt-get clean
    run_cmd apt-get --yes autoremove
    run_cmd rm -rf "/tmp/* ~/.bash_history"
}

make_bootfiles() {
    run_cmd env -i bash -c 'update-initramfs -c -k all'
    _apt_install memtest86+
    [[ -d "${bootfiles_dir}" ]] && rm -r "${bootfiles_dir}"
    mkdir -p ${bootfiles_dir}/{casper,isolinux,install}
    cp "${work_dir}/airootfs/boot/vmlinuz" "${bootfiles_dir}/casper/vmlinuz"
    cp "${work_dir}/airootfs/boot/initrd.img" "${bootfiles_dir}/casper/initrd"
    cp "${work_dir}/airootfs/boot/memtest86+.bin" "${bootfiles_dir}/install/memtest86+"

    if [[ ! -f "${cache_dir}/memtest86-usb.zip" ]]; then
        wget -O ${cache_dir}/memtest86-usb.zip https://www.memtest86.com/downloads/memtest86-usb.zip
    fi

    (unzip -p ${cache_dir}/memtest86-usb.zip memtest86-usb.img > ${bootfiles_dir}/install/memtest86)
}

make_grubcfg() {
    touch "${bootfiles_dir}/ubuntu"
    cp ${isolinux_dir}/grub.cfg ${bootfiles_dir}/isolinux/grub.cfg
}

make_manifest() {
    run_cmd dpkg-query -W --showformat='${Package} ${Version}\n' | tee ${bootfiles_dir}/casper/filesystem.manifest
    cp -v ${bootfiles_dir}/casper/filesystem.manifest "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/ubiquity/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/casper/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/discover/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/laptop-detect/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/os-prober/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
}

make_squashfs() {
    mksquashfs "${work_dir}/airootfs" "${bootfiles_dir}/casper/filesystem.squashfs"
    printf $(du -sx --block-size=1 "${work_dir}/airootfs" | cut -f1) > ${bootfiles_dir}/casper/filesystem.size
}

make_deifnes() {
    cp ${isolinux_dir}/README.diskdefines ${bootfiles_dir}/README.diskdefines
}

make_isolinux() {
    if [[ "${arch}" = "amd64" ]]; then
        local _arch="x86_64"
    else
        local _arch="${arch}"
    fi
    grub-mkstandalone \
        --format=${_arch}-efi \
        --output=isolinux/bootx64.efi \
        --locales="" \
        --fonts="" \
        "boot/grub/grub.cfg=isolinux/grub.cfg"
    (
        cd isolinux && \
        dd if=/dev/zero of=efiboot.img bs=1M count=10 && \
        sudo mkfs.vfat efiboot.img && \
        LC_CTYPE=C mmd -i efiboot.img efi efi/boot && \
        LC_CTYPE=C mcopy -i efiboot.img ./bootx64.efi ::efi/boot/
    )
    grub-mkstandalone \
        --format=i386-pc \
        --output=isolinux/core.img \
        --install-modules="linux16 linux normal iso9660 biosdisk memdisk search tar ls" \
        --modules="linux16 linux normal iso9660 biosdisk search" \
        --locales="" \
        --fonts="" \
        --themes="" \
        "boot/grub/grub.cfg=isolinux/grub.cfg"
    cat /usr/lib/grub/i386-pc/cdboot.img isolinux/core.img > isolinux/bios.img
}

make_md5sum() {
    /bin/bash -c "(find . -type f -print0 | xargs -0 md5sum | grep -v "\./md5sum.txt" > md5sum.txt)"
}

make_iso() {
    xorriso \
        -as mkisofs \
        -iso-level 3 \
        -full-iso9660-filenames \
        -volid "${iso_label}" \
        -appid "${iso_application}" \
        -publisher "${iso_publisher}" \
        -preparer "prepared by LUBS" \
        -eltorito-boot boot/grub/bios.img \
        -no-emul-boot \
        -boot-load-size 4 \
        -boot-info-table \
        --eltorito-catalog boot/grub/boot.cat \
        --grub2-boot-info \
        --grub2-mbr /usr/lib/grub/i386-pc/boot_hybrid.img \
        -eltorito-alt-boot \
        -e EFI/efiboot.img \
        -no-emul-boot \
        -append_partition 2 0xef isolinux/efiboot.img \
        -output "${out_dir}/${iso_filename}" \
        -graft-points \
            "." \
            /boot/grub/bios.img=isolinux/bios.img \
            /EFI/efiboot.img=isolinux/efiboot.img
}

make_checksum() {
    cd ${out_dir}
    _msg_info "Creating md5 checksum ..."
    md5sum "${iso_filename}" > "${iso_filename}.md5"

    _msg_info "Creating sha256 checksum ..."
    sha256sum "${iso_filename}" > "${iso_filename}.sha256"
    cd - > /dev/null 2>&1
}


# 引数解析（）
# 参考記事：https://0e0.pw/ci83 https://0e0.pw/VJlg

_opt_short="w:o:ha:-:m:c:d"
_opt_long="help,arch:,codename:,debug,help,mirror:,out:,work,cache-only"
OPT=$(getopt -o ${_opt_short} -l ${_opt_long} -- "${@}")
if [[ ${?} != 0 ]]; then
    exit 1
fi
eval set -- "${OPT}"
while :; do
    case ${1} in
        -a | --arch)
            if [[ -z ${2} ]]; then
                _msg_error "Please specify the architecture."
                exit 1
            else
                arch="${2}"
            fi
            shift 2
            ;;
        -c | --codename)
            if [[ -z ${2} ]]; then
                _msg_error "Please specify the codename."
                exit 1
            else
                codename="${2}"
            fi
            shift 2
            ;;
        -d | --debug)
            debug=true
            shift 1
            ;;
        -h | --help)
            _usage
            exit 0
            ;;
        -m | --mirror)
            if [[ -z ${2} ]]; then
                _msg_error "Please specify the mirror server."
                exit 1
            else
                mirror="${2}"
            fi
            shift 2
            ;;
        -o | --out)
            if [[ -z ${2} ]]; then
                _msg_error "Please specify the out dir."
                exit 1
            fi
            shift 2
            ;;
        -w | --work)
            if [[ -z ${2} ]]; then
                _msg_error "Please specify the out dir."
                exit 1
            else
                work_dir="${2}"
            fi
            shift 2
            ;;
        --cache-only)
            cache_only=true
            shift 1
            ;;
        --)
            shift
            break
            ;;
        *)
            _msg_error "Invalid argument '${1}'"
            _usage 1
            ;;
    esac
done

bootfiles_dir="${work_dir}/bootfiles"
trap  umount_chroot 0 2 15

if [[ -n "${1}" ]]; then
    channel_name="${1}"

    check_channel() {
        local channel_list
        local i
        channel_list=()
        for _channel in $(ls -l "${channels_dir}" | awk '$1 ~ /d/ {print $9 }'); do
            if [[ -n $(ls "${channels_dir}/${_channel}") ]] && [[ ! "${_channel}" = "share" ]]; then
                channel_list="${channel_list[@]} ${_channel}"
            fi
        done
        for i in ${channel_list[@]}; do
            if [[ "${i}" = "${channel_name}" ]]; then
                echo -n "true"
                return 0
            fi
        done
        echo -n "false"
        return 1
    }

    if [[ $(check_channel ${channel_name}) = false ]]; then
        _msg_error "Invalid channel ${channel_name}"
        exit 1
    fi
fi


prepare_build
run_once make_basefs
run_once make_sourcelist
run_once make_systemd
run_once make_apt_packages
#run_once make_snap_packages
run_once make_config
run_once make_customize_airootfs
run_once make_clean
run_once make_bootfiles
run_once make_grubcfg
run_once make_manifest
run_once make_squashfs
run_once make_deifnes
run_once run_bootfiles make_isolinux
run_once run_bootfiles make_md5sum
run_once run_bootfiles make_iso
run_once make_checksum