#!/bin/bash
set -e

XEN_VERSION="4.18.0"
XEN_URL="https://downloads.xenproject.org/release/xen/${XEN_VERSION}/xen-${XEN_VERSION}.tar.gz"
BUILD_DIR="/tmp/xen-turnstile-build"
TURNSTILE_DIR="$(cd "$(dirname "$0")" && pwd)"

echo "=== Turnstile Xen Build Script ==="
echo "Xen version: ${XEN_VERSION}"
echo "Turnstile source: ${TURNSTILE_DIR}"
echo "Build directory: ${BUILD_DIR}"
echo ""

install_deps_arch() {
    echo "Installing build dependencies for Arch Linux..."
    sudo pacman -Sy --needed --noconfirm \
        base-devel git wget python python-setuptools \
        iasl acpica yajl pixman bridge-utils \
        iproute2 libaio glib2 libpng sdl2 ncurses \
        libnl openssl zlib lzo libseccomp systemd \
        e2fsprogs nasm flex bison ninja meson

    if ! command -v bcc &> /dev/null; then
        echo ""
        echo "dev86 not found. Installing from AUR..."
        if command -v yay &> /dev/null; then
            yay -S --needed --noconfirm dev86
        elif command -v paru &> /dev/null; then
            paru -S --needed --noconfirm dev86
        else
            echo "No AUR helper found. Installing dev86 manually..."
            cd /tmp
            rm -rf dev86
            git clone https://aur.archlinux.org/dev86.git
            cd dev86
            makepkg -si --noconfirm
            cd "${TURNSTILE_DIR}"
        fi
    fi
}

install_deps_debian() {
    echo "Installing build dependencies for Debian/Ubuntu..."
    sudo apt-get update
    sudo apt-get install -y \
        build-essential git wget python3 python3-dev \
        uuid-dev libncurses-dev libglib2.0-dev libyajl-dev \
        libaio-dev libpixman-1-dev libssl-dev libsdl2-dev \
        liblzo2-dev libseccomp-dev bison flex iasl \
        bridge-utils iproute2 e2fsprogs ninja-build meson
}

install_deps() {
    if [ -f /etc/arch-release ]; then
        install_deps_arch
    elif [ -f /etc/debian_version ]; then
        install_deps_debian
    else
        echo "Unknown distribution. Please install Xen build dependencies manually."
        echo "Required: build tools, python, uuid, ncurses, glib2, yajl, aio, pixman, ssl, sdl2, lzo, seccomp, iasl, ninja"
        exit 1
    fi
}

download_xen() {
    echo ""
    echo "=== Downloading Xen ${XEN_VERSION} ==="
    mkdir -p "${BUILD_DIR}"
    cd "${BUILD_DIR}"

    if [ ! -f "xen-${XEN_VERSION}.tar.gz" ]; then
        wget "${XEN_URL}"
    fi

    if [ ! -d "xen-${XEN_VERSION}" ]; then
        tar xzf "xen-${XEN_VERSION}.tar.gz"
    fi
}

copy_turnstile_sources() {
    echo ""
    echo "=== Copying turnstile sources ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    cp "${TURNSTILE_DIR}/xen-patch/turnstile.c" xen/arch/x86/
    cp "${TURNSTILE_DIR}/xen-patch/turnstile.h" xen/arch/x86/include/asm/

    echo "Copied turnstile.c to xen/arch/x86/"
    echo "Copied turnstile.h to xen/arch/x86/include/asm/"
}

apply_patches() {
    echo ""
    echo "=== Applying turnstile patches ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    DOMAIN_H="xen/arch/x86/include/asm/domain.h"
    DOMAIN_C="xen/arch/x86/domain.c"
    VMX_C="xen/arch/x86/hvm/vmx/vmx.c"
    SVM_C="xen/arch/x86/hvm/svm/svm.c"
    MAKEFILE="xen/arch/x86/Makefile"

    echo "Adding turnstile pointer to arch_domain..."
    if ! grep -q "struct turnstile_domain_state" "${DOMAIN_H}"; then
        sed -i '/^struct arch_domain$/i struct turnstile_domain_state;' "${DOMAIN_H}"
        sed -i '/^struct arch_domain$/,/^{$/ { /^{$/a\    struct turnstile_domain_state *turnstile;
        }' "${DOMAIN_H}"
    fi

    echo "Adding turnstile include and cleanup to domain.c..."
    if ! grep -q "asm/turnstile.h" "${DOMAIN_C}"; then
        sed -i '/#include <asm\/pv\/mm.h>/a #include <asm/turnstile.h>' "${DOMAIN_C}"
    fi
    if ! grep -q "turnstile_domain_destroy" "${DOMAIN_C}"; then
        sed -i '/^void arch_domain_destroy(struct domain \*d)/,/^{/ { /^{$/a\    turnstile_domain_destroy(d);
        }' "${DOMAIN_C}"
    fi

    echo "Adding turnstile EPT hook to vmx.c..."
    if ! grep -q "asm/turnstile.h" "${VMX_C}"; then
        sed -i '/#include <asm\/monitor.h>/a #include <asm/turnstile.h>' "${VMX_C}"
    fi
    if ! grep -q "turnstile_check_violation" "${VMX_C}"; then
        sed -i '/case EXIT_REASON_EPT_VIOLATION:/,/ept_handle_violation/ {
            /paddr_t gpa;/a\        uint8_t insn_bytes[16];
        }' "${VMX_C}"
        sed -i 's/        ept_handle_violation(exit_qualification, gpa);/        if ( (exit_qualification \& (1u << 1)) \&\& currd->arch.turnstile )\n        {\n            unsigned long rip, cr3;\n            int block;\n\n            __vmread(GUEST_RIP, \&rip);\n            __vmread(GUEST_CR3, \&cr3);\n\n            memset(insn_bytes, 0, sizeof(insn_bytes));\n            hvm_copy_from_guest_phys(insn_bytes, rip, sizeof(insn_bytes));\n\n            block = turnstile_check_violation(currd, gpa, rip, cr3,\n                                               exit_qualification,\n                                               insn_bytes);\n            if ( block )\n            {\n                hvm_inject_page_fault(PFEC_write_access | PFEC_page_present, gpa);\n                break;\n            }\n        }\n\n        ept_handle_violation(exit_qualification, gpa);/' "${VMX_C}"
    fi

    echo "Adding turnstile NPT hook to svm.c..."
    if ! grep -q "asm/turnstile.h" "${SVM_C}"; then
        sed -i '/#include <asm\/apic.h>/a #include <asm/turnstile.h>' "${SVM_C}"
    fi
    if ! grep -q "turnstile_check_violation" "${SVM_C}"; then
        sed -i '/case VMEXIT_NPF:/,/v->arch.hvm.svm.cached_insn_len = 0;/ {
            /if ( cpu_has_svm_decode )/i\        if ( (vmcb->ei.npf.ec \& PFEC_write_access) \&\& v->domain->arch.turnstile )\n        {\n            uint8_t insn_bytes[16];\n            unsigned long rip = vmcb->rip;\n            unsigned long cr3 = vmcb_get_cr3(vmcb);\n            int block;\n\n            memset(insn_bytes, 0, sizeof(insn_bytes));\n            hvm_copy_from_guest_phys(insn_bytes, rip, sizeof(insn_bytes));\n\n            block = turnstile_check_violation(v->domain, vmcb->ei.npf.gpa, rip, cr3,\n                                               vmcb->ei.npf.ec,\n                                               insn_bytes);\n            if ( block )\n            {\n                hvm_inject_page_fault(PFEC_write_access | PFEC_page_present,\n                                      vmcb->ei.npf.gpa);\n                break;\n            }\n        }\n
        }' "${SVM_C}"
    fi

    echo "Adding turnstile.o to Makefile..."
    if ! grep -q "turnstile.o" "${MAKEFILE}"; then
        sed -i '/^obj-y += usercopy.o/i obj-y += turnstile.o' "${MAKEFILE}"
    fi

    echo "Patches applied."
}

configure_xen() {
    echo ""
    echo "=== Configuring Xen ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    ./configure --prefix=/usr --sysconfdir=/etc --localstatedir=/var \
        --enable-systemd --disable-stubdom --disable-docs
}

build_xen() {
    echo ""
    echo "=== Building Xen hypervisor ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    make -j$(nproc) xen
}

build_tools() {
    echo ""
    echo "=== Building Xen tools ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    make -j$(nproc) tools
}

install_xen() {
    echo ""
    echo "=== Installing Xen hypervisor ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    sudo make install-xen

    echo ""
    echo "=== Updating boot configuration ==="
    if [ -f /etc/arch-release ]; then
        sudo grub-mkconfig -o /boot/grub/grub.cfg
    elif [ -f /etc/debian_version ]; then
        sudo update-grub
    fi
}

install_tools() {
    echo ""
    echo "=== Installing Xen tools ==="
    cd "${BUILD_DIR}/xen-${XEN_VERSION}"

    sudo make install-tools
}

print_usage() {
    echo "Usage: $0 [command]"
    echo ""
    echo "Commands:"
    echo "  deps        - Install build dependencies only"
    echo "  download    - Download Xen source only"
    echo "  patch       - Copy sources and apply patches only"
    echo "  configure   - Configure Xen build"
    echo "  build       - Build Xen hypervisor (assumes patched)"
    echo "  build-tools - Build Xen tools"
    echo "  install     - Install built Xen hypervisor"
    echo "  install-tools - Install Xen tools"
    echo "  hypervisor  - Build and install hypervisor only (no tools)"
    echo "  all         - Do everything (default)"
    echo ""
}

case "${1:-all}" in
    deps)
        install_deps
        ;;
    download)
        download_xen
        ;;
    patch)
        download_xen
        copy_turnstile_sources
        apply_patches
        ;;
    configure)
        configure_xen
        ;;
    build)
        build_xen
        ;;
    build-tools)
        build_tools
        ;;
    install)
        install_xen
        ;;
    install-tools)
        install_tools
        ;;
    hypervisor)
        install_deps
        download_xen
        copy_turnstile_sources
        apply_patches
        configure_xen
        build_xen
        install_xen
        echo ""
        echo "=== Turnstile-patched Xen ${XEN_VERSION} hypervisor installed ==="
        echo ""
        echo "Next steps:"
        echo "1. Reboot and select Xen from the bootloader"
        echo "2. Build turnstile-ctl: cd ${TURNSTILE_DIR}/turnstile-ctl && cargo build --release"
        echo "3. Start a guest domain and run: turnstile-ctl protect <domid> /path/to/vmlinux"
        ;;
    all)
        install_deps
        download_xen
        copy_turnstile_sources
        apply_patches
        configure_xen
        build_xen
        build_tools
        install_xen
        install_tools
        echo ""
        echo "=== Turnstile-patched Xen ${XEN_VERSION} installed ==="
        echo ""
        echo "Next steps:"
        echo "1. Reboot and select Xen from the bootloader"
        echo "2. Build turnstile-ctl: cd ${TURNSTILE_DIR}/turnstile-ctl && cargo build --release"
        echo "3. Start a guest domain and run: turnstile-ctl protect <domid> /path/to/vmlinux"
        ;;
    *)
        print_usage
        exit 1
        ;;
esac