alef 0.58.3

use crate::codegen::generators::trait_bridge::{TraitBridgeGenerator, TraitBridgeSpec, host_function_path};
use crate::core::ir::{MethodDef, TypeRef};
use std::collections::HashMap;

fn exported_pyfunction_symbol(fn_name: &str) -> String {
    fn_name.to_string()
}

/// PyO3-specific trait bridge generator.
/// Implements code generation for bridging Python objects to Rust traits.
pub struct Pyo3BridgeGenerator {
    /// Core crate import path (e.g., `"sample_core"`).
    pub core_import: String,
    /// Map of type name → fully-qualified Rust path for type references.
    pub type_paths: HashMap<String, String>,
    /// Error type name (e.g., `"SampleCrateError"`).
    pub error_type: String,
    /// Callback-param type names that get NATIVE-object marshalling — known serde structs per
    /// the shared [`crate::codegen::generators::trait_bridge::is_native_marshalled_struct`] rule.
    /// For such a param the bridge constructs the binding's native Python object (the `#[pyclass]`
    /// wrapper, via the same `From<core::T>` conversion used for return values) and hands THAT to
    /// the host method, instead of serializing the param to a JSON string. Enums, opaque/handle
    /// types, and excluded/unknown `Named` params are absent and keep their prior representation.
    pub struct_param_types: std::collections::HashSet<String>,
    /// Callback-RETURN type names that get NATIVE-object marshalling — known serde structs returned
    /// directly by a method (per the shared `native_marshalled_struct_returns` rule). For such a
    /// return the bridge first tries to extract the host's native Python object and convert it via
    /// `From<Binding>` for the core type, falling back to the JSON/mapping path otherwise.
    pub struct_return_types: std::collections::HashSet<String>,
    /// Rust-defaulted trait methods the bridge forwards to the host when the Python
    /// object defines them (per the shared `forwardable_defaulted_method_names` rule).
    /// Methods absent here keep the trait's Rust default unconditionally.
    pub forwardable_defaulted: std::collections::HashSet<String>,
    /// Struct types the package exports as options dataclasses. Callback params of
    /// these types are lifted from the native pyclass into the public dataclass via
    /// the generated `options._from_native_*` converter before invoking the host, so
    /// the type the host receives is the type the package exports under that name.
    pub options_dataclass_types: std::collections::HashSet<String>,
    /// Names of enums whose variants are ALL unit variants (no fields) — e.g.
    /// `ProcessingStage`, `OcrBackendType`. A callback returning one of these types is only
    /// ever choosing among a fixed set of bare names, so the bridge accepts the bare variant
    /// name (`"Early"`) in addition to the JSON-quoted form (`"\"Early\""`) required by the
    /// generic mapping path. Structs and data-carrying enums are NOT in this set and keep the
    /// strict JSON/mapping deserialization, since a bare string can't represent their fields.
    pub unit_enum_return_types: std::collections::HashSet<String>,
}

impl TraitBridgeGenerator for Pyo3BridgeGenerator {
    fn foreign_object_type(&self) -> &str {
        "Py<PyAny>"
    }

    fn gen_lifecycle_presence_check(&self, method: &MethodDef, _spec: &TraitBridgeSpec) -> Option<String> {
        Some(format!(
            "Python::attach(|py| self.inner.bind(py).hasattr(\"{}\").unwrap_or(false))",
            method.name
        ))
    }

    fn gen_method_presence_check(&self, method: &MethodDef, _spec: &TraitBridgeSpec) -> Option<String> {
        self.forwardable_defaulted.contains(&method.name).then(|| {
            format!(
                "Python::attach(|py| self.inner.bind(py).hasattr(\"{}\").unwrap_or(false))",
                method.name
            )
        })
    }

    fn bridge_imports(&self) -> Vec<String> {
        vec!["pyo3::prelude::*".to_string(), "std::sync::Arc".to_string()]
    }

    fn gen_sync_method_body(&self, method: &MethodDef, spec: &TraitBridgeSpec) -> String {
        let name = &method.name;
        let has_error = method.error_type.is_some();

        let py_args = self.sync_py_args(method);
        let run_args = if py_args.is_empty() {
            "bound_method,".to_string()
        } else {
            format!("bound_method, {py_args}")
        };
        let call = if py_args.is_empty() {
            format!("self.inner.bind(py).call_method0(\"{name}\")")
        } else {
            format!("self.inner.bind(py).call_method1(\"{name}\", ({py_args}))")
        };
        let error_expr = spec.make_error(&format!(
            "format!(\"Plugin '{{}}' method '{name}' failed: {{}}\", self.cached_name, e)"
        ));

        if matches!(method.return_type, TypeRef::Unit) {
            crate::backends::pyo3::template_env::render(
                "trait_bridge/sync_method_unit_return.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    has_error => has_error,
                    error_expr => error_expr,
                    wrapper => spec.wrapper_name(),
                },
            )
        } else {
            let ext = self.extract_ty(&method.return_type);
            let is_named = matches!(method.return_type, TypeRef::Named(_));
            let is_unit_enum = self.is_unit_enum_return(&method.return_type);
            let return_type_name = self.return_type_display_name(&method.return_type);
            let deserialize_error_expr = if is_unit_enum {
                format!(
                    "pyo3::exceptions::PyRuntimeError::new_err(format!(\"method '{name}' returned a value that does not match the expected return type `{return_type_name}`: {{}}. The returned value must be one of `{return_type_name}`'s variant names (e.g. \\\"Early\\\") or its JSON-quoted form.\", e))"
                )
            } else {
                format!(
                    "pyo3::exceptions::PyRuntimeError::new_err(format!(\"method '{name}' returned a value that does not match the expected return type `{return_type_name}`: {{}}. The returned value must be a mapping matching the fields of `{return_type_name}`.\", e))"
                )
            };
            crate::backends::pyo3::template_env::render(
                "trait_bridge/sync_method_non_unit_return.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    is_named => is_named,
                    is_unit_enum => is_unit_enum,
                    extract_ty => ext,
                    native_return_binding => self.native_struct_return(&method.return_type),
                    has_error => has_error,
                    error_expr => error_expr,
                    deserialize_error_expr => deserialize_error_expr,
                    wrapper => spec.wrapper_name(),
                },
            )
        }
    }

    fn gen_async_method_body(&self, method: &MethodDef, spec: &TraitBridgeSpec) -> String {
        let name = &method.name;

        if let Some((mut_param_name, mut_native_ty, mut_core_ty)) = self.mut_writeback_param(method) {
            let py_args = self.async_py_args(method);
            let run_args = if py_args.is_empty() {
                "bound_method,".to_string()
            } else {
                format!("bound_method, {py_args}")
            };
            let call = if py_args.is_empty() {
                format!("obj.call_method0(\"{name}\")")
            } else {
                format!("obj.call_method1(\"{name}\", ({py_args}))")
            };
            let params: Vec<minijinja::Value> = method
                .params
                .iter()
                .map(|p| {
                    minijinja::context! {
                        name => &p.name,
                        ty => match &p.ty {
                            TypeRef::Bytes => "Bytes",
                            TypeRef::Path => "Path",
                            TypeRef::Named(n) => n.as_str(),
                            _ => "",
                        }.to_string(),
                        ty_is_named => matches!(&p.ty, TypeRef::Named(_)),
                        is_native_struct => matches!(&p.ty, TypeRef::Named(n) if self.is_native_struct_param(n)),
                        is_ref => p.is_ref,
                    }
                })
                .collect();
            let param_cloning = crate::backends::pyo3::template_env::render(
                "trait_bridge/async_param_cloning.jinja",
                minijinja::context! { params => params },
            );
            let error_expr = spec.make_error(&format!(
                "format!(\"Plugin '{{}}' method '{name}' failed: {{}}\", cached_name, e)"
            ));
            let json_error_expr =
                spec.make_error("format!(\"Plugin '{}': JSON serialization failed: {}\", cached_name, e)");
            let deserialize_error_expr = spec.make_error(&format!(
                "format!(\"Plugin '{{}}' method '{name}' returned a value that does not match the expected type `{mut_native_ty}`: {{}}. The returned value must be the (optionally modified) `{mut_native_ty}`, or None to leave it unchanged.\", cached_name, e)"
            ));
            let spawn_error_expr = spec.make_error("format!(\"spawn_blocking failed: {}\", e)");
            return crate::backends::pyo3::template_env::render(
                "trait_bridge/async_method_mut_writeback.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    param_cloning => param_cloning,
                    mut_param_name => mut_param_name,
                    mut_native_ty => mut_native_ty,
                    mut_core_ty => mut_core_ty,
                    error_expr => error_expr,
                    json_error_expr => json_error_expr,
                    deserialize_error_expr => deserialize_error_expr,
                    spawn_error_expr => spawn_error_expr,
                },
            );
        }

        let params: Vec<minijinja::Value> = method
            .params
            .iter()
            .map(|p| {
                minijinja::context! {
                    name => &p.name,
                    ty => match &p.ty {
                        TypeRef::Bytes => "Bytes",
                        TypeRef::Path => "Path",
                        TypeRef::Named(_) => {

                            match &p.ty {
                                TypeRef::Named(n) => n.as_str(),
                                _ => "",
                            }
                        },
                        _ => "",
                    }.to_string(),
                    ty_is_named => matches!(&p.ty, TypeRef::Named(_)),
                    is_native_struct => matches!(&p.ty, TypeRef::Named(n) if self.is_native_struct_param(n)),
                    is_ref => p.is_ref,
                }
            })
            .collect();

        let param_cloning = crate::backends::pyo3::template_env::render(
            "trait_bridge/async_param_cloning.jinja",
            minijinja::context! {
                params => params,
            },
        );

        let py_args = self.async_py_args(method);
        let run_args = if py_args.is_empty() {
            "bound_method,".to_string()
        } else {
            format!("bound_method, {py_args}")
        };
        let call = if py_args.is_empty() {
            format!("obj.call_method0(\"{name}\")")
        } else {
            format!("obj.call_method1(\"{name}\", ({py_args}))")
        };
        let error_expr = spec.make_error(&format!(
            "format!(\"Plugin '{{}}' method '{name}' failed: {{}}\", cached_name, e)"
        ));
        let json_error_expr =
            spec.make_error("format!(\"Plugin '{}': JSON serialization failed: {}\", cached_name, e)");
        let return_type_name = self.return_type_display_name(&method.return_type);
        let deserialize_error_expr = spec.make_error(&format!(
            "format!(\"Plugin '{{}}' method '{name}' returned a value that does not match the expected return type `{return_type_name}`: {{}}. The returned value must be a mapping matching the fields of `{return_type_name}`.\", cached_name, e)"
        ));
        let spawn_error_expr = spec.make_error("format!(\"spawn_blocking failed: {}\", e)");

        if self.is_named(&method.return_type) {
            let return_type =
                crate::codegen::generators::trait_bridge::format_type_ref(&method.return_type, &spec.type_paths);
            crate::backends::pyo3::template_env::render(
                "trait_bridge/async_method_named_return.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    param_cloning => param_cloning,
                    return_type => return_type,
                    native_return_binding => self.native_struct_return(&method.return_type),
                    error_expr => error_expr,
                    json_error_expr => json_error_expr,
                    deserialize_error_expr => deserialize_error_expr,
                    spawn_error_expr => spawn_error_expr,
                },
            )
        } else if matches!(method.return_type, TypeRef::Unit) {
            crate::backends::pyo3::template_env::render(
                "trait_bridge/async_method_unit_return.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    param_cloning => param_cloning,
                    error_expr => error_expr,
                    spawn_error_expr => spawn_error_expr,
                },
            )
        } else {
            let ext = self.extract_ty(&method.return_type);
            crate::backends::pyo3::template_env::render(
                "trait_bridge/async_method_non_unit_return.jinja",
                minijinja::context! {
                    method_name => name,
                    call => call,
                    run_args => run_args,
                    extract_ty => ext,
                    param_cloning => param_cloning,
                    error_expr => error_expr,
                    spawn_error_expr => spawn_error_expr,
                },
            )
        }
    }

    fn gen_constructor(&self, spec: &TraitBridgeSpec) -> String {
        let wrapper = spec.wrapper_name();
        let required_methods = spec.required_methods();
        crate::backends::pyo3::template_env::render(
            "trait_bridge/constructor.jinja",
            minijinja::context! {
                wrapper => wrapper,
                required_methods => required_methods,
            },
        )
    }

    fn gen_unregistration_fn(&self, spec: &TraitBridgeSpec) -> String {
        let Some(unregister_fn) = spec.bridge_config.unregister_fn.as_deref() else {
            return String::new();
        };
        let host_path = host_function_path(spec, unregister_fn);
        let host_symbol = exported_pyfunction_symbol(unregister_fn);
        crate::backends::pyo3::template_env::render(
            "trait_bridge/unregistration_fn.jinja",
            minijinja::context! {
                unregister_fn => unregister_fn,
                host_symbol => host_symbol,
                host_path => host_path,
            },
        )
    }

    fn gen_clear_fn(&self, spec: &TraitBridgeSpec) -> String {
        let Some(clear_fn) = spec.bridge_config.clear_fn.as_deref() else {
            return String::new();
        };
        let host_path = host_function_path(spec, clear_fn);
        let host_symbol = exported_pyfunction_symbol(clear_fn);
        crate::backends::pyo3::template_env::render(
            "trait_bridge/clear_fn.jinja",
            minijinja::context! {
                clear_fn => clear_fn,
                host_symbol => host_symbol,
                host_path => host_path,
            },
        )
    }

    fn gen_registration_fn(&self, spec: &TraitBridgeSpec) -> String {
        let Some(register_fn) = spec.bridge_config.register_fn.as_deref() else {
            return String::new();
        };
        let Some(registry_getter) = spec.bridge_config.registry_getter.as_deref() else {
            return String::new();
        };
        let wrapper = spec.wrapper_name();
        let trait_path = spec.trait_path();

        let req_methods: Vec<&MethodDef> = spec.required_methods();
        let required_methods_str = req_methods
            .iter()
            .map(|m| format!("\"{}\"", m.name))
            .collect::<Vec<_>>()
            .join(", ");

        let register_extra_args = spec
            .bridge_config
            .register_extra_args
            .as_deref()
            .map(|a| format!(", {a}"))
            .unwrap_or_default();

        crate::backends::pyo3::template_env::render(
            "trait_bridge/registration_fn.jinja",
            minijinja::context! {
                register_fn => register_fn,
                wrapper => wrapper,
                trait_path => trait_path,
                registry_getter => registry_getter,
                register_extra_args => register_extra_args,
                has_required_methods => !req_methods.is_empty(),
                required_methods_str => required_methods_str,
            },
        )
    }
}

impl Pyo3BridgeGenerator {
    /// Human-facing name of a return type for callback deserialization error messages.
    /// For a `Named` type this is the bare type name (e.g. `Doc`), not the fully-qualified
    /// Rust path, so the message reads the way a host implementer thinks about their return
    /// value. Other shapes fall back to their Rust rendering.
    fn return_type_display_name(&self, ty: &TypeRef) -> String {
        match ty {
            TypeRef::Named(name) => name.clone(),
            other => self.extract_ty(other),
        }
    }

    /// Extract the Python type that corresponds to a Rust TypeRef.
    fn extract_ty(&self, ty: &TypeRef) -> String {
        match ty {
            TypeRef::Primitive(p) => self.prim(p).to_string(),
            TypeRef::String | TypeRef::Path | TypeRef::Char => "String".into(),
            TypeRef::Bytes => "Vec<u8>".into(),
            TypeRef::Vec(inner) => format!("Vec<{}>", self.extract_ty(inner)),
            TypeRef::Optional(inner) => format!("Option<{}>", self.extract_ty(inner)),
            TypeRef::Named(name) => self
                .type_paths
                .get(name.as_str())
                .map(|p| p.replace('-', "_"))
                .unwrap_or_else(|| format!("{}::{}", self.core_import, name)),
            TypeRef::Unit => "()".into(),
            TypeRef::Map(k, v) => format!(
                "std::collections::HashMap<{}, {}>",
                self.extract_ty(k),
                self.extract_ty(v)
            ),
            TypeRef::Json => "String".into(),
            TypeRef::Duration => "u64".into(),
        }
    }

    /// Get the Rust string representation of a primitive type.
    fn prim(&self, p: &crate::core::ir::PrimitiveType) -> &'static str {
        use crate::core::ir::PrimitiveType::*;
        match p {
            Bool => "bool",
            U8 => "u8",
            U16 => "u16",
            U32 => "u32",
            U64 => "u64",
            I8 => "i8",
            I16 => "i16",
            I32 => "i32",
            I64 => "i64",
            F32 => "f32",
            F64 => "f64",
            Usize => "usize",
            Isize => "isize",
        }
    }

    /// True when `ty` is a `Named` type registered in `unit_enum_return_types` — an enum whose
    /// variants are all unit (fieldless), like `ProcessingStage` or `OcrBackendType`.
    fn is_unit_enum_return(&self, ty: &TypeRef) -> bool {
        matches!(ty, TypeRef::Named(name) if self.unit_enum_return_types.contains(name))
    }

    /// Detects the "in-place mutation" callback pattern: a `Unit`-returning method with a
    /// `&mut Named` parameter whose type is native-marshalled (e.g.
    /// `PostProcessor::process(&self, result: &mut ExtractedDocument, ..)`). Python cannot
    /// mutate a frozen `#[pyclass]` in place, so the bridge instead treats the callback's
    /// *return value* as the (optionally) updated value and writes it back into `*param` after
    /// the call, rather than silently discarding it as the naive `.map(|_| ())` bridge did.
    ///
    /// Returns `(param_name, native_binding_type_name, fully_qualified_core_type_path)` when the
    /// pattern applies. Returns `None` for any other shape (no mut param, non-Unit return, or a
    /// mut param whose type isn't a known native-marshalled struct) so those keep the existing
    /// bridge behavior unchanged.
    fn mut_writeback_param(&self, method: &MethodDef) -> Option<(String, String, String)> {
        if !matches!(method.return_type, TypeRef::Unit) {
            return None;
        }
        let param = method.params.iter().find(|p| p.is_mut)?;
        let TypeRef::Named(name) = &param.ty else {
            return None;
        };
        if !self.struct_param_types.contains(name) {
            return None;
        }
        Some((param.name.clone(), name.clone(), self.extract_ty(&param.ty)))
    }

    /// True when a `Named(name)` param should be handed to the host as the binding's native
    /// Python object rather than a JSON string — i.e. it is a known serde struct per the shared
    /// allowlist. The native object is the `#[pyclass]` wrapper, constructed from the core value
    /// via the same `From<core::T>` conversion the binding uses for function return values.
    fn is_native_struct_param(&self, name: &str) -> bool {
        self.struct_param_types.contains(name)
    }

    /// Build Python call argument expressions for a sync method.
    fn sync_py_args(&self, method: &MethodDef) -> String {
        let args: Vec<String> = method
            .params
            .iter()
            .map(|p| match (&p.ty, p.is_ref) {
                (TypeRef::Bytes, true) => format!("pyo3::types::PyBytes::new(py, {})", p.name),
                (TypeRef::Path, true) => format!("{}.to_str().unwrap_or_default()", p.name),
                // (`{Binding}::from(core_value)`). PyO3 auto-converts the `#[pyclass]` to a Python
                (TypeRef::Named(n), true) if self.is_native_struct_param(n) => {
                    self.options_lift_expr(n, format!("{}::from((*{}).clone())", n, p.name))
                }
                (TypeRef::Named(n), false) if self.is_native_struct_param(n) => {
                    self.options_lift_expr(n, format!("{}::from({}.clone())", n, p.name))
                }
                (TypeRef::Named(_), true) => {
                    format!("serde_json::to_string({}).unwrap_or_default()", p.name)
                }
                _ => p.name.clone(),
            })
            .collect();
        if args.len() == 1 {
            format!("{},", args[0])
        } else {
            args.join(", ")
        }
    }

    /// Build Python call argument expressions for an async method.
    fn async_py_args(&self, method: &MethodDef) -> String {
        let args: Vec<String> = method
            .params
            .iter()
            .map(|p| match (&p.ty, p.is_ref) {
                (TypeRef::Bytes, true) => format!("pyo3::types::PyBytes::new(py, &{})", p.name),
                (TypeRef::Path, true) => format!("{}_str.as_str()", p.name),
                (TypeRef::Named(n), _) if self.is_native_struct_param(n) => {
                    self.options_lift_expr(n, format!("{}::from({}_owned)", n, p.name))
                }
                (TypeRef::Named(n), false) if self.is_native_struct_param(n) => {
                    self.options_lift_expr(n, format!("{}::from({}.clone())", n, p.name))
                }
                (TypeRef::Named(_), true) => format!("{}_json.as_str()", p.name),
                _ => p.name.clone(),
            })
            .collect();
        if args.len() == 1 {
            format!("{},", args[0])
        } else {
            args.join(", ")
        }
    }

    /// Wrap a native-pyclass expression in the options-dataclass lift when the type
    /// is publicly exported as an options dataclass; otherwise pass the native object.
    fn options_lift_expr(&self, type_name: &str, native_expr: String) -> String {
        use heck::ToSnakeCase;
        if self.options_dataclass_types.contains(type_name) {
            format!(
                "__alef_options_from_native(py, \"_from_native_{}\", {})",
                type_name.to_snake_case(),
                native_expr
            )
        } else {
            native_expr
        }
    }

    /// Check if a TypeRef is a Named type.
    fn is_named(&self, ty: &TypeRef) -> bool {
        matches!(ty, TypeRef::Named(_))
    }

    /// Binding pyclass type name to extract for a native-object return, when the return is a bare
    /// `Named` struct on the native-marshalled return allowlist. The bridge tries
    /// `py_result.extract::<Binding>()` and converts via `From<Binding>` for the core type, falling
    /// back to the JSON/mapping path. `None` keeps the mapping path unchanged.
    fn native_struct_return<'a>(&self, ty: &'a TypeRef) -> Option<&'a str> {
        match ty {
            TypeRef::Named(n) if self.struct_return_types.contains(n) => Some(n.as_str()),
            _ => None,
        }
    }
}