miden-base-macros 0.12.0

use std::{
    collections::{BTreeSet, HashSet},
    fs,
    io::ErrorKind,
};

use proc_macro::Span;
use semver::Version;
use syn::spanned::Spanned;

use crate::{
    component_macro::{CORE_TYPES_PACKAGE, ComponentMethod, MethodReturn, to_kebab_case},
    types::{ExportedTypeDef, ExportedTypeKind, ensure_custom_type_defined},
    util::generated_wit_folder,
    wit_builder::WitBuilder,
};

/// Writes the generated component WIT to the crate's `wit` directory so that dependent targets can
/// reference it via manifest metadata.
pub fn write_component_wit_file(
    call_site_span: Span,
    wit_source: &str,
    package_name: &str,
) -> Result<(), syn::Error> {
    let sanitized_package_name = sanitize_package_name(package_name);
    let autogenerated_wit_folder = generated_wit_folder()?;
    let wit_path = autogenerated_wit_folder.join(format!("{sanitized_package_name}.wit"));

    let needs_write = match fs::read_to_string(&wit_path) {
        Ok(existing) => existing != wit_source,
        Err(err) if err.kind() == ErrorKind::NotFound => true,
        Err(err) => {
            return Err(syn::Error::new(
                call_site_span.into(),
                format!("failed to read existing WIT file '{}': {err}", wit_path.display()),
            ));
        }
    };

    if needs_write {
        fs::write(&wit_path, wit_source).map_err(|err| {
            syn::Error::new(
                call_site_span.into(),
                format!("failed to write WIT file '{}': {err}", wit_path.display()),
            )
        })?;
    }

    Ok(())
}

/// Renders the inline WIT source describing the component interface exported by the `impl` block.
pub fn build_component_wit(
    component_package: &str,
    component_version: &Version,
    interface_name: &str,
    world_name: &str,
    type_imports: &BTreeSet<String>,
    methods: &[ComponentMethod],
    exported_types: &[ExportedTypeDef],
) -> Result<String, syn::Error> {
    let exported_type_names: HashSet<String> =
        exported_types.iter().map(|def| def.wit_name.clone()).collect();

    let mut combined_core_imports = type_imports.clone();
    for exported in exported_types {
        match &exported.kind {
            ExportedTypeKind::Record { fields } => {
                for field in fields {
                    ensure_custom_type_defined(
                        &field.ty,
                        &exported_type_names,
                        Span::call_site().into(),
                    )?;
                    if !field.ty.is_custom {
                        combined_core_imports.insert(field.ty.wit_name.clone());
                    }
                }
            }
            ExportedTypeKind::Variant { variants } => {
                for variant in variants {
                    if let Some(payload) = &variant.payload {
                        ensure_custom_type_defined(
                            payload,
                            &exported_type_names,
                            Span::call_site().into(),
                        )?;
                        if !payload.is_custom {
                            combined_core_imports.insert(payload.wit_name.clone());
                        }
                    }
                }
            }
        }
    }

    let mut wit = WitBuilder::new("#[component]", component_package, component_version);
    wit.use_path(CORE_TYPES_PACKAGE);
    wit.blank_line();
    wit.interface(interface_name, |interface| {
        if !combined_core_imports.is_empty() {
            let imports = combined_core_imports.iter().cloned().collect::<Vec<_>>().join(", ");
            interface.line(&format!("use core-types.{{{imports}}};"));
            interface.blank_line();
        }

        for (index, exported) in exported_types.iter().enumerate() {
            if index > 0 {
                interface.blank_line();
            }

            match &exported.kind {
                ExportedTypeKind::Record { fields } => {
                    interface.block(&format!("record {} {{", exported.wit_name), |record| {
                        for field in fields {
                            let field_name = to_kebab_case(&field.name);
                            record.line(&format!("{field_name}: {},", field.ty.wit_name));
                        }
                    });
                }
                ExportedTypeKind::Variant { variants } => {
                    interface.block(
                        &format!("variant {} {{", exported.wit_name),
                        |variant_block| {
                            for variant in variants {
                                if let Some(payload) = &variant.payload {
                                    variant_block.line(&format!(
                                        "{}({}),",
                                        variant.wit_name, payload.wit_name
                                    ));
                                } else {
                                    variant_block.line(&format!("{},", variant.wit_name));
                                }
                            }
                        },
                    );
                }
            }
        }

        if !exported_types.is_empty() && !methods.is_empty() {
            interface.blank_line();
        }

        for method in methods {
            let signature = component_method_signature(method, &exported_type_names)?;
            interface.line(&signature);
        }

        Ok::<(), syn::Error>(())
    })?;
    wit.blank_line();
    wit.world(world_name, |world| {
        world.line(&format!("export {interface_name};"));
    });

    Ok(wit.finish())
}

/// Renders the WIT function signature for a component method.
fn component_method_signature(
    method: &ComponentMethod,
    exported_type_names: &HashSet<String>,
) -> Result<String, syn::Error> {
    for param in &method.params {
        ensure_custom_type_defined(&param.type_ref, exported_type_names, param.user_ty.span())?;
    }
    if let MethodReturn::Type { type_ref, user_ty } = &method.return_info {
        ensure_custom_type_defined(type_ref, exported_type_names, user_ty.span())?;
    }

    let signature = if method.params.is_empty() {
        match &method.return_info {
            MethodReturn::Unit => format!("{}: func();", method.wit_name),
            MethodReturn::Type { type_ref, .. } => {
                format!("{}: func() -> {};", method.wit_name, type_ref.wit_name)
            }
        }
    } else {
        let params = method
            .params
            .iter()
            .map(|param| format!("{}: {}", param.wit_param_name, param.type_ref.wit_name))
            .collect::<Vec<_>>()
            .join(", ");
        match &method.return_info {
            MethodReturn::Unit => format!("{}: func({params});", method.wit_name),
            MethodReturn::Type { type_ref, .. } => {
                format!("{}: func({params}) -> {};", method.wit_name, type_ref.wit_name)
            }
        }
    };

    Ok(signature)
}

fn sanitize_package_name(package_name: &str) -> String {
    let mut sanitized = package_name
        .chars()
        .map(|ch| match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' => ch,
            _ => '-',
        })
        .collect::<String>();

    if sanitized.is_empty() {
        sanitized.push_str("component");
    }

    sanitized
}