#!/usr/bin/env perl

use strict;
use warnings;
use feature 'say';
use File::Copy qw(copy);
use File::Path qw(rmtree make_path);
use Cwd        qw(abs_path);
use FindBin;
use lib $FindBin::Bin;
use Util;

# ========================================================================================= #
# Project root
# ========================================================================================= #

my $root = abs_path("$FindBin::Bin/..");
chdir($root) or die "Cannot chdir to $root: $!\n";

# ========================================================================================= #
# Environment set by "nix develop"
# ========================================================================================= #

sub require_env {
    my ($name) = @_;
    return $ENV{$name}
      // die "$name is not set (run from 'nix develop' shell)\n";
}

my $wasi_root = require_env('WASI_ROOT');
my $wamr_root = require_env('WAMR_ROOT');
my $cross_cc  = require_env('CROSS_CC');
my $linux_cc  = require_env('LINUX_CC');

my $use_aot_in_text  = ( $ENV{WAMR_USE_AOT_IN_TEXT}  // 'false' ) eq 'true';
my $use_mmap_in_text = ( $ENV{WAMR_USE_MMAP_IN_TEXT} // 'false' ) eq 'true';
my $use_xip          = ( $ENV{WAMR_USE_XIP}          // 'false' ) eq 'true';
my $use_allocator    = ( $ENV{WAMR_USE_ALLOCATOR}    // 'false' ) eq 'true';

# ========================================================================================= #
# WAMR cmake configuration
# ========================================================================================= #

# Flags common to all platforms
my @wamr_cmake_base = (
    '-DCMAKE_BUILD_TYPE=Debug',    '-DWAMR_BUILD_TARGET=X86_32',
    '-DWAMR_BUILD_AOT=1',          '-DWAMR_BUILD_WAMR_COMPILER=0',
    '-DWAMR_BUILD_INTERP=1',       '-DWAMR_BUILD_FAST_INTERP=0',
    '-DWAMR_BUILD_JIT=0',          '-DWAMR_BUILD_FAST_JIT=0',
    '-DWAMR_BUILD_LIBC_BUILTIN=1', '-DWAMR_BUILD_LIBC_WASI=0',
    '-DWAMR_BUILD_SIMD=0',
);

my @wamr_cmake_baremetal = (
    '-DWAMR_BUILD_PLATFORM=baremetal',
    '-DCMAKE_SYSTEM_NAME=Generic',
    '-DCMAKE_SYSTEM_PROCESSOR=i386',
    '-DCMAKE_TRY_COMPILE_TARGET_TYPE=STATIC_LIBRARY',
);

my @wamr_cmake_linux = ( '-DWAMR_BUILD_PLATFORM=linux', );

# Variant defines passed as CMAKE_C_FLAGS
my @variant_cflags =
  ( '-Wno-error=implicit-function-declaration', '-O0', '-ggdb3' );
push @variant_cflags, '-DWAMR_MMAP_IN_TEXT=0' unless $use_mmap_in_text;
push @variant_cflags, '-DWASM_MEM_ALLOC_WITH_USAGE=1' if $use_allocator;
my $cmake_c_flags = join( ' ', @variant_cflags );

# Variant options passed as cmake variables
my @variant_cmake_flags;
push @variant_cmake_flags, '-DWAMR_BUILD_ALLOC_WITH_USAGE=1' if $use_allocator;

# ========================================================================================= #
# Compiler / linker flags
# ========================================================================================= #

my @wasi_cflags = (
    '--target=wasm32',          "--sysroot=$wasi_root/share/wasi-sysroot",
    '-z',                       'stack-size=4096',
    '-O0',                      '-nostdlib',
    '-Wl,--no-entry',           '-Wl,--export=wasm_module',
    '-Wl,--no-gc-sections',     '-Wl,--initial-memory=65536',
    '-Wl,--export=__heap_base', '-Wl,--export=__data_end',
);

my @cross_cflags =
  qw(-O0 -m32 -ffunction-sections -fdata-sections -ffreestanding -fpermissive -ggdb3);
my @linux_cflags =
  qw(-O0 -m32 -ffunction-sections -fdata-sections -fpermissive -ggdb3);
my @linux_baremetal_cflags =
  qw(-O0 -m32 -ffunction-sections -fdata-sections -ffreestanding -ggdb3);

# libiwasm.a is passed as a direct file path (no -L/-liwasm)
my @cross_ldflags_base = (
    '-Wl,--build-id=none', '-static', '-nostdlib', '-m32',
    '-lc', '-lgcc', '-lm'
);
my @linux_ldflags_base     = ( '-Wl,--build-id=none', '-m32', '-lm' );
my @baremetal_ldflags_base = (
    '-Wl,--build-id=none', '-static', '-nostdlib', '-m32',
    '-lc',                 '-lgcc',   '-lm',       '--entry',
    'main'
);
my @cross_ldflags_nowasm = (
    '-Wl,--build-id=none', '-static', '-nostdlib', '-m32',
    '-lc', '-lgcc', '-lm'
);
my @linux_ldflags_nowasm = ( '-Wl,--build-id=none', '-m32', '-lm' );

my @wamr_inc_baremetal = (
    "-I$wamr_root/core/iwasm/include",
    "-I$wamr_root/core/shared/utils",
    "-I$wamr_root/core/shared/platform/baremetal",
);
my @wamr_inc_linux = (
    "-I$wamr_root/core/iwasm/include",
    "-I$wamr_root/core/shared/utils",
    "-I$wamr_root/core/shared/platform/linux",
);

my @cross_wamrcflags = ( '--target=i386', '--cpu=generic', '--opt-level=0' );
push @cross_wamrcflags, '--xip' if $use_xip;
my @linux_wamrcflags = ( '--target=i386', '--cpu=generic', '--opt-level=0' );

# ========================================================================================= #
# Entry point
# ========================================================================================= #

die "Usage: compile.pl <module> <target> <mode>\n"
  . "  target: fail | linux | linux-baremetal\n"
  . "  mode:   aot  | interp | c\n"
  unless @ARGV == 3;

my ( $module, $target, $mode ) = @ARGV;

build( $module, $target, $mode );

# ========================================================================================= #
# Build
# ========================================================================================= #

sub build {
    my ( $module, $target, $mode ) = @_;
    my $bd = "build-$module";

    rmtree($bd);
    make_path($bd);
    copy_auxiliary( $module, $bd );
    build_libiwasm( $module, $bd, $target );

    if ( $mode eq 'aot' || $mode eq 'interp' ) {
        compile_wasm_module( $module, $bd );
        if ( $mode eq 'aot' ) {
            compile_wasm_aot( $module, $bd, $target );
            make_aot_array( $module, $bd );
        }
        else {
            make_interp_array( $module, $bd );
        }
        prepare_wasm_host( $module, $bd, $mode );
        compile_wasm_host( $module, $bd, $target );
        compile_startup( $module, $bd, $target );
        compile_syscalls( $module, $bd, $target );
        link_wasm( $module, $bd, $target );
    }
    elsif ( $mode eq 'c' ) {
        compile_c_module( $module, $bd, $target );
        compile_c_host( $module, $bd, $target );
        compile_startup( $module, $bd, $target );
        link_c( $module, $bd, $target );
    }
    else {
        die "Unknown mode '$mode'; expected: aot, interp, c\n";
    }

    build_iso( $module, $bd );
}

# ========================================================================================= #
# Steps
# ========================================================================================= #

sub build_libiwasm {
    my ( $module, $bd, $target ) = @_;
    my $cmake_bd = "$bd/wamr-build";
    make_path($cmake_bd);

    my ( $cmake_cc, @sys_flags ) =
        ( $target eq 'linux' )
      ? ( $linux_cc, @wamr_cmake_linux )
      : ( $cross_cc, @wamr_cmake_baremetal );

    Util::run(
        'cmake',              '-S',
        $wamr_root,           '-B',
        $cmake_bd,            "-DCMAKE_C_COMPILER=$cmake_cc",
        @wamr_cmake_base,     @sys_flags,
        @variant_cmake_flags, "-DCMAKE_C_FLAGS=$cmake_c_flags",
    );
    Util::run( 'cmake', '--build', $cmake_bd );
    copy( "$cmake_bd/libiwasm.a", "$bd/libiwasm.a" )
      or die "failed to copy $cmake_bd/libiwasm.a: $!";
}

sub copy_auxiliary {
    my ( $module, $bd ) = @_;
    my @files = (
        [ 'flake.nix',                       "$bd/flake.nix" ],
        [ 'scripts/runner.pl',               "$bd/runner.pl" ],
        [ 'scripts/compile.pl',              "$bd/compile.pl" ],
        [ 'targets/lib.h',                   "$bd/lib.h" ],
        [ 'targets/linker.ld',               "$bd/linker.ld" ],
        [ 'targets/startup.s',               "$bd/startup.s" ],
        [ 'targets/syscalls.c',              "$bd/syscalls.c" ],
        [ "targets/wasm-module/$module.cpp", "$bd/wasm-module.cpp" ],
    );
    for my $pair (@files) {
        copy( $pair->[0], $pair->[1] )
          or die "failed to copy $pair->[0] -> $pair->[1]: $!";
    }
}

sub compile_wasm_module {
    my ( $module, $bd ) = @_;
    Util::run( "$wasi_root/bin/clang", @wasi_cflags,
        "targets/wasm-module/$module.cpp",
        '-o', "$bd/wasm_module.wasm" );
}

sub compile_wasm_aot {
    my ( $module, $bd, $target ) = @_;
    my @flags = ( $target eq 'linux' ) ? @linux_wamrcflags : @cross_wamrcflags;
    Util::run( 'wamrc', @flags, '-o', "$bd/wasm_module.aot", "$bd/wasm_module.wasm" );
}

sub make_aot_array {
    my ( $module, $bd ) = @_;
    open( my $xxd, '-|', 'xxd', '-i', "$bd/wasm_module.aot" )
      or die "failed to run xxd: $!";
    my $content = do { local $/; <$xxd> };
    close($xxd);

    $content =
      qq{__attribute__((section(".text.wamr_aot"), aligned(4096)))\n}
      . $content
      if $use_aot_in_text;

    Util::write_file( "$bd/wasm_aot_array.c", $content );
}

sub make_interp_array {
    my ( $module, $bd ) = @_;
    open( my $xxd, '-|', 'xxd', '-i', "$bd/wasm_module.wasm" )
      or die "Cannot run xxd: $!\n";
    my $content = do { local $/; <$xxd> };
    close($xxd);
    Util::write_file( "$bd/wasm_interp_array.c", $content );
}

sub prepare_wasm_host {
    my ( $module, $bd, $mode ) = @_;
    my $template =
      $use_allocator
      ? 'targets/wasm-host/wasm_host_allocator.c'
      : 'targets/wasm-host/wasm_host_pool.c';

    my $mod_c = ( $module =~ s/-/_/gr );
    my ( $array_file, $array_sym, $len_sym ) =
      ( $mode eq 'aot' )
      ? (
        'wasm_aot_array.c',
        "build_${mod_c}_wasm_module_aot",
        "build_${mod_c}_wasm_module_aot_len"
      )
      : (
        'wasm_interp_array.c',
        "build_${mod_c}_wasm_module_wasm",
        "build_${mod_c}_wasm_module_wasm_len"
      );

    my $content = Util::read_file($template);
    $content =~ s/__WASM_ARRAY_FILE__/$array_file/g;
    $content =~ s/__WASM_ARRAY__/$array_sym/g;
    $content =~ s/__WASM_ARRAY_LEN__/$len_sym/g;
    Util::write_file( "$bd/module_host.c", $content );
}

sub compile_wasm_host {
    my ( $module, $bd, $target ) = @_;
    if ( $target eq 'fail' ) {
        Util::run(
            $cross_cc,           '-I./targets/wasm-host',
            @cross_cflags,       @wamr_inc_baremetal,
            '-DTARGET_FAIL',     '-c',
            "$bd/module_host.c", '-o',
            "$bd/system.o"
        );
    }
    elsif ( $target eq 'linux' ) {
        Util::run(
            $linux_cc,           '-I./targets/wasm-host',
            @linux_cflags,       @wamr_inc_linux,
            '-DTARGET_LINUX',    '-c',
            "$bd/module_host.c", '-o',
            "$bd/system.o"
        );
    }
    elsif ( $target eq 'linux-baremetal' ) {
        Util::run(
            $cross_cc,                  '-I./targets/wasm-host',
            @linux_baremetal_cflags,    @wamr_inc_baremetal,
            '-DTARGET_LINUX_BAREMETAL', '-c',
            "$bd/module_host.c",        '-o',
            "$bd/system.o"
        );
    }
    else {
        die "Unknown target '$target'\n";
    }
}

sub compile_c_module {
    my ( $module, $bd, $target ) = @_;
    my ( $cc, @flags ) =
        ( $target eq 'linux' )
      ? ( $linux_cc, @linux_cflags )
      : ( $cross_cc, @cross_cflags );
    Util::run( $cc, @flags, '-c', "targets/wasm-module/$module.cpp",
        '-o', "$bd/c_module.o" );
}

sub compile_c_host {
    my ( $module, $bd, $target ) = @_;
    my ( $cc, @flags ) =
        ( $target eq 'linux' )
      ? ( $linux_cc, @linux_cflags, '-DTARGET_LINUX' )
      : ( $cross_cc, @cross_cflags, '-DTARGET_FAIL' );
    Util::run( $cc, @flags, '-c', 'targets/c-host/c_host.c', '-o', "$bd/c_host.o" );
    copy( 'targets/c-host/c_host.c', "$bd/module_host.c" );
}

sub compile_startup {
    my ( $module, $bd, $target ) = @_;
    return unless $target eq 'fail';
    Util::run( $cross_cc, 'targets/startup.s', '-I./targets/wasm-host', @cross_cflags,
        '-c', '-o', "$bd/startup.o" );
}

sub compile_syscalls {
    my ( $module, $bd, $target ) = @_;
    if ( $target eq 'fail' ) {
        Util::run( $cross_cc, 'targets/syscalls.c',
            '-I./targets/wasm-host', @cross_cflags,
            '-c', '-o', "$bd/syscalls.o" );
    }
    elsif ( $target eq 'linux-baremetal' ) {
        Util::run( $cross_cc, 'targets/syscalls.c', @linux_baremetal_cflags,
            '-c', '-o', "$bd/syscalls.o" );
    }

    # Linux needs neither startup nor syscall stubs
}

sub link_wasm {
    my ( $module, $bd, $target ) = @_;
    if ( $target eq 'fail' ) {
        Util::run(
            $cross_cc,           '-Wl,-T',
            'targets/linker.ld', "$bd/system.o",
            "$bd/startup.o",     "$bd/syscalls.o",
            "$bd/libiwasm.a",    @cross_ldflags_base,
            '-o',                "$bd/system.elf"
        );
    }
    elsif ( $target eq 'linux' ) {
        Util::run( $linux_cc, "$bd/system.o", "$bd/libiwasm.a", @linux_ldflags_base,
            '-o', "$bd/system.elf" );
    }
    elsif ( $target eq 'linux-baremetal' ) {
        Util::run( $cross_cc, "$bd/system.o", "$bd/syscalls.o", "$bd/libiwasm.a",
            @baremetal_ldflags_base, '-o', "$bd/system.elf" );
    }
    else {
        die "Unknown target '$target'\n";
    }
}

sub link_c {
    my ( $module, $bd, $target ) = @_;
    if ( $target eq 'fail' ) {
        Util::run(
            $cross_cc,             '-Wl,-T',
            'targets/linker.ld',   "$bd/c_host.o",
            "$bd/startup.o",       "$bd/c_module.o",
            @cross_ldflags_nowasm, '-o',
            "$bd/system.elf"
        );
    }
    elsif ( $target eq 'linux' ) {
        Util::run( $linux_cc, "$bd/c_host.o", "$bd/c_module.o", @linux_ldflags_nowasm,
            '-o', "$bd/system.elf" );
    }
    else {
        die "C mode is not supported for target '$target'\n";
    }
}

sub build_iso {
    my ( $module, $bd ) = @_;
    make_path("$bd/grub/boot/grub");
    copy( 'targets/grub.cfg', "$bd/grub/boot/grub/grub.cfg" )
      or die "failed to copy grub.cfg: $!\n";
    copy( "$bd/system.elf", "$bd/grub/boot/system.elf" )
      or die "failed to copy system.elf: $!\n";
    Util::run( 'grub-mkrescue', '-o', "$bd/system.iso", "$bd/grub" );
}