#!/usr/bin/env nu
# WebAssembly optimization and compilation module
# Supports Rust, C/C++, AssemblyScript, and Go compilation to WASM

# Main WASM compilation function
export def compile_wasm [config: record, dev: bool, verbose: bool] {
    if not ("wasm" in ($config | columns)) or not $config.wasm.enabled {
        if $verbose { print "📦 WASM not enabled, skipping..." }
        return { size: 0, modules: [] }
    }

    if $verbose { print "🦀 Compiling WebAssembly modules..." }

    let wasm_config = $config.wasm
    mut compiled_modules = []
    mut total_size = 0

    # Validate WASM build environment
    let toolchain_check = (validate_wasm_toolchain $wasm_config $verbose)
    if not $toolchain_check.valid {
        print $"❌ WASM toolchain validation failed: ($toolchain_check.error)"
        return { size: 0, modules: [] }
    }

    # Create WASM output directory
    mkdir $wasm_config.output_dir

    # Compile each WASM module
    for module in $wasm_config.modules {
        let module_result = (compile_wasm_module $module $wasm_config $dev $verbose)

        if $module_result.success {
            $compiled_modules = ($compiled_modules | append $module_result)
            $total_size = ($total_size + $module_result.size)
        } else {
            print $"❌ Failed to compile WASM module: ($module.name)"
        }
    }

    # Generate JavaScript loader if modules were compiled
    if ($compiled_modules | length) > 0 {
        let loader_result = (generate_wasm_loader $compiled_modules $wasm_config $verbose)
        if $loader_result.success {
            $total_size = ($total_size + $loader_result.size)
        }
    }

    if $verbose {
        print $"✅ Compiled ($compiled_modules | length) WASM modules, total size: ($total_size) bytes"
    }

    {
        size: $total_size,
        modules: $compiled_modules
    }
}

# Compile individual WASM module
def compile_wasm_module [module: record, wasm_config: record, dev: bool, verbose: bool] {
    let source_path = ($wasm_config.source_dir | path join $module.source)

    if not ($source_path | path exists) {
        return {
            name: $module.name,
            success: false,
            error: $"Source file not found: ($source_path)"
        }
    }

    let output_path = ($wasm_config.output_dir | path join $"($module.name).wasm")

    match $module.lang {
        "rust" => { compile_rust_wasm $module $source_path $output_path $wasm_config $dev $verbose }
        "c" | "cpp" | "c++" => { compile_c_wasm $module $source_path $output_path $wasm_config $dev $verbose }
        "assemblyscript" | "as" => { compile_assemblyscript_wasm $module $source_path $output_path $wasm_config $dev $verbose }
        "go" => { compile_go_wasm $module $source_path $output_path $wasm_config $dev $verbose }
        _ => {
            return {
                name: $module.name,
                success: false,
                error: $"Unsupported language: ($module.lang)"
            }
        }
    }
}

# Compile Rust to WASM
def compile_rust_wasm [module: record, source_path: string, output_path: string, wasm_config: record, dev: bool, verbose: bool] {
    let rust_config = if "rust" in ($wasm_config | columns) { $wasm_config.rust } else { {} }
    let toolchain = if "toolchain" in ($rust_config | columns) { $rust_config.toolchain } else { "wasm-pack" }
    let target = if "target" in ($rust_config | columns) { $rust_config.target } else { "web" }

    if $verbose { print $"  🦀 Compiling Rust module: ($module.name)" }

    match $toolchain {
        "wasm-pack" => {
            let build_args = if $dev { ["--dev"] } else { ["--release"] }
            let scope = if "scope" in ($rust_config | columns) { ["--scope", $rust_config.scope] } else { [] }

            try {
                let args = ["build", "--target", $target, "--out-dir", ($wasm_config.output_dir | path join $module.name)] ++ $build_args ++ $scope ++ [($source_path | path dirname)]
                run-external "wasm-pack" ...$args

                # Copy main WASM file to expected location
                let pkg_wasm = ($wasm_config.output_dir | path join $module.name | path join $"($module.name)_bg.wasm")
                if ($pkg_wasm | path exists) {
                    cp $pkg_wasm $output_path
                }

                let wasm_size = if ($output_path | path exists) {
                    (ls $output_path | get size | first | into int)
                } else { 0 }

                {
                    name: $module.name,
                    success: true,
                    path: $output_path,
                    size: $wasm_size,
                    lang: "rust",
                    loader_path: ($wasm_config.output_dir | path join $module.name | path join $"($module.name).js")
                }
            } catch { |e|
                {
                    name: $module.name,
                    success: false,
                    error: $"wasm-pack build failed: ($e.msg)"
                }
            }
        }
        "cargo" => {
            try {
                # Use cargo for basic WASM compilation
                let release_args = if $dev { [] } else { ["--release"] }
                let args = ["build", "--target", "wasm32-unknown-unknown"] ++ $release_args
                run-external "cargo" ...$args

                # Find compiled WASM file
                let target_dir = if $dev { "debug" } else { "release" }
                let cargo_wasm = ($source_path | path dirname | path join "target" "wasm32-unknown-unknown" $target_dir $"($module.name).wasm")

                if ($cargo_wasm | path exists) {
                    cp $cargo_wasm $output_path

                    let wasm_size = (ls $output_path | get size | first | into int)

                    {
                        name: $module.name,
                        success: true,
                        path: $output_path,
                        size: $wasm_size,
                        lang: "rust"
                    }
                } else {
                    {
                        name: $module.name,
                        success: false,
                        error: "Compiled WASM file not found"
                    }
                }
            } catch { |e|
                {
                    name: $module.name,
                    success: false,
                    error: $"cargo build failed: ($e.msg)"
                }
            }
        }
        _ => {
            {
                name: $module.name,
                success: false,
                error: $"Unsupported Rust toolchain: ($toolchain)"
            }
        }
    }
}

# Compile C/C++ to WASM using Emscripten
def compile_c_wasm [module: record, source_path: string, output_path: string, wasm_config: record, dev: bool, verbose: bool] {
    if $verbose { print $"  🔧 Compiling C/C++ module: ($module.name)" }

    let c_config = if "c" in ($wasm_config | columns) { $wasm_config.c } else { {} }
    let optimization = if $dev { "-O0" } else { "-O3" }
    let exports = if "exports" in ($c_config | columns) { $c_config.exports } else { ["main"] }
    let export_args = ($exports | each { |fn| ["-s", $"EXPORTED_FUNCTIONS=['_($fn)']"] } | flatten)

    try {
        let args = [$source_path, "-o", $output_path, $optimization, "-s", "WASM=1", "-s", "MODULARIZE=1", "-s", "EXPORT_NAME=initWasm"] ++ $export_args
        run-external "emcc" ...$args

        let wasm_size = if ($output_path | path exists) {
            (ls $output_path | get size | first | into int)
        } else { 0 }

        {
            name: $module.name,
            success: true,
            path: $output_path,
            size: $wasm_size,
            lang: "c"
        }
    } catch { |e|
        {
            name: $module.name,
            success: false,
            error: $"emcc compilation failed: ($e.msg)"
        }
    }
}

# Compile AssemblyScript to WASM
def compile_assemblyscript_wasm [module: record, source_path: string, output_path: string, wasm_config: record, dev: bool, verbose: bool] {
    if $verbose { print $"  📝 Compiling AssemblyScript module: ($module.name)" }

    let optimization = if $dev { "" } else { "--optimize" }
    let args = if $dev { [] } else { ["--optimize"] }

    try {
        let asc_args = [$source_path, "--outFile", $output_path, "--bindings", "esm"] ++ $args
        run-external "asc" ...$asc_args

        let wasm_size = if ($output_path | path exists) {
            (ls $output_path | get size | first | into int)
        } else { 0 }

        {
            name: $module.name,
            success: true,
            path: $output_path,
            size: $wasm_size,
            lang: "assemblyscript"
        }
    } catch { |e|
        {
            name: $module.name,
            success: false,
            error: $"AssemblyScript compilation failed: ($e.msg)"
        }
    }
}

# Compile Go to WASM
def compile_go_wasm [module: record, source_path: string, output_path: string, wasm_config: record, dev: bool, verbose: bool] {
    if $verbose { print $"  🐹 Compiling Go module: ($module.name)" }

    try {
        # Set environment variables for Go WASM compilation
        with-env {
            GOOS: "js",
            GOARCH: "wasm"
        } {
            run-external "go" "build" "-o" $output_path $source_path
        }

        let wasm_size = if ($output_path | path exists) {
            (ls $output_path | get size | first | into int)
        } else { 0 }

        {
            name: $module.name,
            success: true,
            path: $output_path,
            size: $wasm_size,
            lang: "go"
        }
    } catch { |e|
        {
            name: $module.name,
            success: false,
            error: $"Go compilation failed: ($e.msg)"
        }
    }
}

# Validate WASM toolchain
def validate_wasm_toolchain [wasm_config: record, verbose: bool] {
    mut missing_tools = []

    # Check for common WASM tools
    let tools_to_check = [
        { name: "wasm-pack", command: "wasm-pack", required_for: "Rust" },
        { name: "cargo", command: "cargo", required_for: "Rust" },
        { name: "emcc", command: "emcc", required_for: "C/C++" },
        { name: "asc", command: "asc", required_for: "AssemblyScript" },
        { name: "go", command: "go", required_for: "Go" },
        { name: "wasm-opt", command: "wasm-opt", required_for: "Optimization" }
    ]

    for tool in $tools_to_check {
        let available = (which $tool.command | length) > 0
        if not $available {
            $missing_tools = ($missing_tools | append $tool)
        } else if $verbose {
            print $"  ✅ ($tool.name) available"
        }
    }

    # Check which languages are actually needed
    let needed_languages = if "modules" in ($wasm_config | columns) {
        ($wasm_config.modules | get lang | uniq)
    } else { [] }

    let critical_missing = ($missing_tools | where { |tool|
        $needed_languages | any { |lang| $lang in $tool.required_for }
    })

    if ($critical_missing | length) > 0 {
        {
            valid: false,
            error: $"Missing required tools: ($critical_missing | get name | str join ', ')"
        }
    } else {
        { valid: true }
    }
}

# Generate JavaScript loader for WASM modules
def generate_wasm_loader [modules: list, wasm_config: record, verbose: bool] {
    if $verbose { print "  📜 Generating WASM loader..." }

    # Find framework root by looking for framework directory
    let current_dir = $env.PWD
    let framework_templates = ($current_dir | path join "framework" "templates")

    let template_path = ($framework_templates | path join "wasm-loader.js.template")
    if not ($template_path | path exists) {
        return {
            success: false,
            error: $"WASM loader template not found: ($template_path)"
        }
    }

    # Create module loading calls
    let module_calls = ($modules | each { |mod|
        let module_name = $mod.name
        $"            this.loadModule\('($module_name)', './($module_name).wasm'\),"
    } | str join "\n")

    # Load template and replace placeholder
    let template_content = (open $template_path)
    let loader_content = ($template_content | str replace "{{MODULE_PROMISES}}" $module_calls)

    let loader_path = ($wasm_config.output_dir | path join "wasm-loader.js")
    $loader_content | save --force $loader_path

    if $verbose { print $"  ✅ Generated WASM loader: ($loader_path)" }

    {
        success: true,
        path: $loader_path,
        size: ($loader_content | str length)
    }
}

# Optimize WASM with wasm-opt (if available)
export def optimize_wasm_binary [wasm_path: string, config: record] {
    let wasm_opt_available = (which wasm-opt | length) > 0

    if not $wasm_opt_available {
        return { optimized: false, reason: "wasm-opt not available" }
    }

    let wasm_config = if "wasm" in ($config | columns) { $config.wasm } else { {} }
    let optimization = if "optimization" in ($wasm_config | columns) { $wasm_config.optimization } else { {} }

    let opt_level = if "opt_level" in ($optimization | columns) { $optimization.opt_level } else { 3 }
    let shrink_level = if "shrink_level" in ($optimization | columns) { $optimization.shrink_level } else { 2 }

    let temp_path = $"($wasm_path).opt"

    try {
        let wasm_opt_args = [$wasm_path, "-o", $temp_path, $"-O($opt_level)", $"-s($shrink_level)", "--enable-simd", "--enable-bulk-memory"]
        run-external "wasm-opt" ...$wasm_opt_args

        # Replace original with optimized version
        mv $temp_path $wasm_path

        { optimized: true }
    } catch { |e|
        # Clean up temp file if it exists
        if ($temp_path | path exists) {
            rm $temp_path
        }

        { optimized: false, reason: $"wasm-opt failed: ($e.msg)" }
    }
}

# Get WASM module information
export def analyze_wasm [wasm_path: string] {
    if not ($wasm_path | path exists) {
        return null
    }

    let size = (ls $wasm_path | get size | first | into int)
    let size_kb = ($size / 1024 | math round -p 1)

    {
        path: $wasm_path,
        size: $size,
        size_kb: $size_kb,
        exists: true
    }
}