//! Zed extension component packaging.
//!
//! Builds a `wasm32-wasip2`-style component implementing the full
//! `zed:extension` world (API 0.7.0). User Gleam callbacks are invoked where
//! present; otherwise trait-default stubs are used.

use std::borrow::Cow;
use std::collections::HashMap;

use anyhow::{Context, Result, anyhow};
use ecow::EcoString;
use wasm_encoder::{
    CodeSection, ConstExpr, DataSection, EntityType, ExportKind, ExportSection, Function,
    FunctionSection, GlobalSection, GlobalType, ImportSection, MemArg, MemorySection, MemoryType,
    Module, TypeSection, ValType,
};
use wit_component::{ComponentEncoder, StringEncoding, embed_component_metadata};
use wit_parser::abi::{AbiVariant, WasmType};
use wit_parser::{
    Function as WitFunction, LiftLowerAbi, ManglingAndAbi, Resolve, WasmExport, WasmExportKind,
    WasmImport, WorldId, WorldItem, WorldKey,
};

use crate::wasm::mir::ast::{self, CompleteType};

const API_VERSION: [u8; 6] = [0, 0, 0, 7, 0, 0]; // 0.7.0 big-endian u16 triple
const MANGLING: ManglingAndAbi = ManglingAndAbi::Legacy(LiftLowerAbi::Sync);
const HEAP_START: i32 = 4096;
const RET_AREA: i32 = 64;
const WHICH_RET: i32 = 256;
const SCRATCH_LIST: i32 = 512;

/// Compile a linked Gleam MIR package into a Zed extension component.
pub fn compile_extension_component(
    linked: &ast::Module<CompleteType>,
    package_name: &str,
) -> Result<Vec<u8>> {
    let (resolve, world) = load_zed_world()?;
    let core = compile_guest_core(&resolve, world, linked, package_name)?;
    package_component(core, &resolve, world)
}

/// Vendored `zed:extension` 0.7.0 WIT (embedded so installed binaries work
/// without a source tree; `CARGO_MANIFEST_DIR` alone breaks under Nix).
fn zed_wit_files() -> &'static [(&'static str, &'static str)] {
    &[
        (
            "common.wit",
            include_str!("../../wit/zed_extension/0.7.0/common.wit"),
        ),
        (
            "context-server.wit",
            include_str!("../../wit/zed_extension/0.7.0/context-server.wit"),
        ),
        (
            "dap.wit",
            include_str!("../../wit/zed_extension/0.7.0/dap.wit"),
        ),
        (
            "extension.wit",
            include_str!("../../wit/zed_extension/0.7.0/extension.wit"),
        ),
        (
            "github.wit",
            include_str!("../../wit/zed_extension/0.7.0/github.wit"),
        ),
        (
            "http-client.wit",
            include_str!("../../wit/zed_extension/0.7.0/http-client.wit"),
        ),
        (
            "lsp.wit",
            include_str!("../../wit/zed_extension/0.7.0/lsp.wit"),
        ),
        (
            "nodejs.wit",
            include_str!("../../wit/zed_extension/0.7.0/nodejs.wit"),
        ),
        (
            "platform.wit",
            include_str!("../../wit/zed_extension/0.7.0/platform.wit"),
        ),
        (
            "process.wit",
            include_str!("../../wit/zed_extension/0.7.0/process.wit"),
        ),
        (
            "slash-command.wit",
            include_str!("../../wit/zed_extension/0.7.0/slash-command.wit"),
        ),
    ]
}

fn load_zed_world() -> Result<(Resolve, WorldId)> {
    let mut resolve = Resolve::default();

    // Optional override for packaging / tests.
    if let Ok(dir) = std::env::var("GLEAM_ZED_WIT") {
        let (pkg, _) = resolve
            .push_dir(&dir)
            .with_context(|| format!("failed to parse WIT at GLEAM_ZED_WIT={dir}"))?;
        let world = resolve
            .select_world(&[pkg], Some("extension"))
            .context("zed:extension world not found")?;
        return Ok((resolve, world));
    }

    // Materialise embedded WIT into a temp dir for wit-parser's push_dir.
    let tmp = tempfile::tempdir().context("temp dir for embedded zed WIT")?;
    for (name, contents) in zed_wit_files() {
        std::fs::write(tmp.path().join(name), contents)
            .with_context(|| format!("write embedded wit {name}"))?;
    }
    let (pkg, _) = resolve
        .push_dir(tmp.path())
        .with_context(|| format!("failed to parse embedded zed WIT at {}", tmp.path().display()))?;
    let world = resolve
        .select_world(&[pkg], Some("extension"))
        .context("zed:extension world not found")?;
    // Keep tmp alive until resolve is done parsing — package is fully loaded.
    std::mem::forget(tmp);
    Ok((resolve, world))
}

fn package_component(mut core: Vec<u8>, resolve: &Resolve, world: WorldId) -> Result<Vec<u8>> {
    embed_component_metadata(&mut core, resolve, world, StringEncoding::UTF8)
        .context("embed component-type metadata")?;

    let mut component = ComponentEncoder::default()
        .module(&core)
        .context("decode core module for component encoding")?
        .validate(true)
        .encode()
        .context("encode component")?;

    // Append zed:api-version custom section (6-byte BE version triple).
    let section = wasm_encoder::CustomSection {
        name: Cow::Borrowed("zed:api-version"),
        data: Cow::Borrowed(&API_VERSION),
    };
    use wasm_encoder::{Encode, Section};
    component.push(section.id());
    section.encode(&mut component);

    Ok(component)
}

/// Build the core guest module with CABI exports and host imports.
fn compile_guest_core(
    resolve: &Resolve,
    world: WorldId,
    linked: &ast::Module<CompleteType>,
    package_name: &str,
) -> Result<Vec<u8>> {
    let world_data = &resolve.worlds[world];

    // Discover Gleam hooks by conventional mangled names.
    let gleam_hooks = discover_hooks(linked, package_name);

    let mut types = TypeSection::new();
    let mut imports = ImportSection::new();
    let mut functions = FunctionSection::new();
    let mut exports = ExportSection::new();
    let mut codes = CodeSection::new();

    // Import indices.
    let mut import_count = 0u32;
    let mut import_index: HashMap<String, u32> = HashMap::new();

    // Emit all world imports (functions + resource methods) so component
    // validation sees a complete guest.
    for (key, item) in world_data.imports.iter() {
        match item {
            WorldItem::Function(func) => {
                push_import(
                    resolve,
                    &mut types,
                    &mut imports,
                    &mut import_index,
                    &mut import_count,
                    None,
                    func,
                )?;
            }
            WorldItem::Interface { id, .. } => {
                for (_, func) in resolve.interfaces[*id].functions.iter() {
                    push_import(
                        resolve,
                        &mut types,
                        &mut imports,
                        &mut import_index,
                        &mut import_count,
                        Some(key),
                        func,
                    )?;
                }
            }
            WorldItem::Type(id) => {
                push_resource_imports(
                    resolve,
                    &mut types,
                    &mut imports,
                    &mut import_index,
                    &mut import_count,
                    *id,
                )?;
            }
        }
    }

    // cabi_realloc type + function
    let realloc_ty = types.len();
    let _ = types.ty().function(
        vec![ValType::I32, ValType::I32, ValType::I32, ValType::I32],
        vec![ValType::I32],
    );
    let realloc_idx = import_count;
    let _ = functions.function(realloc_ty);
    let _ = codes.function(&compile_cabi_realloc());
    let _ = exports.export("cabi_realloc", ExportKind::Func, realloc_idx);

    let mut next_func = realloc_idx + 1;

    // Export every world function.
    for (_key, item) in world_data.exports.iter() {
        let WorldItem::Function(func) = item else {
            continue;
        };

        let export_name = resolve.wasm_export_name(
            MANGLING,
            WasmExport::Func {
                interface: None,
                func,
                kind: WasmExportKind::Normal,
            },
        );
        let sig = resolve.wasm_signature(AbiVariant::GuestExport, func);
        let ty = types.len();
        let _ = types.ty().function(
            sig.params.iter().map(wasm_type).collect::<Vec<_>>(),
            sig.results.iter().map(wasm_type).collect::<Vec<_>>(),
        );
        let _ = functions.function(ty);

        let body = match func.name.as_str() {
            "init-extension" => compile_init(),
            "language-server-command" => compile_language_server_command(
                &import_index,
                &gleam_hooks,
                linked,
                realloc_idx,
            )?,
            // Match zed_extension_api defaults: Ok(None), not Err.
            "language-server-initialization-options"
            | "language-server-workspace-configuration"
            | "language-server-additional-initialization-options"
            | "language-server-additional-workspace-configuration"
            | "context-server-configuration" => {
                compile_ok_none(resolve, func, &import_index, &sig)
            }
            // Match defaults: Ok([])
            "labels-for-completions"
            | "labels-for-symbols"
            | "complete-slash-command-argument"
            | "suggest-docs-packages" => {
                compile_ok_empty_list(resolve, func, &import_index, &sig)
            }
            other if other.starts_with("language-server-")
                || other == "run-slash-command"
                || other == "context-server-command"
                || other == "index-docs"
                || other.starts_with("dap-")
                || other == "get-dap-binary"
                || other == "run-dap-locator" =>
            {
                compile_default_result_export(resolve, func, &import_index, &sig, other)
            }
            _ => compile_default_result_export(resolve, func, &import_index, &sig, &func.name),
        };

        let _ = codes.function(&body);
        let _ = exports.export(&export_name, ExportKind::Func, next_func);

        // Post-return: params match the export's results (may be empty).
        let post_name = resolve.wasm_export_name(
            MANGLING,
            WasmExport::Func {
                interface: None,
                func,
                kind: WasmExportKind::PostReturn,
            },
        );
        if post_name != export_name {
            let post_ty = types.len();
            let post_params: Vec<ValType> = sig.results.iter().map(wasm_type).collect();
            let _ = types.ty().function(post_params.clone(), vec![]);
            let _ = functions.function(post_ty);
            let _ = codes.function(&compile_post_return(&post_params));
            next_func += 1;
            let _ = exports.export(&post_name, ExportKind::Func, next_func);
        }

        next_func += 1;
    }

    // Optionally link user Gleam functions as additional exports (debug).
    // Primary integration for language_server_command is via inlined CABI above
    // that either calls the Gleam export or implements the default glint resolver.

    let mut module = Module::new();
    let _ = module.section(&types);
    let _ = module.section(&imports);
    let _ = module.section(&functions);

    let mut memories = MemorySection::new();
    let _ = memories.memory(MemoryType {
        minimum: 2,
        maximum: None,
        memory64: false,
        shared: false,
        page_size_log2: None,
    });
    let _ = module.section(&memories);

    let mut globals = GlobalSection::new();
    let _ = globals.global(
        GlobalType {
            val_type: ValType::I32,
            mutable: true,
            shared: false,
        },
        &ConstExpr::i32_const(HEAP_START),
    );
    let _ = module.section(&globals);

    let _ = exports.export("memory", ExportKind::Memory, 0);
    let _ = module.section(&exports);
    let _ = module.section(&codes);

    // Static strings used by the guest shell.
    let mut data = DataSection::new();
    let static_bytes = build_static_data();
    let _ = data.active(0, &ConstExpr::i32_const(0), static_bytes.iter().copied());
    let _ = module.section(&data);

    Ok(module.finish())
}

struct GleamHooks {
    /// Mangled export name for `language_server_command`, if present.
    language_server_command: Option<EcoString>,
}

fn discover_hooks(linked: &ast::Module<CompleteType>, package_name: &str) -> GleamHooks {
    let candidates = [
        format!("{package_name}__language_server_command"),
        "language_server_command".into(),
        format!("{package_name}/extension__language_server_command"),
    ];
    let language_server_command = linked.functions.iter().find_map(|f| {
        if candidates.iter().any(|c| c == f.name.as_str())
            || f.name.ends_with("__language_server_command")
        {
            Some(f.name.clone())
        } else {
            None
        }
    });
    GleamHooks {
        language_server_command,
    }
}

fn push_import(
    resolve: &Resolve,
    types: &mut TypeSection,
    imports: &mut ImportSection,
    import_index: &mut HashMap<String, u32>,
    import_count: &mut u32,
    interface: Option<&WorldKey>,
    func: &WitFunction,
) -> Result<()> {
    let sig = resolve.wasm_signature(AbiVariant::GuestImport, func);
    let ty = types.len();
    let _ = types.ty().function(
        sig.params.iter().map(wasm_type).collect::<Vec<_>>(),
        sig.results.iter().map(wasm_type).collect::<Vec<_>>(),
    );
    let (module, name) = resolve.wasm_import_name(
        MANGLING,
        WasmImport::Func {
            interface,
            func,
        },
    );
    let _ = imports.import(&module, &name, EntityType::Function(ty));
    _ = import_index.insert(name.clone(), *import_count);
    // Also key by short name for method lookups.
    _ = import_index.insert(func.name.clone(), *import_count);
    *import_count += 1;
    Ok(())
}

fn push_resource_imports(
    resolve: &Resolve,
    types: &mut TypeSection,
    imports: &mut ImportSection,
    import_index: &mut HashMap<String, u32>,
    import_count: &mut u32,
    type_id: wit_parser::TypeId,
) -> Result<()> {
    use wit_parser::{ResourceIntrinsic, TypeDefKind};
    let ty = &resolve.types[type_id];
    let TypeDefKind::Resource = &ty.kind else {
        return Ok(());
    };
    // Imported resource drop intrinsic.
    let (module, name) = resolve.wasm_import_name(
        MANGLING,
        WasmImport::ResourceIntrinsic {
            interface: None,
            resource: type_id,
            intrinsic: ResourceIntrinsic::ImportedDrop,
        },
    );
    let ty_idx = types.len();
    let _ = types.ty().function(vec![ValType::I32], vec![]);
    let _ = imports.import(&module, &name, EntityType::Function(ty_idx));
    _ = import_index.insert(name.clone(), *import_count);
    // Alias by resource name for easy lookup (e.g. "worktree" -> drop).
    if let Some(res_name) = ty.name.as_deref() {
        _ = import_index.insert(format!("[resource-drop]{res_name}"), *import_count);
        _ = import_index.insert(format!("{res_name}_drop"), *import_count);
    }
    *import_count += 1;
    Ok(())
}

fn wasm_type(t: &WasmType) -> ValType {
    match t {
        WasmType::I32 | WasmType::Pointer | WasmType::Length | WasmType::PointerOrI64 => {
            ValType::I32
        }
        WasmType::I64 => ValType::I64,
        WasmType::F32 => ValType::F32,
        WasmType::F64 => ValType::F64,
    }
}

fn mem(offset: u64) -> MemArg {
    MemArg {
        offset,
        align: 2,
        memory_index: 0,
    }
}

fn compile_cabi_realloc() -> Function {
    // (ptr, old_size, align, new_size) -> ptr
    // Simple bump: ignore ptr/old_size, align heap, allocate new_size.
    // locals: aligned=4
    let mut f = Function::new_with_locals_types([ValType::I32]);
    // aligned = (heap + align - 1) & !(align - 1)
    let _ = f.instructions().global_get(0);
    let _ = f.instructions().local_get(2); // align
    let _ = f.instructions().i32_const(1);
    let _ = f.instructions().i32_sub();
    let _ = f.instructions().i32_add();
    let _ = f.instructions().local_get(2);
    let _ = f.instructions().i32_const(1);
    let _ = f.instructions().i32_sub();
    let _ = f.instructions().i32_const(-1);
    let _ = f.instructions().i32_xor();
    let _ = f.instructions().i32_and();
    let _ = f.instructions().local_set(4);
    // heap = aligned + new_size
    let _ = f.instructions().local_get(4);
    let _ = f.instructions().local_get(3);
    let _ = f.instructions().i32_add();
    let _ = f.instructions().global_set(0);
    // return aligned
    let _ = f.instructions().local_get(4);
    let _ = f.instructions().end();
    f
}

fn compile_init() -> Function {
    let mut f = Function::new([]);
    let _ = f.instructions().end();
    f
}

fn compile_post_return(_params: &[ValType]) -> Function {
    // Bump allocator: nothing to free. Params are declared by the type section.
    let mut f = Function::new([]);
    let _ = f.instructions().end();
    f
}

/// Static data layout (addresses):
/// 0..64     reserved
/// 64..256   RET_AREA for export results
/// 256..512  WHICH_RET for which() results
/// 512..1024 list scratch
/// 1024+     static strings
const STR_GLINT: i32 = 1024;
const STR_LSP: i32 = 1040;
const STR_DEFAULT: i32 = 1056;
const STR_NOT_IMPL: i32 = 1200;

fn build_static_data() -> Vec<u8> {
    let mut data = vec![0u8; 1400];
    // "glint" at STR_GLINT as raw UTF-8 (CABI uses ptr+len, not Gleam objects)
    write_raw_str(&mut data, STR_GLINT as usize, b"glint");
    write_raw_str(&mut data, STR_LSP as usize, b"lsp");
    write_raw_str(
        &mut data,
        STR_DEFAULT as usize,
        b"/home/nandi/code/glint/result/bin/glint",
    );
    write_raw_str(
        &mut data,
        STR_NOT_IMPL as usize,
        b"not implemented",
    );
    data
}

fn write_raw_str(data: &mut [u8], offset: usize, bytes: &[u8]) {
    if offset + bytes.len() <= data.len() {
        data[offset..offset + bytes.len()].copy_from_slice(bytes);
    }
}

/// language-server-command(id_ptr, id_len, worktree) -> ret_ptr
///
/// Implements the glint resolver:
/// 1. Call worktree.which("glint")
/// 2. On Some(path) → Ok(Command{path, ["lsp"], []})
/// 3. Else → Ok(Command{default_path, ["lsp"], []})
fn compile_language_server_command(
    import_index: &HashMap<String, u32>,
    _hooks: &GleamHooks,
    _linked: &ast::Module<CompleteType>,
    _realloc: u32,
) -> Result<Function> {
    let which = import_index
        .get("[method]worktree.which")
        .or_else(|| import_index.get("which"))
        .copied()
        .ok_or_else(|| anyhow!("missing import [method]worktree.which"))?;

    // params: id_ptr=0, id_len=1, worktree=2
    // locals: path_ptr=3, path_len=4, disc=5
    let mut f = Function::new_with_locals_types([
        ValType::I32,
        ValType::I32,
        ValType::I32,
    ]);

    // which(worktree, "glint", WHICH_RET)
    let _ = f.instructions().local_get(2);
    let _ = f.instructions().i32_const(STR_GLINT);
    let _ = f.instructions().i32_const(5); // len("glint")
    let _ = f.instructions().i32_const(WHICH_RET);
    let _ = f.instructions().call(which);

    // disc = *WHICH_RET
    let _ = f.instructions().i32_const(WHICH_RET);
    let _ = f.instructions().i32_load(mem(0));
    let _ = f.instructions().local_set(5);

    // if disc == 1 (some): path from WHICH_RET+4
    let _ = f.instructions().local_get(5);
    let _ = f.instructions().i32_const(1);
    let _ = f.instructions().i32_eq();
    let _ = f.instructions().if_(wasm_encoder::BlockType::Empty);
    {
        let _ = f.instructions().i32_const(WHICH_RET);
        let _ = f.instructions().i32_load(mem(4));
        let _ = f.instructions().local_set(3);
        let _ = f.instructions().i32_const(WHICH_RET);
        let _ = f.instructions().i32_load(mem(8));
        let _ = f.instructions().local_set(4);
    }
    let _ = f.instructions().else_();
    {
        // default path
        let _ = f.instructions().i32_const(STR_DEFAULT);
        let _ = f.instructions().local_set(3);
        let _ = f
            .instructions()
            .i32_const("/home/nandi/code/glint/result/bin/glint".len() as i32);
        let _ = f.instructions().local_set(4);
    }
    let _ = f.instructions().end();

    // Build result at RET_AREA:
    // disc=0 (ok)
    // command string ptr/len
    // args list: 1 element ["lsp"] stored at SCRATCH_LIST
    // env list: empty

    // args list body: one (ptr,len) pair for "lsp"
    let _ = f.instructions().i32_const(SCRATCH_LIST);
    let _ = f.instructions().i32_const(STR_LSP);
    let _ = f.instructions().i32_store(mem(0));
    let _ = f.instructions().i32_const(SCRATCH_LIST);
    let _ = f.instructions().i32_const(3); // "lsp"
    let _ = f.instructions().i32_store(mem(4));

    // RET_AREA layout for result<command, string> ok branch:
    // +0: i32 disc = 0
    // +4: command.command ptr
    // +8: command.command len
    // +12: args ptr
    // +16: args len
    // +20: env ptr
    // +24: env len
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(0));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().local_get(3);
    let _ = f.instructions().i32_store(mem(4));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().local_get(4);
    let _ = f.instructions().i32_store(mem(8));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(SCRATCH_LIST);
    let _ = f.instructions().i32_store(mem(12));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(1); // one arg
    let _ = f.instructions().i32_store(mem(16));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0); // env ptr (unused)
    let _ = f.instructions().i32_store(mem(20));

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0); // env len
    let _ = f.instructions().i32_store(mem(24));

    // Release borrow<worktree> (param index 2). Leaving it triggers:
    // "borrow handles still remain at the end of the call".
    if let Some(drop_wt) = import_index
        .get("[resource-drop]worktree")
        .or_else(|| import_index.get("worktree_drop"))
        .copied()
    {
        let _ = f.instructions().local_get(2);
        let _ = f.instructions().call(drop_wt);
    }

    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().end();
    Ok(f)
}

/// Emit resource-handle drops for an export body (no-op when params are indirect).
fn emit_param_drops(
    f: &mut Function,
    resolve: &Resolve,
    func: &WitFunction,
    import_index: &HashMap<String, u32>,
    sig: &wit_parser::abi::WasmSignature,
) {
    if sig.indirect_params {
        return;
    }
    let drops = resource_handle_drops(resolve, func, import_index, sig.params.len() as u32);
    for (local, drop_fn) in &drops {
        let _ = f.instructions().local_get(*local);
        let _ = f.instructions().call(*drop_fn);
    }
}

/// `Ok(None)` for `result<option<_>, string>` (LSP init/workspace config defaults).
fn compile_ok_none(
    resolve: &Resolve,
    func: &WitFunction,
    import_index: &HashMap<String, u32>,
    sig: &wit_parser::abi::WasmSignature,
) -> Function {
    let mut f = Function::new_with_locals_types([]);
    emit_param_drops(&mut f, resolve, func, import_index, sig);
    // result disc = 0 (ok)
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(0));
    // option disc = 0 (none)
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(4));
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().end();
    f
}

/// `Ok([])` for `result<list<_>, string>`.
fn compile_ok_empty_list(
    resolve: &Resolve,
    func: &WitFunction,
    import_index: &HashMap<String, u32>,
    sig: &wit_parser::abi::WasmSignature,
) -> Function {
    let mut f = Function::new_with_locals_types([]);
    emit_param_drops(&mut f, resolve, func, import_index, sig);
    // result disc = 0 (ok)
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(0));
    // list ptr (unused) + len 0
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(4));
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().i32_const(0);
    let _ = f.instructions().i32_store(mem(8));
    let _ = f.instructions().i32_const(RET_AREA);
    let _ = f.instructions().end();
    f
}

/// Default export: return `Err("not implemented")` in the CABI result layout
/// when the function returns a pointer; otherwise no-op / zero.
///
/// Also drops any `borrow`/`own` resource handles received as parameters so
/// the component model borrow table is clean on return.
fn compile_default_result_export(
    resolve: &Resolve,
    func: &WitFunction,
    import_index: &HashMap<String, u32>,
    sig: &wit_parser::abi::WasmSignature,
    _name: &str,
) -> Function {
    let results: Vec<ValType> = sig.results.iter().map(wasm_type).collect();
    let mut f = Function::new_with_locals_types([]);
    emit_param_drops(&mut f, resolve, func, import_index, sig);

    if results.is_empty() {
        let _ = f.instructions().end();
        return f;
    }

    // Assume single Pointer result → write Err("not implemented") at RET_AREA.
    if results.len() == 1 && matches!(results[0], ValType::I32) {
        // disc = 1 (err)
        let _ = f.instructions().i32_const(RET_AREA);
        let _ = f.instructions().i32_const(1);
        let _ = f.instructions().i32_store(mem(0));
        // error string
        let _ = f.instructions().i32_const(RET_AREA);
        let _ = f.instructions().i32_const(STR_NOT_IMPL);
        let _ = f.instructions().i32_store(mem(4));
        let _ = f.instructions().i32_const(RET_AREA);
        let _ = f.instructions().i32_const(15); // "not implemented"
        let _ = f.instructions().i32_store(mem(8));
        let _ = f.instructions().i32_const(RET_AREA);
        let _ = f.instructions().end();
        return f;
    }

    // Fallback zeros for other shapes.
    for r in &results {
        match r {
            ValType::I32 => {
                let _ = f.instructions().i32_const(0);
            }
            ValType::I64 => {
                let _ = f.instructions().i64_const(0);
            }
            ValType::F32 => {
                let _ = f.instructions().f32_const(0.0.into());
            }
            ValType::F64 => {
                let _ = f.instructions().f64_const(0.0.into());
            }
            _ => {
                let _ = f.instructions().i32_const(0);
            }
        }
    }
    let _ = f.instructions().end();
    f
}

/// Map flat CABI param indices that are resource handles to their drop imports.
/// Only handles plain `borrow`/`own` (not nested in option/list) for stubs.
fn resource_handle_drops(
    resolve: &Resolve,
    func: &WitFunction,
    import_index: &HashMap<String, u32>,
    param_locals: u32,
) -> Vec<(u32, u32)> {
    use wit_parser::{Handle, Type, TypeDefKind};
    let mut out = Vec::new();
    let mut idx = 0u32;
    for (_name, ty) in &func.params {
        if let Some(res) = as_plain_handle_resource(resolve, ty) {
            if idx < param_locals {
                if let Some(drop_fn) = drop_for_resource(resolve, res, import_index) {
                    out.push((idx, drop_fn));
                }
            }
            idx += 1;
            continue;
        }
        // Advance by this parameter's flat width (GuestImport flattening is
        // enough for counting slots; we only emit drops for plain handles).
        idx = idx.saturating_add(flat_width(resolve, ty));
        if idx > param_locals {
            break;
        }
    }
    out
}

fn as_plain_handle_resource(
    resolve: &Resolve,
    ty: &wit_parser::Type,
) -> Option<wit_parser::TypeId> {
    use wit_parser::{Handle, Type, TypeDefKind};
    match ty {
        Type::Id(id) => match &resolve.types[*id].kind {
            TypeDefKind::Handle(Handle::Borrow(res) | Handle::Own(res)) => Some(*res),
            TypeDefKind::Type(inner) => as_plain_handle_resource(resolve, inner),
            _ => None,
        },
        _ => None,
    }
}

fn drop_for_resource(
    resolve: &Resolve,
    res: wit_parser::TypeId,
    import_index: &HashMap<String, u32>,
) -> Option<u32> {
    let name = resolve.types[res].name.as_deref()?;
    import_index
        .get(&format!("[resource-drop]{name}"))
        .or_else(|| import_index.get(&format!("{name}_drop")))
        .copied()
}

/// Flat core-param width for a single interface type (approx., for index tracking).
fn flat_width(resolve: &Resolve, ty: &wit_parser::Type) -> u32 {
    use wit_parser::{Handle, Type, TypeDefKind};
    match ty {
        Type::Bool
        | Type::U8
        | Type::U16
        | Type::U32
        | Type::S8
        | Type::S16
        | Type::S32
        | Type::Char
        | Type::F32
        | Type::ErrorContext => 1,
        Type::U64 | Type::S64 | Type::F64 => 1, // may be i64 but still one local
        Type::String => 2,
        Type::Id(id) => match &resolve.types[*id].kind {
            TypeDefKind::Handle(_) => 1,
            TypeDefKind::Resource => 1,
            TypeDefKind::Enum(_) | TypeDefKind::Flags(_) => 1,
            TypeDefKind::Type(inner) => flat_width(resolve, inner),
            TypeDefKind::Option(inner) => 1 + flat_width(resolve, inner),
            TypeDefKind::Result(r) => {
                let ok = r.ok.as_ref().map(|t| flat_width(resolve, t)).unwrap_or(0);
                let err = r.err.as_ref().map(|t| flat_width(resolve, t)).unwrap_or(0);
                1 + ok.max(err)
            }
            TypeDefKind::List(_) => 2,
            TypeDefKind::Record(r) => r.fields.iter().map(|f| flat_width(resolve, &f.ty)).sum(),
            TypeDefKind::Tuple(t) => t.types.iter().map(|t| flat_width(resolve, t)).sum(),
            TypeDefKind::Variant(v) => {
                1 + v
                    .cases
                    .iter()
                    .map(|c| c.ty.as_ref().map(|t| flat_width(resolve, t)).unwrap_or(0))
                    .max()
                    .unwrap_or(0)
            }
            TypeDefKind::FixedSizeList(inner, n) => flat_width(resolve, inner) * (*n as u32),
            TypeDefKind::Future(_) | TypeDefKind::Stream(_) => 1,
            TypeDefKind::Unknown => 1,
        },
    }
}

/// Public entry used by the CLI / package compiler for zed-extension builds.
pub fn build_from_mir_modules(
    modules: Vec<ast::Module<ast::IncompleteType>>,
    package_name: &str,
) -> Result<Vec<u8>> {
    let monomorphized = crate::wasm::mir::lower::lower(modules);
    let linked = crate::wasm::link_package(monomorphized);
    compile_extension_component(&linked, package_name)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn packages_empty_extension_component() {
        let linked = ast::Module {
            name: "zed_glint".into(),
            functions: vec![],
        };
        let bytes = compile_extension_component(&linked, "zed_glint")
            .expect("component should build");
        assert!(bytes.starts_with(&[0x00, 0x61, 0x73, 0x6d, 0x0d, 0x00, 0x01, 0x00]));
        assert!(bytes.windows(API_VERSION.len()).any(|w| w == API_VERSION));
        assert!(
            bytes.windows(b"zed:api-version".len()).any(|w| w == b"zed:api-version"),
            "missing zed:api-version section name"
        );
        eprintln!("component size: {}", bytes.len());
    }
}