obeli-sk-utils 0.40.0

use crate::wasm_tools::ExIm;
use concepts::{
    FnName, FunctionExtension, FunctionMetadata, IfcFqnName, PackageExtension, PackageIfcFns,
    PkgFqn,
};
use const_format::formatcp;
use hashbrown::HashMap;
use id_arena::Arena;
use indexmap::IndexMap;
use semver::{BuildMetadata, Prerelease, Version};
use std::{fmt::Display, path::PathBuf};
use tracing::{error, warn};
use wit_component::WitPrinter;
use wit_parser::{
    Function, FunctionKind, Handle, Interface, InterfaceId, PackageId, PackageName, Param, Resolve,
    Span, Stability, Type, TypeDef, TypeDefKind, TypeOwner, UnresolvedPackageGroup, World,
    WorldItem, WorldKey,
};

const OBELISK_TYPES_VERSION_MAJOR: u64 = 5;
const OBELISK_TYPES_VERSION_MINOR: u64 = 0;
const OBELISK_TYPES_VERSION_PATCH: u64 = 0;
const OBELISK_TYPES_VERSION: &str = formatcp!(
    "{OBELISK_TYPES_VERSION_MAJOR}.{OBELISK_TYPES_VERSION_MINOR}.{OBELISK_TYPES_VERSION_PATCH}"
);
const OBELISK_TYPES_PACKAGE_NAME: &str = formatcp!("obelisk:types@{OBELISK_TYPES_VERSION}");

pub const WIT_OBELISK_LOG_PACKAGE: [&str; 3] = [
    "obelisk_log@1.0.0",
    "obelisk_log@1.0.0.wit",
    include_str!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/wit/obelisk_log@1.0.0/obelisk_log@1.0.0.wit"
    )),
];
const WIT_OBELISK_TYPES_PACKAGE_CONTENT: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/wit/obelisk_types@5.0.0/obelisk_types@5.0.0.wit"
));
pub const WIT_OBELISK_TYPES_PACKAGE: [&str; 3] = [
    "obelisk_types@5.0.0",
    "obelisk_types@5.0.0.wit",
    WIT_OBELISK_TYPES_PACKAGE_CONTENT,
];
pub const WIT_OBELISK_WORKFLOW_PACKAGE: [&str; 3] = [
    "obelisk_workflow@6.0.0",
    "obelisk_workflow@6.0.0.wit",
    include_str!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/wit/obelisk_workflow@6.0.0/obelisk_workflow@6.0.0.wit"
    )),
];

pub const WIT_OBELISK_WEBHOOK_PACKAGE: [&str; 3] = [
    "obelisk_webhook@6.0.0",
    "obelisk_webhook@6.0.0.wit",
    include_str!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/wit/obelisk_webhook@6.0.0/obelisk_webhook@6.0.0.wit"
    )),
];

pub(crate) fn wit(resolve: &Resolve, main_package: PackageId) -> Result<String, anyhow::Error> {
    // print all packages, with the main package as root, others as nested.
    let ids = packages_except_main(resolve, main_package, false);
    let mut printer = WitPrinter::new(OutputToString::default());
    printer.print(resolve, main_package, &ids)?;
    let wit = printer.output.to_string();
    Ok(wit)
}

#[expect(
    clippy::needless_pass_by_value,
    reason = "original resolve is consumed; the enriched resolve is returned"
)]
pub(crate) fn rebuild_resolve(
    exim: &ExIm,
    resolve: Resolve,
    main_package: PackageId,
) -> Result<(Resolve, PackageId), anyhow::Error> {
    let ids = packages_except_main(&resolve, main_package, true);
    let mut printer = WitPrinter::new(OutputToString::default());
    printer.print(&resolve, main_package, &ids)?;
    let wit = printer.output.to_string();

    let (mut resolve, main_pkg_id) = {
        let wit = replace_obelisk_types(&wit);
        let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit)?;
        let mut resolve = Resolve::new();
        let main_pkg_id = resolve.push_group(group)?;
        (resolve, main_pkg_id)
    };
    let world_id = resolve
        .select_world(&[main_pkg_id], None)
        .expect("default world must be found");
    let added_interfaces = add_extended_interfaces(exim.get_exports_hierarchy_ext(), &mut resolve)?;
    resolve
        .worlds
        .get_mut(world_id)
        .expect("id belongs to this resolve")
        .exports
        .extend(added_interfaces.into_iter().map(|ifc_id| {
            (
                WorldKey::Interface(ifc_id),
                WorldItem::Interface {
                    id: ifc_id,
                    stability: Stability::Unknown,
                    span: Span::default(),
                },
            )
        }));

    Ok((resolve, main_pkg_id))
}

fn add_extended_interfaces(
    exports_hierarchy_ext: &[PackageIfcFns],
    resolve: &mut Resolve,
) -> Result<Vec<InterfaceId>, semver::Error> {
    let mut added_interfaces = Vec::new();
    // Find necessary handles
    // Get obelisk:types
    let obelisk_types_package_name = PackageName {
        namespace: "obelisk".to_string(),
        name: "types".to_string(),
        version: Some(Version {
            major: OBELISK_TYPES_VERSION_MAJOR,
            minor: OBELISK_TYPES_VERSION_MINOR,
            patch: OBELISK_TYPES_VERSION_PATCH,
            pre: Prerelease::EMPTY,
            build: BuildMetadata::EMPTY,
        }),
    };
    let obelisk_types_pkg_id =
        if let Some(id) = resolve.package_names.get(&obelisk_types_package_name) {
            *id
        } else {
            let pkg = wit_parser::Package {
                name: obelisk_types_package_name,
                docs: wit_parser::Docs::default(),
                interfaces: IndexMap::default(),
                worlds: IndexMap::default(),
            };
            let package_name = pkg.name.clone();
            let ext_pkg_id = resolve.packages.alloc(pkg);
            resolve.package_names.insert(package_name, ext_pkg_id);
            ext_pkg_id
        };
    // Get obelisk:types/time@VERSION
    let time_ifc_id = *resolve.packages[obelisk_types_pkg_id]
        .interfaces
        .get("time")
        .expect("`time` interface was added");
    let time_ifc = &resolve.interfaces[time_ifc_id];

    let (execution_ifc_id, execution_ifc) = find_interface(
        &IfcFqnName::from_parts("obelisk", "types", "execution", Some(OBELISK_TYPES_VERSION)),
        resolve,
        &resolve.interfaces,
    )
    .expect(formatcp!(
        "{OBELISK_TYPES_PACKAGE_NAME} with interface `execution` must be found"
    ));

    let (join_set_ifc_id, join_set_ifc) = find_interface(
        &IfcFqnName::from_parts("obelisk", "types", "join-set", Some(OBELISK_TYPES_VERSION)),
        resolve,
        &resolve.interfaces,
    )
    .expect(formatcp!(
        "{OBELISK_TYPES_PACKAGE_NAME} with interface `join-set` must be found"
    ));

    // obelisk:types/execution@VERSION.{execution-id}
    let type_id_execution_id = {
        let actual_type_id = *execution_ifc
            .types
            .get("execution-id")
            .expect("`execution-id` must exist");
        // Create a reference to the type.
        resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(execution_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        })
    };
    // obelisk:types/execution@VERSION.{join-set}
    let (type_id_join_set_id, type_id_join_set_id_borrow_handle) = {
        let actual_type_id = *join_set_ifc
            .types
            .get("join-set")
            .expect("`join-set` must exist");
        // Create a reference to the type.
        let type_id_join_set_id = resolve.types.alloc(TypeDef {
            name: Some("join-set".to_string()),
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(join_set_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        });
        // Create a Handle::Borrow to the reference.
        let type_id_join_set_id_borrow_handle = resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Handle(Handle::Borrow(type_id_join_set_id)),
            owner: TypeOwner::Interface(join_set_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        });
        (type_id_join_set_id, type_id_join_set_id_borrow_handle)
    };
    // obelisk:types/execution.{await-next-extension-error}
    let type_id_await_next_extension_error = {
        let actual_type_id = *execution_ifc
            .types
            .get("await-next-extension-error")
            .expect("`await-next-extension-error` must exist");
        // Create a reference to the type.
        resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(execution_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        })
    };
    // obelisk:types/execution.{get-extension-error}
    let type_id_get_extension_error = {
        let actual_type_id = *execution_ifc
            .types
            .get("get-extension-error")
            .expect("`get-extension-error` must exist");
        // Create a reference to the type.
        resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(execution_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        })
    };
    // obelisk:types/execution.{stub-error}
    let type_id_stub_error = {
        let actual_type_id = *execution_ifc
            .types
            .get("stub-error")
            .expect("`stub-error` must exist");
        // Create a reference to the type.
        resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(execution_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        })
    };
    let type_id_await_next_err_part = type_id_await_next_extension_error;
    // obelisk:types/time.{schedule-at}
    let type_id_schedule_at = {
        let actual_type_id = *time_ifc
            .types
            .get("schedule-at")
            .expect("`schedule-at` must exist");
        // Create a reference to the type.
        resolve.types.alloc(TypeDef {
            name: None,
            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
            owner: TypeOwner::Interface(execution_ifc_id),
            docs: wit_parser::Docs::default(),
            stability: wit_parser::Stability::default(),
            span: Span::default(),
        })
    };

    for (pkg_fqn, ifc_to_fns) in get_ext_pkg_to_ifc_to_details_map(exports_hierarchy_ext) {
        let (orig_pkg_fqn, pkg_ext) = pkg_fqn
            .split_ext()
            .expect("`get_pkg_to_ifc_to_details_map` filtered by ext packages");
        let pkg_id = get_or_create_package(pkg_fqn, resolve)?;
        let (orig_pkg_id, _) = resolve
            .packages
            .iter()
            .find(|(_, found_pkg)| {
                from_wit_package_name_to_pkg_fqn(&found_pkg.name) == orig_pkg_fqn
            })
            .unwrap_or_else(|| {
                panic!("original package must be not found in resolve: {orig_pkg_fqn}")
            });

        for (ifc_fqn, fns) in ifc_to_fns {
            let orig_pkg = resolve.packages.get(orig_pkg_id).expect("id is fresh");

            let orig_ifc_id = *orig_pkg
                .interfaces
                .get(ifc_fqn.ifc_name())
                .unwrap_or_else(|| {
                    panic!("interface must be found in original resolve: {ifc_fqn}")
                });
            let orig_ifc = resolve
                .interfaces
                .get(orig_ifc_id)
                .expect("orig_ifc obtained from orig_resolve");

            let mut types = copy_or_refer_original_types(orig_ifc_id, orig_ifc, &mut resolve.types);
            match pkg_ext {
                PackageExtension::ObeliskExt => {
                    types.insert("execution-id".to_string(), type_id_execution_id);
                    types.insert("join-set".to_string(), type_id_join_set_id);
                    types.insert(
                        "await-next-extension-error".to_string(),
                        type_id_await_next_extension_error,
                    );
                    types.insert(
                        "get-extension-error".to_string(),
                        type_id_get_extension_error,
                    );
                }
                PackageExtension::ObeliskSchedule => {
                    types.insert("execution-id".to_string(), type_id_execution_id);
                    types.insert("schedule-at".to_string(), type_id_schedule_at);
                }
                PackageExtension::ObeliskStub => {
                    types.insert("execution-id".to_string(), type_id_execution_id);
                    types.insert("stub-error".to_string(), type_id_stub_error);
                }
            }

            let mut ifc = Interface {
                name: Some(ifc_fqn.ifc_name().to_string()),
                types,
                functions: IndexMap::default(),
                docs: wit_parser::Docs::default(),
                stability: Stability::default(),
                package: Some(pkg_id),
                span: Span::default(),
                clone_of: None,
            };
            for (fn_name, fn_meta) in fns {
                let (prefix, fn_ext) = fn_meta.split_extension().expect("filtered by ext package");
                let original_fn = orig_ifc.functions.get(prefix).unwrap_or_else(|| {
                    panic!("original function {prefix} must be found based on {fn_meta:?}")
                });
                let (params, result) = match fn_ext {
                    FunctionExtension::Submit => {
                        // -submit: func(join-set: borrow<join-set>, <params>) -> execution-id;
                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
                        let mut params = vec![Param {
                            name: generate_param_name("join-set", &original_fn.params),
                            ty: Type::Id(type_id_join_set_id_borrow_handle),
                            span: Span::default(),
                        }];
                        params.extend_from_slice(&original_fn.params);

                        (params, Some(Type::Id(type_id_execution_id)))
                    }
                    FunctionExtension::AwaitNext => {
                        // -await-next: func(join-set: borrow<join-set>) ->
                        //  result<return-type, await-next-extension-error>;
                        // The response id is no longer in the tuple; read it via
                        // `join-set.last-id`. Structurally identical to `-get`.
                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
                        let params = vec![Param {
                            name: "join-set".to_string(),
                            ty: Type::Id(type_id_join_set_id_borrow_handle),
                            span: Span::default(),
                        }];
                        let result = {
                            let type_id_result = resolve.types.alloc(TypeDef {
                                name: None,
                                kind: TypeDefKind::Result(wit_parser::Result_ {
                                    ok: original_fn.result,
                                    err: Some(Type::Id(type_id_await_next_err_part)),
                                }),
                                owner: TypeOwner::None,
                                docs: wit_parser::Docs::default(),
                                stability: wit_parser::Stability::default(),
                                span: Span::default(),
                            });
                            Some(Type::Id(type_id_result))
                        };
                        (params, result)
                    }
                    FunctionExtension::Schedule => {
                        // -schedule: func(schedule-at: schedule-at, <params>) -> execution-id;
                        assert_eq!(pkg_ext, PackageExtension::ObeliskSchedule);
                        let schedule_at_param_name =
                            generate_param_name("schedule-at", &original_fn.params);
                        let mut params = vec![Param {
                            name: schedule_at_param_name.clone(),
                            ty: Type::Id(type_id_schedule_at),
                            span: Span::default(),
                        }];
                        params.extend_from_slice(&original_fn.params);
                        let result = Some(Type::Id(type_id_execution_id));
                        (params, result)
                    }
                    FunctionExtension::Stub => {
                        // -stub: func(execution_id: execution-id, original retval) -> result<_, stub-error>;
                        assert_eq!(pkg_ext, PackageExtension::ObeliskStub);
                        let mut params = vec![Param {
                            name: "execution-id".to_string(),
                            ty: Type::Id(type_id_execution_id),
                            span: Span::default(),
                        }];
                        let Some(return_type) = &original_fn.result else {
                            unreachable!(
                                "return types of exported functions are validated in ExImLite"
                            )
                        };
                        let return_type_id = Type::Id(resolve.types.alloc(TypeDef {
                            name: None,
                            kind: TypeDefKind::Type(*return_type),
                            owner: TypeOwner::None,
                            docs: wit_parser::Docs::default(),
                            stability: wit_parser::Stability::default(),
                            span: Span::default(),
                        }));
                        params.push(Param {
                            name: "execution-result".to_string(),
                            ty: return_type_id,
                            span: Span::default(),
                        });

                        let result = {
                            let type_id_result = resolve.types.alloc(TypeDef {
                                name: None,
                                kind: TypeDefKind::Result(wit_parser::Result_ {
                                    ok: None,
                                    err: Some(Type::Id(type_id_stub_error)),
                                }),
                                owner: TypeOwner::None,
                                docs: wit_parser::Docs::default(),
                                stability: wit_parser::Stability::default(),
                                span: Span::default(),
                            });
                            Some(Type::Id(type_id_result))
                        };
                        (params, result)
                    }
                    FunctionExtension::Get => {
                        // -get(execution-id) -> result<originalreturn type, get-extension-error>
                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
                        let params = vec![Param {
                            name: "execution-id".to_string(),
                            ty: Type::Id(type_id_execution_id),
                            span: Span::default(),
                        }];
                        let result = Some(Type::Id(resolve.types.alloc(TypeDef {
                            name: None,
                            kind: TypeDefKind::Result(wit_parser::Result_ {
                                ok: original_fn.result,
                                err: Some(Type::Id(type_id_get_extension_error)),
                            }),
                            owner: TypeOwner::None,
                            docs: wit_parser::Docs::default(),
                            stability: wit_parser::Stability::default(),
                            span: Span::default(),
                        })));
                        (params, result)
                    }
                };
                let wit_fun = Function {
                    name: fn_name.to_string(),
                    kind: FunctionKind::Freestanding,
                    params,
                    result,
                    docs: wit_parser::Docs::default(),
                    stability: Stability::default(),
                    span: Span::default(),
                };
                ifc.functions.insert(fn_name.to_string(), wit_fun);
            }
            // Add Interface to `resolve`.
            let ifc_id = resolve.interfaces.alloc(ifc);
            resolve
                .packages
                .get_mut(pkg_id)
                .expect("found or inserted already")
                .interfaces
                .insert(ifc_fqn.ifc_name().to_string(), ifc_id);

            // Add the interface
            added_interfaces.push(ifc_id);
        }
    }
    Ok(added_interfaces)
}

pub(crate) fn packages_except_main(
    resolve: &Resolve,
    main_package: PackageId,
    sorted: bool,
) -> Vec<PackageId> {
    let mut packages = resolve
        .packages
        .iter()
        .map(|(id, _)| id)
        // The main package would show as a nested package as well
        .filter(|id| *id != main_package)
        .collect::<Vec<_>>();
    if sorted {
        packages.sort();
    }
    packages
}

// Replace obelisk:types from the actual WASM file because it may not contain all types we are going to need in exported functions.
fn replace_obelisk_types(wit: &str) -> String {
    // Replace last character of the first line from ; to {
    let types_nesting = {
        let mut types_nesting = WIT_OBELISK_TYPES_PACKAGE_CONTENT.replacen(';', "{", 1);
        types_nesting.push('}');
        types_nesting
    };
    const TYPES_NESTED_PACKAGE_FIRST_LINE: &str =
        formatcp!("package {OBELISK_TYPES_PACKAGE_NAME} {{");
    let wit = remove_nested_package(wit, TYPES_NESTED_PACKAGE_FIRST_LINE);
    let wit = format!("{wit}\n{types_nesting}");
    wit
}

fn generate_param_name(param_name: &str, params: &[Param]) -> String {
    let orig_param_names: hashbrown::HashSet<&str> = params
        .iter()
        .map(|Param { name, .. }| name.as_str())
        .collect();
    if orig_param_names.contains(param_name) {
        for my_char in 'a'..='z' {
            let name = format!("{param_name}-{my_char}");
            if !orig_param_names.contains(name.as_str()) {
                return name;
            }
        }
        warn!("Parameter name `{param_name}` collides with other params {orig_param_names:?}");
    }
    param_name.to_string()
}

fn copy_or_refer_original_types(
    orig_ifc_id: InterfaceId,
    orig_ifc: &Interface,
    resolve_types: &mut Arena<TypeDef>,
) -> IndexMap<String, id_arena::Id<TypeDef>> {
    let mut target_types = IndexMap::new();
    // Copy all imports from original to ext interface. Declared types like records will be referenced instead.
    for (name, orig_type_id) in &orig_ifc.types {
        let type_def = resolve_types
            .get(*orig_type_id)
            .expect("type def must be found in resolve");

        let allocated_type_id = match type_def.kind {
            TypeDefKind::Type(_) => resolve_types.alloc(type_def.clone()),
            _ => {
                // Create a reference to the type.
                resolve_types.alloc(TypeDef {
                    name: None,
                    kind: TypeDefKind::Type(Type::Id(*orig_type_id)),
                    owner: TypeOwner::Interface(orig_ifc_id),
                    docs: wit_parser::Docs::default(),
                    stability: wit_parser::Stability::default(),
                    span: Span::default(),
                })
            }
        };
        target_types.insert(name.clone(), allocated_type_id);
    }
    target_types
}

fn get_or_create_package(
    pkg_fqn: PkgFqn,
    resolve: &mut Resolve,
) -> Result<PackageId, semver::Error> {
    if let Some((pkg_id, _)) = resolve
        .packages
        .iter()
        .find(|(_, found_pkg)| from_wit_package_name_to_pkg_fqn(&found_pkg.name) == pkg_fqn)
    {
        Ok(pkg_id)
    } else {
        let pkg = wit_parser::Package {
            name: from_pkg_fqn_to_wit_package_name(pkg_fqn)?,
            docs: wit_parser::Docs::default(),
            interfaces: IndexMap::default(),
            worlds: IndexMap::default(),
        };
        let package_name = pkg.name.clone();
        let pkg_id = resolve.packages.alloc(pkg);
        resolve.package_names.insert(package_name, pkg_id);
        Ok(pkg_id)
    }
}

fn get_ext_pkg_to_ifc_to_details_map(
    exports: &[PackageIfcFns],
) -> IndexMap<PkgFqn, IndexMap<IfcFqnName, IndexMap<FnName, FunctionMetadata>>> {
    // Consistent iteration order so that the WIT output is deterministic.
    // Interfaces are sorted already.
    let mut pkg_to_ifc_to_details_map: IndexMap<PkgFqn, IndexMap<IfcFqnName, IndexMap<FnName, _>>> =
        IndexMap::new();
    for pkg_ifc_fns in exports {
        if pkg_ifc_fns.ifc_fqn.pkg_fqn_name().is_extension() {
            let inner_map = pkg_to_ifc_to_details_map
                .entry(pkg_ifc_fns.ifc_fqn.pkg_fqn_name())
                .or_default();
            inner_map.insert(pkg_ifc_fns.ifc_fqn.clone(), pkg_ifc_fns.fns.clone());
        }
    }
    pkg_to_ifc_to_details_map
}

fn find_interface<'a>(
    ifc_fqn: &IfcFqnName,
    resolve: &'_ Resolve,
    interfaces: &'a Arena<Interface>,
) -> Option<(InterfaceId, &'a Interface)> {
    let pkg_id = *resolve
        .package_names
        .get(&try_from_ifc_fqn_name(ifc_fqn).ok()?)?;
    let ifc_id = *resolve.packages[pkg_id]
        .interfaces
        .get(ifc_fqn.ifc_name())?;
    interfaces.get(ifc_id).map(|ifc| (ifc_id, ifc))
}

fn remove_nested_package(wit_string: &str, nested_package_to_remove: &str) -> String {
    // Find the start of the namespace
    let Some(nested_package_start) = wit_string.find(nested_package_to_remove) else {
        return wit_string.to_string();
    };

    // Find the opening brace after the namespace
    let Some(open_brace_index) = wit_string[nested_package_start..]
        .find('{')
        .map(|idx| nested_package_start + idx)
    else {
        panic!("nested namespace must contain '{{'");
    };

    // Track brace nesting to find the matching closing brace
    let mut brace_count = 1;
    let mut current_index = open_brace_index + 1;

    while current_index < wit_string.len() && brace_count > 0 {
        match wit_string.chars().nth(current_index) {
            Some('{') => brace_count += 1,
            Some('}') => brace_count -= 1,
            _ => {}
        }
        current_index += 1;
    }

    // If we didn't find the matching brace, return the original string
    if brace_count != 0 {
        warn!("Cannot remove the nested package {nested_package_to_remove}");
        return wit_string.to_string();
    }

    // Remove the package and its contents
    format!(
        "{}\n{}",
        wit_string[..nested_package_start].trim(),
        wit_string[current_index..].trim()
    )
}

fn try_from_ifc_fqn_name(ifc_fqn: &IfcFqnName) -> Result<PackageName, anyhow::Error> {
    Ok(PackageName {
        namespace: ifc_fqn.namespace().to_string(),
        name: ifc_fqn.package_name().to_string(),
        version: ifc_fqn
            .version()
            .map(semver::Version::parse)
            .transpose()
            .inspect_err(|err| {
                error!(
                    "cannot parse the version `{:?}` - {err:?}",
                    ifc_fqn.version()
                );
            })?,
    })
}

pub(crate) fn from_wit_package_name_to_pkg_fqn(package_name: &PackageName) -> PkgFqn {
    PkgFqn {
        namespace: package_name.namespace.clone(),
        package_name: package_name.name.clone(),
        version: package_name.version.as_ref().map(ToString::to_string),
    }
}

fn from_pkg_fqn_to_wit_package_name(pkg_fqn: PkgFqn) -> Result<PackageName, semver::Error> {
    Ok(PackageName {
        namespace: pkg_fqn.namespace,
        name: pkg_fqn.package_name,
        version: pkg_fqn
            .version
            .as_ref()
            .map(|v| v.parse())
            .transpose()
            .inspect_err(|err| error!("Cannot convert version {:?} - {err:?}", pkg_fqn.version))?,
    })
}

/// Build a WIT text map for all exported packages (primary + extension + `obelisk:types`
/// dependency) from the registry's already-merged export hierarchy.
pub fn build_wit_deps_map(
    all_exports: &[PackageIfcFns],
) -> Result<HashMap<PkgFqn, String /* WIT */>, anyhow::Error> {
    let mut result = HashMap::new();

    // Group non-extension entries by primary package.
    let mut primary_pkgs: IndexMap<PkgFqn, Vec<&PackageIfcFns>> = IndexMap::new();
    for pkg_ifc_fns in all_exports.iter().filter(|e| !e.extension) {
        primary_pkgs
            .entry(pkg_ifc_fns.ifc_fqn.pkg_fqn_name())
            .or_default()
            .push(pkg_ifc_fns);
    }

    let has_any_extension = all_exports.iter().any(|e| e.extension);

    for (pkg_fqn, ifc_fns_list) in &primary_pkgs {
        // Build primary Resolve and print it.
        let (primary_resolve, primary_pkg_id) = build_primary_resolve(pkg_fqn, ifc_fns_list, None)?;
        let primary_wit = {
            let mut printer = WitPrinter::new(OutputToString::default());
            printer.print(&primary_resolve, primary_pkg_id, &[])?;
            printer.output.to_string()
        };
        result.insert(pkg_fqn.clone(), primary_wit.clone());

        // Collect extension entries whose primary package is pkg_fqn.
        let ext_for_pkg: Vec<PackageIfcFns> = all_exports
            .iter()
            .filter(|e| {
                e.extension
                    && e.ifc_fqn
                        .pkg_fqn_name()
                        .split_ext()
                        .map(|(orig, _)| orig == *pkg_fqn)
                        .unwrap_or(false)
            })
            .cloned()
            .collect();

        if !ext_for_pkg.is_empty() {
            // Roundtrip through WIT text so that obelisk:types is available in the resolve
            // (required by add_extended_interfaces to build borrow<join-set> handles etc.).
            let wit_with_types = replace_obelisk_types(&primary_wit);
            let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit_with_types)?;
            let mut resolve_with_types = Resolve::new();
            resolve_with_types.push_group(group)?;

            add_extended_interfaces(&ext_for_pkg, &mut resolve_with_types)?;

            // Print each extension package that was just added.
            for (ext_pkg_id, ext_pkg) in &resolve_with_types.packages {
                let ext_pkg_fqn = from_wit_package_name_to_pkg_fqn(&ext_pkg.name);
                if ext_pkg_fqn.is_extension() {
                    let mut printer = WitPrinter::new(OutputToString::default());
                    printer.print(&resolve_with_types, ext_pkg_id, &[])?;
                    result.insert(ext_pkg_fqn, printer.output.to_string());
                }
            }
        }
    }

    // Always include obelisk:types when extension packages are present.
    if has_any_extension {
        let types_pkg_fqn = PkgFqn {
            namespace: "obelisk".to_string(),
            package_name: "types".to_string(),
            version: Some(OBELISK_TYPES_VERSION.to_string()),
        };
        result.insert(types_pkg_fqn, WIT_OBELISK_TYPES_PACKAGE_CONTENT.to_string());
    }

    Ok(result)
}

/// Build a [`Resolve`] that contains a single primary package with all its interfaces.
///
/// Each [`PackageIfcFns`] entry contributes one interface; functions are allocated using
/// [`crate::wit_builder::allocate_type`] so that named types (record, variant, enum, flags) get
/// proper `TypeDef` entries.
pub(crate) fn build_primary_resolve(
    pkg_fqn: &PkgFqn,
    ifc_fns_list: &[&PackageIfcFns],
    world_name: Option<&str>,
) -> Result<(Resolve, PackageId), anyhow::Error> {
    use crate::wit_builder::allocate_type;

    let mut resolve = Resolve::new();

    let pkg = wit_parser::Package {
        name: from_pkg_fqn_to_wit_package_name(pkg_fqn.clone())?,
        docs: wit_parser::Docs::default(),
        interfaces: IndexMap::default(),
        worlds: IndexMap::default(),
    };
    let pkg_id = resolve.packages.alloc(pkg);
    resolve
        .package_names
        .insert(resolve.packages[pkg_id].name.clone(), pkg_id);

    // Track already-created interfaces so that multiple PackageIfcFns entries
    // sharing the same interface name get their functions merged into one interface.
    let mut ifc_name_to_id: HashMap<String, InterfaceId> = HashMap::new();

    for &pkg_ifc_fns in ifc_fns_list {
        let ifc_name = pkg_ifc_fns.ifc_fqn.ifc_name().to_string();
        let ifc_id = if let Some(&existing_id) = ifc_name_to_id.get(&ifc_name) {
            existing_id
        } else {
            let iface = Interface {
                name: Some(ifc_name.clone()),
                types: IndexMap::new(),
                functions: IndexMap::new(),
                docs: wit_parser::Docs::default(),
                stability: Stability::default(),
                package: Some(pkg_id),
                span: Span::default(),
                clone_of: None,
            };
            let new_id = resolve.interfaces.alloc(iface);
            resolve
                .packages
                .get_mut(pkg_id)
                .unwrap()
                .interfaces
                .insert(ifc_name.clone(), new_id);
            ifc_name_to_id.insert(ifc_name, new_id);
            new_id
        };

        let mut dedup = HashMap::new();

        for (fn_name, fn_metadata) in &pkg_ifc_fns.fns {
            let wit_params: Vec<Param> = fn_metadata
                .parameter_types
                .iter()
                .map(|p| {
                    let ty = allocate_type(&mut resolve, ifc_id, &p.type_wrapper, &mut dedup);
                    Param {
                        name: p.name.as_ref().to_string(),
                        ty,
                        span: Span::default(),
                    }
                })
                .collect();

            let return_tw = fn_metadata.return_type.type_wrapper();
            let result_type = allocate_type(&mut resolve, ifc_id, &return_tw, &mut dedup);

            let wit_fn = Function {
                name: fn_name.to_string(),
                kind: FunctionKind::Freestanding,
                params: wit_params,
                result: Some(result_type),
                docs: wit_parser::Docs::default(),
                stability: Stability::default(),
                span: Span::default(),
            };
            resolve
                .interfaces
                .get_mut(ifc_id)
                .unwrap()
                .functions
                .insert(fn_name.to_string(), wit_fn);
        }

        // Collect named TypeDefs into the interface's types map.
        let mut types = IndexMap::new();
        for (type_id, type_def) in &resolve.types {
            if type_def.owner == TypeOwner::Interface(ifc_id)
                && let Some(name) = &type_def.name
            {
                types.insert(name.clone(), type_id);
            }
        }
        resolve.interfaces.get_mut(ifc_id).unwrap().types = types;
    }

    let main_pkg_id = if let Some(world_name) = world_name {
        let world_exports: IndexMap<WorldKey, WorldItem> = resolve.packages[pkg_id]
            .interfaces
            .values()
            .map(|&ifc_id| {
                (
                    WorldKey::Interface(ifc_id),
                    WorldItem::Interface {
                        id: ifc_id,
                        stability: Stability::Unknown,
                        span: Span::default(),
                    },
                )
            })
            .collect();
        // Put the synthesized world in a dedicated `root:component` package to avoid
        // a package-level dependency cycle.
        let root_pkg = wit_parser::Package {
            name: PackageName {
                namespace: "root".to_string(),
                name: "component".to_string(),
                version: None,
            },
            docs: wit_parser::Docs::default(),
            interfaces: IndexMap::default(),
            worlds: IndexMap::default(),
        };
        let root_pkg_id = resolve.packages.alloc(root_pkg);
        resolve
            .package_names
            .insert(resolve.packages[root_pkg_id].name.clone(), root_pkg_id);
        let world = World {
            name: world_name.to_string(),
            docs: wit_parser::Docs::default(),
            imports: IndexMap::default(),
            exports: world_exports,
            package: Some(root_pkg_id),
            span: Span::default(),
            includes: vec![],
            stability: Stability::Unknown,
        };
        let world_id = resolve.worlds.alloc(world);
        resolve
            .packages
            .get_mut(root_pkg_id)
            .unwrap()
            .worlds
            .insert(world_name.to_string(), world_id);
        root_pkg_id
    } else {
        pkg_id
    };

    Ok((resolve, main_pkg_id))
}

// TODO: Make the wit_component's OutputToString configurable for number of spaces
#[derive(Default)]
pub struct OutputToString {
    indent: usize,
    output: String,
    // set to true after newline, then to false after first item is indented.
    needs_indent: bool,
}

impl wit_component::Output for OutputToString {
    fn push_str(&mut self, src: &str) {
        self.output.push_str(src);
    }

    fn indent_if_needed(&mut self) -> bool {
        if self.needs_indent {
            for _ in 0..self.indent {
                // Indenting by 4 spaces.
                self.output.push_str("    ");
            }
            self.needs_indent = false;
            true
        } else {
            false
        }
    }

    fn indent_start(&mut self) {
        assert!(
            !self.needs_indent,
            "`indent_start` is never called after newline"
        );
        self.output.push_str(" {");
        self.indent += 1;
        self.newline();
    }

    fn indent_end(&mut self) {
        // Note that a `saturating_sub` is used here to prevent a panic
        // here in the case of invalid code being generated in debug
        // mode. It's typically easier to debug those issues through
        // looking at the source code rather than getting a panic.
        self.indent = self.indent.saturating_sub(1);
        self.indent_if_needed();
        self.output.push('}');
        self.newline();
    }

    fn newline(&mut self) {
        self.output.push('\n');
        self.needs_indent = true;
    }
}
impl Display for OutputToString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.output.fmt(f)
    }
}

#[cfg(test)]
mod tests {
    use crate::{wasm_tools::WasmComponent, wasm_tools::strip_wasm_hash, wit::OutputToString};
    use concepts::ComponentType;
    use rstest::rstest;
    use std::path::PathBuf;
    use wit_component::WitPrinter;
    use wit_parser::{Resolve, UnresolvedPackageGroup};

    #[rstest]
    #[case(
        test_programs_fibo_workflow_builder::TEST_PROGRAMS_FIBO_WORKFLOW,
        ComponentType::Workflow
    )]
    #[case(
        test_programs_fibo_activity_builder::TEST_PROGRAMS_FIBO_ACTIVITY,
        ComponentType::Activity
    )]
    #[case(
        test_programs_fibo_webhook_builder::TEST_PROGRAMS_FIBO_WEBHOOK,
        ComponentType::WebhookEndpoint
    )]
    #[case(
        test_programs_http_get_activity_builder::TEST_PROGRAMS_HTTP_GET_ACTIVITY,
        ComponentType::Activity
    )]
    #[case(
        test_programs_http_get_workflow_builder::TEST_PROGRAMS_HTTP_GET_WORKFLOW,
        ComponentType::Workflow
    )]
    #[case(
        test_programs_stub_activity_builder::TEST_PROGRAMS_STUB_ACTIVITY,
        ComponentType::ActivityStub
    )]
    #[case(
        test_programs_stub_workflow_builder::TEST_PROGRAMS_STUB_WORKFLOW,
        ComponentType::Workflow
    )]
    #[case(
        test_programs_sleep_activity_builder::TEST_PROGRAMS_SLEEP_ACTIVITY,
        ComponentType::Workflow
    )]
    #[case(
        test_programs_sleep_workflow_builder::TEST_PROGRAMS_SLEEP_WORKFLOW,
        ComponentType::Workflow
    )]
    fn wit_should_contain_extensions(
        #[case] wasm_path: &'static str,
        #[case] component_type: ComponentType,
    ) {
        test_utils::set_up();

        let component = WasmComponent::new(wasm_path, component_type).unwrap();
        let wasm_path = PathBuf::from(wasm_path);
        let wasm_file = strip_wasm_hash(&wasm_path.file_name().unwrap().to_string_lossy());
        let wit = component.wit();
        // Verify that the generated WIT parses.
        let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit).unwrap();
        let mut resolve = Resolve::new();
        let main_id = resolve.push_group(group).unwrap();
        let ids = resolve
            .packages
            .iter()
            .map(|(id, _)| id)
            // The main package would show as a nested package as well
            .filter(|id| *id != main_id)
            .collect::<Vec<_>>();
        let mut printer = WitPrinter::new(OutputToString::default());
        printer.print(&resolve, main_id, &ids).unwrap(); // verify it parses
        // store original WIT string in snapshots, because that is the `wit()` output.
        insta::with_settings!({sort_maps => true, snapshot_suffix => format!("{wasm_file}_wit")}, {insta::assert_snapshot!(wit)});
    }
}