#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-3.0
#
# mk-linux419
# Twitter: @fascoder_4
# Email  : m.k419sabuaka@gmail.com
#
# Yamada Hayao
# Twitter: @Hayao0819
# Email  : hayao@fascode.net
#
# (c) 2019-2020 Fascode Network.
#
# lubs
#

set -e
# set -u

export LANG=C

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

cache_dir="${script_path}/cache"
channels_dir="${script_path}/channels"
nfb_dir="${script_path}/nfb"

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
        local _time
        _current_time="$(date +%s)"
        _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
    _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
        _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
    local _error
    _msg="${1}"
    _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
    name="$1"

    if [[ ! -e "${work_dir}/build.${name}" ]]; then
        _msg_info "$(echo $name | sed "s@_@ @g") is starting."
        "${1}"
        _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
        if [[ "${mount}" == "run/systemd/resolve/stub-resolv.conf" ]]; then
            mount --bind /etc/resolv.conf "${work_dir}/airootfs/${mount}"
        else
            mount --bind /${mount} "${work_dir}/airootfs/${mount}"
        fi
    done
    
    chroot "${work_dir}/airootfs" "${@}"

    for mount in $(mount | awk '{print $3}' | grep "$(realpath "${work_dir}")" | sort -r); do
        umount -fl "${mount}"
    done
}

_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}" )
        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}" --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
    sed -i "s@http://archive.ubuntu.com/ubuntu/@${mirror}@g"  ${work_dir}/airootfs/etc/apt/sources.list
    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 --ensure=/etc/machine-id
    run_cmd ln -sf /etc/machine-id /var/lib/dbus/machine-id
}

make_apt_packages() {
    remove "${work_dir}/airootfs/aptpkglist"
    _apt_install initramfs-tools
    run_cmd env -i bash -c 'DEBIAN_FRONTEND=noninteractive apt-get dist-upgrade --yes'

    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 [[ -s "${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 [[ -s "${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
    _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
    if ${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() {
    sed -i "s@${mirror}@http://archive.ubuntu.com/ubuntu/@g"  ${work_dir}/airootfs/etc/apt/sources.list
    run_cmd apt-get update
    run_cmd apt-get clean
    run_cmd apt-get --yes autoremove
    run_cmd rm -rf "/tmp/* ~/.bash_history"
}

make_squashfs() {
    # prepare
    [[ -d "${bootfiles_dir}" ]] && rm -r "${bootfiles_dir}"
    mkdir -p "${bootfiles_dir}"/{casper,isolinux,install}

    # make squashfs
    mksquashfs "${work_dir}/airootfs" "${bootfiles_dir}/casper/filesystem.squashfs"
}

make_nfb() {
    # prepare
    touch "${bootfiles_dir}/ubuntu"
    sed "s|%OS_NAME%|${os_name}|g" "${nfb_dir}/README.diskdefines" > "${bootfiles_dir}/README.diskdefines"

    # casper setup
    cp "${work_dir}/airootfs/boot/vmlinuz" "${bootfiles_dir}/casper/vmlinuz"
    cp "${work_dir}/airootfs/boot/initrd.img" "${bootfiles_dir}/casper/initrd"
    run_cmd dpkg-query -W --showformat='${Package} ${Version}\n' > ${bootfiles_dir}/casper/filesystem.manifest
    cp -v "${bootfiles_dir}/casper/filesystem.manifest" "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    sed -i '/casper/d;/discover/d;/laptop-detect/d;/os-prober/d;/ubiquity/d' "${bootfiles_dir}/casper/filesystem.manifest-desktop"
    printf $(du -sx --block-size=1 "${work_dir}/airootfs" | cut -f1) > "${bootfiles_dir}/casper/filesystem.size"

    # isolinux setup
    cp "${work_dir}"/airootfs/usr/lib/syslinux/modules/bios/*.c32 "${bootfiles_dir}/isolinux/"
    cp "${work_dir}/airootfs/usr/lib/ISOLINUX/isolinux.bin" "${bootfiles_dir}/isolinux/"
    cp "${work_dir}/airootfs/usr/lib/ISOLINUX/isohdpfx.bin" "${bootfiles_dir}/isolinux/"
    sed "s|%OS_NAME%|${os_name}|g" "${nfb_dir}/grub.cfg" > "${bootfiles_dir}/isolinux/grub.cfg"
    sed "s|%OS_NAME%|${os_name}|g" "${nfb_dir}/isolinux.cfg" > "${bootfiles_dir}/isolinux/isolinux.cfg"

    # memtest86+ setup
    _apt_install memtest86+
    cp "${work_dir}/airootfs/boot/memtest86+.bin" "${bootfiles_dir}/install/memtest86+"

    # memtest86 setup
    if [[ ! -f "${cache_dir}/memtest86-usb.zip" ]]; then
        wget -O "${cache_dir}/memtest86-usb.zip" "https://www.memtest86.com/downloads/memtest86-usb.zip"
        bash -c "(unzip -p ${cache_dir}/memtest86-usb.zip memtest86-usb.img > ${cache_dir}/memtest86)"
    fi

    cp "${cache_dir}/memtest86" "${bootfiles_dir}/install/memtest86"
}

make_efi() {
    # UEFI 32bit (ia32)
    grub-mkstandalone \
        --format=i386-efi \
        --output="${bootfiles_dir}/isolinux/bootia32.efi" \
        --locales="" \
        --fonts="" \
        "boot/grub/grub.cfg=${bootfiles_dir}/isolinux/grub.cfg"
    
    # UEFI 64bit (x64)
    grub-mkstandalone \
        --format=x86_64-efi \
        --output="${bootfiles_dir}/isolinux/bootx64.efi" \
        --locales="" \
        --fonts="" \
        "boot/grub/grub.cfg=${bootfiles_dir}/isolinux/grub.cfg"

    # create efiboot.img
    truncate -s 10M "${bootfiles_dir}/isolinux/efiboot.img"
    mkfs.fat -F 16 -f 1 -r 112 "${bootfiles_dir}/isolinux/efiboot.img"
    mkdir "${bootfiles_dir}/mnt"
    mount "${bootfiles_dir}/isolinux/efiboot.img" "${bootfiles_dir}/mnt"
    cp "${bootfiles_dir}/isolinux/bootia32.efi" "${bootfiles_dir}/mnt"
    cp "${bootfiles_dir}/isolinux/bootx64.efi" "${bootfiles_dir}/mnt"
    umount -d "${bootfiles_dir}/mnt"
    rm -r "${bootfiles_dir}/mnt"
}

make_iso() {
    cd "${bootfiles_dir}"

    # create checksum (using at booting)
    bash -c "(find . -type f -print0 | xargs -0 md5sum | grep -v "\./md5sum.txt" > md5sum.txt)"

    # create 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 isolinux/isolinux.bin \
            -no-emul-boot \
            -boot-load-size 4 \
            -boot-info-table \
        -eltorito-alt-boot \
            -eltorito-platform efi \
            -eltorito-boot EFI/efiboot.img \
            -no-emul-boot \
        -isohybrid-mbr "${bootfiles_dir}/isolinux/isohdpfx.bin" \
        -isohybrid-gpt-basdat \
        -eltorito-catalog isolinux/boot.cat \
        -output "${out_dir}/${iso_filename}" \
        -graft-points \
            "." \
            "/isolinux/isolinux.bin=isolinux/isolinux.bin" \
            "/EFI/efiboot.img=isolinux/efiboot.img"
    
    cd - > /dev/null
}

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
            else
                out_dir="${2}"
            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}" )
            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_squashfs
run_once make_nfb
run_once make_efi
run_once make_iso
run_once make_checksum