olai-codegen 0.0.4

//! PyO3 wrapper `#[pyclass]` generation for the protobuf-defined model types.
//!
//! ## Why wrappers
//!
//! Models are generated by `buf` (the buffa or prost plugin), not by trestle, so
//! we can neither add `#[pyclass]` to the struct definition nor a `#[derive]`.
//! An earlier approach implemented PyO3's [`FromPyObject`]/[`IntoPyObject`]
//! directly on the model types via [`pythonize`], so a model crossed the FFI
//! boundary as a plain Python **dict**. That compiled, but the Python side never
//! got *real objects*: no attribute access (`catalog.name`), no `isinstance`, no
//! `__repr__`, and enums arrived as bare proto-name strings.
//!
//! Instead we emit, for every project message, a newtype `#[pyclass]` wrapper
//! (`PyCatalog(Catalog)`) with native, typed getters/setters, a keyword
//! constructor, and `__repr__`/`__eq__`. The wrapper is *local* to the model
//! crate (orphan rule satisfied), and `From`/`Into` conversions bridge it to the
//! bare model type so the generated client bindings can convert at the boundary.
//! Enums get their own `#[pyclass(eq, eq_int)]` wrapper (see [`enum_wrappers`]).
//!
//! Everything is emitted into a sibling `pyo3_impls.rs` of the generated models
//! module and gated behind `#[cfg(feature = "python")]`, so pure-Rust /
//! server-only / WASM consumers never pull `pyo3` into their dependency tree.
//!
//! ## Runtime differences
//!
//! Unlike the old pythonize impls, the wrapper *bodies* are runtime-aware: buffa
//! stores singular messages as `MessageField<T>` and enums as `EnumValue<E>`,
//! while prost uses `Option<T>` (or `Option<Box<T>>` for fields prost boxes to
//! break a recursion cycle) and bare `i32`. The accessor expressions are selected
//! by [`Runtime`] and, for messages, by per-field boxing (see [`boxed_fields`]);
//! the *public* wrapper surface (class + method names + Python-visible types) is
//! identical across runtimes.
//!
//! [`FromPyObject`]: pyo3::FromPyObject
//! [`IntoPyObject`]: pyo3::IntoPyObject
//! [`pythonize`]: https://docs.rs/pythonize

use std::collections::BTreeSet;

use convert_case::{Case, Casing};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

use super::super::format_tokens;
use crate::analysis::GenerationPlan;
use crate::codegen::Runtime;
use crate::parsing::types::{BaseType, UnifiedType};
use crate::parsing::{CodeGenMetadata, EnumInfo, MessageField, MessageInfo};
use crate::utils::{extract_qualified_type_name, extract_simple_type_name};

/// Generate the `pyo3_impls.rs` content: a `#[pyclass]` wrapper for every message
/// and enum defined in the project's own packages, plus their `From`/`Into`
/// bridges to the bare model types.
///
/// Returns an empty string when there are no project messages/enums to emit (e.g.
/// a service-less descriptor), so the caller can skip writing the file.
pub(crate) fn generate_pyo3_impls(
    plan: &GenerationPlan,
    metadata: &CodeGenMetadata,
    runtime: Runtime,
) -> crate::error::Result<String> {
    // The project's own packages, taken straight from the plan's services — the
    // same set whose model modules `generate_models_mod` emits. Types in any other
    // package (google/gnostic/etc.) are not in our crate and must be skipped
    // (orphan rule + they have no in-crate path).
    let project_packages: BTreeSet<&str> =
        plan.services.iter().map(|s| s.package.as_str()).collect();

    let mut items: Vec<TokenStream> = Vec::new();

    // Enums first so message getters can reference the enum wrappers.
    let mut enum_names: Vec<&String> = metadata.enums.keys().collect();
    enum_names.sort();
    for name in enum_names {
        let Some(rust_path) = type_rust_path(name, &project_packages) else {
            continue;
        };
        let info = &metadata.enums[name];
        // An enum with no values has no representable wrapper; skip it.
        if info.values.is_empty() {
            continue;
        }
        items.push(enum_wrapper(name, info, &rust_path, runtime)?);
    }

    // Emit a wrapper for *every* message in the project's own packages, not just
    // resource-annotated ones: request/response and other non-resource messages
    // also cross the FFI boundary. Sort by name for deterministic output.
    let mut msg_names: Vec<&String> = metadata.messages.keys().collect();
    msg_names.sort();
    for name in msg_names {
        let Some(rust_path) = type_rust_path(name, &project_packages) else {
            continue;
        };
        let info = &metadata.messages[name];
        items.push(message_wrapper(name, info, &rust_path, runtime, metadata)?);
    }

    if items.is_empty() {
        return Ok(String::new());
    }

    // No file-level inner attributes (`#![..]`) or inner doc comments (`//!`):
    // this file is mounted via `#[cfg(feature = "python")] mod pyo3_impls;` in the
    // generated `mod.rs` *and* `include!`d directly in some setups; inner
    // attributes are illegal inside an `include!`d module body, so every allow is
    // applied as an outer attribute on the individual emitted items instead.
    let tokens = quote! {
        // PyO3 wrapper `#[pyclass]` types for the generated models.
        //
        // Each model message/enum gets a newtype wrapper exposing native, typed
        // getters/setters and a keyword constructor, plus `From`/`Into` bridges to
        // the bare model type. Mounted under the `python` feature; see the model
        // module's `mod.rs`.

        #(#items)*
    };

    format_tokens(tokens)
}

/// The `Py`-prefixed wrapper ident for a fully-qualified proto type name, using
/// the parent-qualified emitted name (collision-free for nested types).
fn wrapper_ident(fq_name: &str) -> proc_macro2::Ident {
    format_ident!("Py{}", extract_qualified_type_name(fq_name))
}

/// The Python-visible class name (no `Py` prefix) for a fully-qualified proto name.
fn python_class_name(fq_name: &str) -> String {
    extract_qualified_type_name(fq_name)
}

// ---------------------------------------------------------------------------
// Enum wrappers
// ---------------------------------------------------------------------------

/// Emit a `#[pyclass(eq, eq_int)]` enum mirroring a proto enum, plus `From`
/// conversions to/from the bare Rust enum.
///
/// The wrapper is an integer-valued Python enum (so `==` and `int()` work). The
/// boundary to the buffa/prost runtime is by proto value: each variant carries
/// its proto number, and `From`/`Into` map through the runtime enum's `i32` repr.
fn enum_wrapper(
    fq_name: &str,
    info: &EnumInfo,
    rust_path: &str,
    runtime: Runtime,
) -> crate::error::Result<TokenStream> {
    let wrapper = wrapper_ident(fq_name);
    let class_name = python_class_name(fq_name);
    let inner: syn::Path = parse_path(rust_path, fq_name)?;

    // Variant idents mirror the proto value names (UPPER_SNAKE), which are valid
    // Python identifiers and stay stable regardless of the runtime's own variant
    // renaming (prost PascalCases + strips the type prefix; buffa keeps the proto
    // name). We therefore never reference the *runtime* enum's variant idents —
    // conversions go through the proto `i32` value, which both runtimes expose.
    let variants: Vec<proc_macro2::Ident> = info
        .values
        .iter()
        .map(|v| format_ident!("{}", v.name))
        .collect();
    // pyo3's `#[pyclass(eq_int)]` enums require `isize` discriminants.
    let numbers: Vec<proc_macro2::Literal> = info
        .values
        .iter()
        .map(|v| proc_macro2::Literal::isize_suffixed(v.number as isize))
        .collect();

    // The wrapper -> i32 arms and i32 -> wrapper arms (shared by both runtimes).
    let to_i32_arms = info.values.iter().map(|v| {
        let id = format_ident!("{}", v.name);
        let n = proc_macro2::Literal::i32_suffixed(v.number);
        quote! { #wrapper::#id => #n }
    });
    let from_i32_arms = info.values.iter().map(|v| {
        let id = format_ident!("{}", v.name);
        let n = proc_macro2::Literal::i32_suffixed(v.number);
        quote! { #n => #wrapper::#id }
    });
    let fallback = &variants[0];

    // Runtime-specific bridges between the wrapper's `i32` and the bare model enum.
    // buffa exposes `Enumeration::{to_i32, from_i32}`; prost enums are
    // `TryFrom<i32>` + `as i32`.
    let (inner_to_i32, i32_to_inner): (TokenStream, TokenStream) = match runtime {
        Runtime::Buffa => (
            quote! { <#inner as ::buffa::Enumeration>::to_i32(&value) },
            quote! { <#inner as ::buffa::Enumeration>::from_i32(n).unwrap_or_default() },
        ),
        Runtime::Prost => (
            quote! { value as i32 },
            quote! { <#inner as ::core::convert::TryFrom<i32>>::try_from(n).unwrap_or_default() },
        ),
    };

    Ok(quote! {
        // Variant idents mirror the proto value names (UPPER_SNAKE), which trip the
        // camel-case lint; the conversions are intentionally exhaustive.
        #[allow(non_camel_case_types)]
        #[::pyo3::pyclass(eq, eq_int, name = #class_name)]
        #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
        pub enum #wrapper {
            #(#variants = #numbers,)*
        }

        impl #wrapper {
            fn __to_proto_i32(self) -> i32 {
                match self {
                    #(#to_i32_arms,)*
                }
            }
            fn __from_proto_i32(value: i32) -> Self {
                match value {
                    #(#from_i32_arms,)*
                    // Unknown values fall back to the first declared variant.
                    _ => #wrapper::#fallback,
                }
            }
        }

        impl ::core::convert::From<#inner> for #wrapper {
            fn from(value: #inner) -> Self {
                #wrapper::__from_proto_i32(#inner_to_i32)
            }
        }

        impl ::core::convert::From<#wrapper> for #inner {
            fn from(value: #wrapper) -> Self {
                let n = value.__to_proto_i32();
                #i32_to_inner
            }
        }
    })
}

// ---------------------------------------------------------------------------
// Message wrappers
// ---------------------------------------------------------------------------

/// Emit a newtype `#[pyclass]` wrapping a model message: keyword constructor,
/// native typed getters/setters per field, `__repr__`/`__eq__`, and the
/// `From`/`Into` bridges to the bare model type.
fn message_wrapper(
    fq_name: &str,
    info: &MessageInfo,
    rust_path: &str,
    runtime: Runtime,
    metadata: &CodeGenMetadata,
) -> crate::error::Result<TokenStream> {
    let wrapper = wrapper_ident(fq_name);
    let class_name = python_class_name(fq_name);
    let inner: syn::Path = parse_path(rust_path, fq_name)?;

    // Build per-field accessor descriptors. Oneof fields are not yet expressed as
    // typed getters; they are skipped here (PR-follow-up) but still round-trip via
    // the inner type's serde when constructed elsewhere.
    let fields: Vec<&MessageField> = info
        .fields
        .iter()
        .filter(|f| f.oneof_variants.is_none())
        .collect();

    let packages = metadata_packages(metadata);
    let project_pkg_refs: BTreeSet<&str> = packages.iter().map(String::as_str).collect();
    let boxed = boxed_fields(metadata, runtime);
    let accessors: Vec<FieldAccessor> = fields
        .iter()
        .map(|f| {
            let is_boxed = boxed.contains(&(fq_name.to_string(), f.name.clone()));
            FieldAccessor::new(f, runtime, &packages, &project_pkg_refs, is_boxed)
        })
        .collect();

    let getters = accessors.iter().map(FieldAccessor::getter);
    let setters = accessors.iter().map(FieldAccessor::setter);
    let ctor = constructor(&accessors, &inner);

    Ok(quote! {
        #[::pyo3::pyclass(name = #class_name)]
        #[derive(Clone, Debug)]
        pub struct #wrapper(pub #inner);

        #[allow(clippy::too_many_arguments, clippy::useless_conversion)]
        #[::pyo3::pymethods]
        impl #wrapper {
            #ctor

            #(#getters)*
            #(#setters)*

            fn __repr__(&self) -> ::std::string::String {
                ::std::format!("{:?}", self.0)
            }

            fn __eq__(&self, other: &Self) -> bool {
                self.0 == other.0
            }
        }

        impl ::core::convert::From<#inner> for #wrapper {
            fn from(value: #inner) -> Self {
                Self(value)
            }
        }

        impl ::core::convert::From<#wrapper> for #inner {
            fn from(value: #wrapper) -> Self {
                value.0
            }
        }
    })
}

/// Project package set, recomputed lazily for `FieldAccessor` so it can decide
/// whether a referenced message/enum is in-project (gets a wrapper) or foreign.
fn metadata_packages(metadata: &CodeGenMetadata) -> BTreeSet<String> {
    metadata
        .services
        .values()
        .map(|s| s.package.clone())
        .collect()
}

/// The set of `(message_fq_name, field_name)` pairs that the prost model plugin
/// stores as `Option<Box<T>>` rather than `Option<T>`.
///
/// prost only boxes a singular message field when leaving it unboxed would make
/// the struct infinitely sized — i.e. the field's message type can transitively
/// reach back to the containing message through a chain of singular (non-repeated,
/// non-map) message fields. Repeated and map fields already introduce an
/// indirection (`Vec`/`HashMap`), so they break the cycle and are never boxed.
///
/// We replicate that rule here so the PyO3 read/write accessors deref/box exactly
/// the fields prost boxed. Returns an empty set for the buffa runtime, which uses
/// `MessageField<T>` uniformly and never boxes.
fn boxed_fields(metadata: &CodeGenMetadata, runtime: Runtime) -> BTreeSet<(String, String)> {
    let mut boxed = BTreeSet::new();
    if runtime != Runtime::Prost {
        return boxed;
    }
    for (msg_name, info) in &metadata.messages {
        for field in &info.fields {
            // Only singular message fields can be boxed; repeated/map fields carry
            // their own indirection, and scalars/enums are never boxed.
            if field.unified_type.is_repeated
                || matches!(field.unified_type.base_type, BaseType::Map(_, _))
            {
                continue;
            }
            let BaseType::Message(target) = &field.unified_type.base_type else {
                continue;
            };
            // Box iff the field's type can reach back to the containing message.
            if reaches(metadata, target, msg_name) {
                boxed.insert((msg_name.clone(), field.name.clone()));
            }
        }
    }
    boxed
}

/// Whether message `from` can transitively reach message `target` by following
/// singular (non-repeated, non-map) message fields. Used to detect the recursion
/// cycles that drive prost's automatic field boxing.
fn reaches(metadata: &CodeGenMetadata, from: &str, target: &str) -> bool {
    let mut stack = vec![from.to_string()];
    let mut seen = BTreeSet::new();
    while let Some(cur) = stack.pop() {
        if cur == target {
            return true;
        }
        if !seen.insert(cur.clone()) {
            continue;
        }
        let Some(info) = metadata.messages.get(&cur) else {
            continue;
        };
        for field in &info.fields {
            if field.unified_type.is_repeated
                || matches!(field.unified_type.base_type, BaseType::Map(_, _))
            {
                continue;
            }
            if let BaseType::Message(next) = &field.unified_type.base_type {
                stack.push(next.clone());
            }
        }
    }
    false
}

/// A single message field, resolved into the Rust expressions needed to expose it
/// across the PyO3 boundary as a native, typed Python attribute.
struct FieldAccessor {
    /// The Python/Rust-visible field identifier (snake_case, keyword-sanitized).
    ident: proc_macro2::Ident,
    /// The original field, for cardinality flags and base type.
    ty: UnifiedType,
    runtime: Runtime,
    /// Whether the referenced message/enum is in the project (gets a wrapper).
    in_project: bool,
    /// In-crate Rust path of a referenced message/enum (e.g. `super::catalog::v1::Color`),
    /// when in-project; used by enum/message read expressions. `None` otherwise.
    inner_rust_path: Option<TokenStream>,
    /// Whether the prost model stores this singular message field as
    /// `Option<Box<T>>` (it sits on a recursion cycle). Drives whether the read
    /// path derefs the box and the write path re-boxes. Always `false` for buffa.
    is_boxed: bool,
}

impl FieldAccessor {
    fn new(
        field: &MessageField,
        runtime: Runtime,
        packages: &BTreeSet<String>,
        project_pkg_refs: &BTreeSet<&str>,
        is_boxed: bool,
    ) -> Self {
        let ident = format_ident!("{}", super::sanitize_python_field_name(&field.name));
        let in_project = match &field.unified_type.base_type {
            BaseType::Message(n) | BaseType::Enum(n) => is_in_project(n, packages),
            _ => false,
        };
        let inner_rust_path = match &field.unified_type.base_type {
            BaseType::Message(n) | BaseType::Enum(n) if in_project => {
                type_rust_path(n, project_pkg_refs)
                    .and_then(|p| syn::parse_str::<syn::Path>(&p).ok())
                    .map(|p| quote! { #p })
            }
            _ => None,
        };
        Self {
            ident,
            ty: field.unified_type.clone(),
            runtime,
            in_project,
            inner_rust_path,
            is_boxed,
        }
    }

    /// The Python-visible Rust type for this field (the getter return / setter
    /// arg). Messages/enums map to their `Py`-prefixed wrapper; collections wrap
    /// in `Vec`/`HashMap`; optional non-collection fields wrap in `Option`.
    fn py_type(&self) -> TokenStream {
        let base = self.base_py_type();
        if self.ty.is_repeated {
            return quote! { ::std::vec::Vec<#base> };
        }
        if let BaseType::Map(_, v) = &self.ty.base_type {
            let val = py_type_for(v, self.in_project);
            return quote! { ::std::collections::HashMap<::std::string::String, #val> };
        }
        // Singular (non-repeated) message fields carry proto presence semantics, so
        // they are always Python-visible as `Optional[...]` — the read path returns
        // `Option<Wrapper>` for both runtimes (buffa `MessageField::into_option`,
        // prost `Option<T>` or `Option<Box<T>>` for recursion-boxed fields).
        if self.is_singular_message() {
            return quote! { ::core::option::Option<#base> };
        }
        // Explicit-presence scalars/enums (proto3 `optional`) are stored as
        // `Option<T>` by both the buffa and prost model plugins, so the wrapper
        // surfaces them as `Optional[...]` too.
        if self.is_optional_scalar() {
            return quote! { ::core::option::Option<#base> };
        }
        base
    }

    fn base_py_type(&self) -> TokenStream {
        py_type_for(&self.ty, self.in_project)
    }

    fn is_singular_message(&self) -> bool {
        !self.ty.is_repeated
            && matches!(self.ty.base_type, BaseType::Message(_) | BaseType::OneOf(_))
    }

    /// Whether [`Self::py_type`] is itself an `Option<..>` — i.e. the field is a
    /// singular message or an explicit-presence scalar/enum. Such fields must not be
    /// re-wrapped in the keyword constructor (would yield ambiguous `Option<Option<T>>`).
    fn py_type_is_optional(&self) -> bool {
        self.is_singular_message() || self.is_optional_scalar()
    }

    /// An explicit-presence scalar or enum (proto3 `optional`): stored as
    /// `Option<T>` (scalar) or `Option<EnumValue<E>>`/`Option<i32>` (enum) by the
    /// model plugins, as opposed to messages whose presence is intrinsic.
    fn is_optional_scalar(&self) -> bool {
        self.ty.is_optional
            && !self.ty.is_repeated
            && !matches!(
                self.ty.base_type,
                BaseType::Message(_) | BaseType::OneOf(_) | BaseType::Map(_, _)
            )
    }

    /// `#[getter]` returning the field as its Python-visible type.
    fn getter(&self) -> TokenStream {
        let ident = &self.ident;
        let ret = self.py_type();
        let body = self.read_expr();
        quote! {
            #[getter]
            fn #ident(&self) -> #ret {
                #body
            }
        }
    }

    /// `#[setter]` accepting the Python-visible type and storing into the inner.
    fn setter(&self) -> TokenStream {
        let ident = &self.ident;
        let setter_name = format_ident!("set_{}", self.ident);
        let arg = self.py_type();
        let store = self.write_expr();
        quote! {
            #[setter(#ident)]
            fn #setter_name(&mut self, value: #arg) {
                self.0.#ident = #store;
            }
        }
    }

    /// Expression that reads `self.0.<field>` and converts it to the Python type.
    fn read_expr(&self) -> TokenStream {
        let ident = &self.ident;
        let access = quote! { self.0.#ident };

        if self.ty.is_repeated {
            return self.read_repeated(&access);
        }
        if let BaseType::Map(_, v) = &self.ty.base_type {
            return read_map(&access, v, self.in_project);
        }
        if self.is_optional_scalar() {
            return self.read_optional_scalar(&access);
        }

        match &self.ty.base_type {
            BaseType::String => quote! { #access.clone() },
            BaseType::Bytes => quote! { #access.clone() },
            BaseType::Int32
            | BaseType::Int64
            | BaseType::Bool
            | BaseType::Float32
            | BaseType::Float64 => quote! { #access },
            BaseType::Enum(_) => self.read_enum(&access),
            BaseType::Message(_) | BaseType::OneOf(_) => self.read_message(&access),
            BaseType::Unit | BaseType::Map(_, _) => quote! { #access.clone() },
        }
    }

    fn read_repeated(&self, access: &TokenStream) -> TokenStream {
        let base = self.base_py_type();
        let inner = self.inner_rust_path.as_ref();
        match &self.ty.base_type {
            BaseType::Enum(_) => match self.runtime {
                // buffa repeated enum: Vec<EnumValue<E>>; keep known values only.
                Runtime::Buffa => quote! {
                    #access.iter().filter_map(|e| e.as_known().map(#base::from)).collect()
                },
                // prost repeated enum: Vec<i32>; map through the bare Rust enum.
                Runtime::Prost => quote! {
                    #access.iter().filter_map(|&v| <#inner as ::core::convert::TryFrom<i32>>::try_from(v).ok().map(#base::from)).collect()
                },
            },
            BaseType::Message(_) => quote! {
                #access.iter().cloned().map(#base::from).collect()
            },
            BaseType::String | BaseType::Bytes => quote! { #access.clone() },
            _ => quote! { #access.clone() },
        }
    }

    fn read_enum(&self, access: &TokenStream) -> TokenStream {
        let base = self.base_py_type();
        let inner = self.inner_rust_path.as_ref();
        match self.runtime {
            // buffa: EnumValue<E> -> Option<E> via as_known(); fall back to the inner
            // enum's Default (the zero/unspecified variant), then to the wrapper.
            Runtime::Buffa => quote! {
                #base::from(#access.as_known().unwrap_or_default())
            },
            // prost: i32 -> E via TryFrom, fall back to the inner enum's Default.
            Runtime::Prost => quote! {
                #base::from(
                    <#inner as ::core::convert::TryFrom<i32>>::try_from(#access)
                        .unwrap_or_default()
                )
            },
        }
    }

    fn read_message(&self, access: &TokenStream) -> TokenStream {
        let base = self.base_py_type();
        match self.runtime {
            // buffa: MessageField<T> -> Option<T>
            Runtime::Buffa => quote! {
                #access.clone().into_option().map(#base::from)
            },
            // prost recursion-boxed field: Option<Box<T>> — deref the box.
            Runtime::Prost if self.is_boxed => quote! {
                #access.clone().map(|b| #base::from(*b))
            },
            // prost plain singular message: Option<T>.
            Runtime::Prost => quote! {
                #access.clone().map(#base::from)
            },
        }
    }

    /// Read an explicit-presence (proto3 `optional`) scalar/enum field, whose
    /// stored type is `Option<..>` — surfaced to Python as `Option<base>`.
    fn read_optional_scalar(&self, access: &TokenStream) -> TokenStream {
        let base = self.base_py_type();
        let inner = self.inner_rust_path.as_ref();
        match &self.ty.base_type {
            // buffa: Option<EnumValue<E>>; prost: Option<i32>.
            BaseType::Enum(_) => match self.runtime {
                Runtime::Buffa => quote! {
                    #access.as_ref().and_then(|e| e.as_known()).map(#base::from)
                },
                Runtime::Prost => quote! {
                    #access.and_then(|v| <#inner as ::core::convert::TryFrom<i32>>::try_from(v).ok())
                        .map(#base::from)
                },
            },
            // Option<{int,bool,float}>: the inner scalar is `Copy`, so the `Option`
            // copies by value — `.clone()` here would trip `clippy::clone_on_copy`.
            BaseType::Int32
            | BaseType::Int64
            | BaseType::Bool
            | BaseType::Float32
            | BaseType::Float64 => quote! { #access },
            // Option<String>/Option<Vec<u8>>: non-Copy, clone through.
            _ => quote! { #access.clone() },
        }
    }

    /// Expression converting a Python-typed `value` back into the inner field
    /// representation for storage.
    fn write_expr(&self) -> TokenStream {
        if self.ty.is_repeated {
            return self.write_repeated();
        }
        if let BaseType::Map(_, _) = &self.ty.base_type {
            return quote! { value };
        }
        if self.is_optional_scalar() {
            return self.write_optional_scalar();
        }
        match &self.ty.base_type {
            BaseType::Enum(_) => self.write_enum(),
            BaseType::Message(_) | BaseType::OneOf(_) => self.write_message(),
            _ => quote! { value },
        }
    }

    /// Store an explicit-presence scalar/enum: `value` is `Option<base>`.
    fn write_optional_scalar(&self) -> TokenStream {
        let inner = self.inner_rust_path.as_ref();
        match &self.ty.base_type {
            // buffa: Option<EnumValue<E>>; prost: Option<i32>.
            BaseType::Enum(_) => match self.runtime {
                Runtime::Buffa => quote! {
                    value.map(|e| ::buffa::EnumValue::Known(<#inner as ::core::convert::From<_>>::from(e)))
                },
                Runtime::Prost => quote! {
                    value.map(|e| <#inner as ::core::convert::From<_>>::from(e) as i32)
                },
            },
            // Option<String> etc. store directly.
            _ => quote! { value },
        }
    }

    fn write_repeated(&self) -> TokenStream {
        let inner = self.inner_rust_path.as_ref();
        match &self.ty.base_type {
            BaseType::Enum(_) => match self.runtime {
                Runtime::Buffa => quote! {
                    value.into_iter()
                        .map(|e| ::buffa::EnumValue::Known(<#inner as ::core::convert::From<_>>::from(e)))
                        .collect()
                },
                Runtime::Prost => quote! {
                    value.into_iter()
                        .map(|e| <#inner as ::core::convert::From<_>>::from(e) as i32)
                        .collect()
                },
            },
            BaseType::Message(_) => quote! {
                value.into_iter().map(::core::convert::Into::into).collect()
            },
            _ => quote! { value },
        }
    }

    fn write_enum(&self) -> TokenStream {
        let inner = self.inner_rust_path.as_ref();
        match self.runtime {
            Runtime::Buffa => quote! {
                ::buffa::EnumValue::Known(<#inner as ::core::convert::From<_>>::from(value))
            },
            Runtime::Prost => quote! {
                <#inner as ::core::convert::From<_>>::from(value) as i32
            },
        }
    }

    /// Singular message write: `value` is always `Option<Wrapper>` (see
    /// [`Self::is_singular_message`]); store into the runtime's optional-message
    /// representation.
    fn write_message(&self) -> TokenStream {
        match self.runtime {
            // buffa: MessageField<T>
            Runtime::Buffa => quote! {
                value.map(|w| ::buffa::MessageField::some(w.into())).unwrap_or_default()
            },
            // prost recursion-boxed field: Option<Box<T>> — re-box on store.
            Runtime::Prost if self.is_boxed => quote! {
                value.map(|w| ::std::boxed::Box::new(w.into()))
            },
            // prost plain singular message: Option<T>.
            Runtime::Prost => quote! {
                value.map(|w| w.into())
            },
        }
    }
}

/// The Python-visible base type (no cardinality wrapping) for a unified type.
fn py_type_for(ty: &UnifiedType, in_project: bool) -> TokenStream {
    match &ty.base_type {
        BaseType::String => quote! { ::std::string::String },
        BaseType::Bytes => quote! { ::std::vec::Vec<u8> },
        BaseType::Int32 => quote! { i32 },
        BaseType::Int64 => quote! { i64 },
        BaseType::Bool => quote! { bool },
        BaseType::Float32 => quote! { f32 },
        BaseType::Float64 => quote! { f64 },
        BaseType::Unit => quote! { () },
        BaseType::Message(n) | BaseType::Enum(n) => {
            if in_project {
                let w = wrapper_ident(n);
                quote! { #w }
            } else {
                // Foreign type (no wrapper): fall back to the bare simple name.
                let id = format_ident!("{}", extract_simple_type_name(n));
                quote! { #id }
            }
        }
        BaseType::OneOf(n) => {
            let id = format_ident!("{}", extract_simple_type_name(n));
            quote! { #id }
        }
        BaseType::Map(_, v) => {
            let val = py_type_for(v, in_project);
            quote! { ::std::collections::HashMap<::std::string::String, #val> }
        }
    }
}

/// Read a `map<string, V>` field into a `HashMap`, converting message values to
/// their wrappers when in-project.
fn read_map(access: &TokenStream, value: &UnifiedType, in_project: bool) -> TokenStream {
    match &value.base_type {
        BaseType::Message(_) if in_project => {
            let w = py_type_for(value, true);
            quote! {
                #access.iter().map(|(k, v)| (k.clone(), #w::from(v.clone()))).collect()
            }
        }
        _ => quote! { #access.clone() },
    }
}

/// The keyword `#[new]` constructor: every field is an optional keyword argument
/// defaulting to `None`; provided values are stored, omitted ones fall back to the
/// inner type's `Default`.
fn constructor(accessors: &[FieldAccessor], inner: &syn::Path) -> TokenStream {
    if accessors.is_empty() {
        return quote! {
            #[new]
            fn new() -> Self {
                Self(::core::default::Default::default())
            }
        };
    }

    let params = accessors.iter().map(|a| {
        let ident = &a.ident;
        let ty = a.py_type();
        // Fields whose Python type is *already* `Option<..>` (singular messages,
        // explicit-presence scalars) must not be double-wrapped: `Option<Option<T>>`
        // makes pyo3's `#[new]` arg-coercion ambiguous (`SomeWrap`). For those the
        // single `Option` already encodes the "omitted/None" keyword default.
        if a.py_type_is_optional() {
            quote! { #ident: #ty }
        } else {
            quote! { #ident: ::core::option::Option<#ty> }
        }
    });

    let sig_parts = accessors
        .iter()
        .map(|a| format!("{} = None", a.ident))
        .collect::<Vec<_>>()
        .join(", ");
    let sig: TokenStream = sig_parts.parse().expect("field idents are valid tokens");

    let assigns = accessors.iter().map(|a| {
        let ident = &a.ident;
        // Reuse the setter's write expression by binding `value`.
        let store = a.write_expr();
        if a.py_type_is_optional() {
            // `value` is the whole (already-`Option`) argument; store unconditionally
            // (an omitted kwarg arrives as `None`, the field's natural absent value).
            quote! {
                {
                    let value = #ident;
                    inner.#ident = #store;
                }
            }
        } else {
            quote! {
                if let ::core::option::Option::Some(value) = #ident {
                    inner.#ident = #store;
                }
            }
        }
    });

    quote! {
        #[new]
        #[pyo3(signature = (#sig))]
        fn new(#(#params,)*) -> Self {
            let mut inner = <#inner as ::core::default::Default>::default();
            #(#assigns)*
            Self(inner)
        }
    }
}

// ---------------------------------------------------------------------------
// Path helpers (shared with the former pythonize impl)
// ---------------------------------------------------------------------------

fn parse_path(rust_path: &str, fq_name: &str) -> crate::error::Result<syn::Path> {
    syn::parse_str(rust_path).map_err(|e| crate::error::Error::InvalidRustPath {
        path: rust_path.to_string(),
        message: fq_name.to_string(),
        source: e,
    })
}

/// Whether a fully-qualified type name belongs to one of the project's packages.
fn is_in_project(fq_name: &str, packages: &BTreeSet<String>) -> bool {
    let stripped = fq_name.strip_prefix('.').unwrap_or(fq_name);
    packages.iter().any(|pkg| {
        stripped == pkg
            || stripped
                .strip_prefix(pkg.as_str())
                .is_some_and(|rest| rest.starts_with('.'))
    })
}

/// Map a fully-qualified proto type name to its in-crate Rust path, as exposed by
/// the model module tree `generate_models_mod` emits.
///
/// `pyo3_impls.rs` is a sibling of the per-package model modules inside the models
/// subdir, so the path is `super::<svc_seg>::<ver_seg>::<...nested>::<Name>` where
/// `(svc_seg, ver_seg)` are the second-to-last and last segments of the owning
/// package — exactly the module nesting `generate_models_mod` produces.
///
/// Returns `None` for types outside the project's own packages (e.g. imported
/// google/gnostic types), which have no in-crate path and must be skipped.
fn type_rust_path(fq_name: &str, project_packages: &BTreeSet<&str>) -> Option<String> {
    let stripped = fq_name.strip_prefix('.').unwrap_or(fq_name);

    let package = project_packages
        .iter()
        .filter(|pkg| {
            stripped == **pkg
                || stripped
                    .strip_prefix(**pkg)
                    .is_some_and(|rest| rest.starts_with('.'))
        })
        .max_by_key(|pkg| pkg.len())?;

    let remainder = stripped.strip_prefix(package)?.trim_start_matches('.');
    if remainder.is_empty() {
        return None;
    }

    let pkg_parts: Vec<&str> = package.split('.').collect();
    let module_path = if pkg_parts.len() >= 2 {
        format!(
            "{}::{}",
            pkg_parts[pkg_parts.len() - 2],
            pkg_parts[pkg_parts.len() - 1]
        )
    } else {
        pkg_parts[0].to_string()
    };

    // `remainder` is the dot-joined type path within the package, e.g.
    // `GenerateTableCredentialsRequest.Operation` for a message-nested enum. Both
    // the prost and buffa model trees nest a type defined inside a message under a
    // *snake_case module* named after that parent message, with only the leaf type
    // keeping its proto (PascalCase) name. So snake_case every segment except the
    // last (`generate_table_credentials_request::Operation`), not a flat `::` join.
    let segments: Vec<&str> = remainder.split('.').collect();
    let (leaf, parents) = segments.split_last()?;
    let type_path = parents
        .iter()
        .map(|seg| seg.to_case(Case::Snake))
        .chain(std::iter::once((*leaf).to_string()))
        .collect::<Vec<_>>()
        .join("::");
    Some(format!("super::{}::{}", module_path, type_path))
}

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

    fn pkgs(list: &[&'static str]) -> BTreeSet<&'static str> {
        list.iter().copied().collect()
    }

    #[test]
    fn three_segment_package_maps_to_service_version_module() {
        let p = pkgs(&["example.catalog.v1"]);
        assert_eq!(
            type_rust_path(".example.catalog.v1.Catalog", &p).as_deref(),
            Some("super::catalog::v1::Catalog")
        );
    }

    #[test]
    fn two_segment_package_keeps_first_segment_as_module() {
        let p = pkgs(&["demo.v1"]);
        assert_eq!(
            type_rust_path(".demo.v1.Widget", &p).as_deref(),
            Some("super::demo::v1::Widget")
        );
    }

    #[test]
    fn nested_type_lives_in_snake_cased_parent_module() {
        let p = pkgs(&["example.catalog.v1"]);
        // A type defined inside a message is nested under a snake_case module named
        // after the parent message; only the leaf keeps its proto name.
        assert_eq!(
            type_rust_path(".example.catalog.v1.Catalog.Nested", &p).as_deref(),
            Some("super::catalog::v1::catalog::Nested")
        );
    }

    #[test]
    fn deeply_nested_type_snake_cases_every_parent() {
        let p = pkgs(&["example.creds.v1"]);
        assert_eq!(
            type_rust_path(
                ".example.creds.v1.GenerateTableCredentialsRequest.Operation",
                &p
            )
            .as_deref(),
            Some("super::creds::v1::generate_table_credentials_request::Operation")
        );
    }

    #[test]
    fn foreign_package_is_skipped() {
        let p = pkgs(&["example.catalog.v1"]);
        assert_eq!(type_rust_path(".google.protobuf.Any", &p), None);
    }

    #[test]
    fn wrapper_ident_is_py_prefixed_qualified_name() {
        assert_eq!(
            wrapper_ident(".example.catalog.v1.Catalog").to_string(),
            "PyCatalog"
        );
        assert_eq!(
            wrapper_ident(".example.catalog.v1.Catalog.Nested").to_string(),
            "PyCatalogNested"
        );
    }

    // --- boxing detection -------------------------------------------------

    fn msg_field(name: &str, ty: UnifiedType) -> MessageField {
        MessageField {
            name: name.to_string(),
            unified_type: ty,
            documentation: None,
            oneof_variants: None,
            field_behavior: vec![],
            is_sensitive: false,
            resource_reference: None,
        }
    }

    fn singular_message(target: &str) -> UnifiedType {
        UnifiedType {
            base_type: BaseType::Message(target.to_string()),
            is_optional: true,
            is_repeated: false,
        }
    }

    fn repeated_message(target: &str) -> UnifiedType {
        UnifiedType {
            base_type: BaseType::Message(target.to_string()),
            is_optional: false,
            is_repeated: true,
        }
    }

    fn metadata_with(messages: Vec<(&str, Vec<MessageField>)>) -> CodeGenMetadata {
        let mut meta = CodeGenMetadata::default();
        for (name, fields) in messages {
            meta.messages.insert(
                name.to_string(),
                MessageInfo {
                    name: name.to_string(),
                    fields,
                    resource_descriptor: None,
                    documentation: None,
                },
            );
        }
        meta
    }

    #[test]
    fn non_recursive_singular_message_is_not_boxed() {
        // A.f: B, B has no message fields -> no cycle -> unboxed.
        let meta = metadata_with(vec![
            (".pkg.A", vec![msg_field("f", singular_message(".pkg.B"))]),
            (".pkg.B", vec![]),
        ]);
        assert!(boxed_fields(&meta, Runtime::Prost).is_empty());
    }

    #[test]
    fn direct_self_recursion_boxes_the_field() {
        // A.next: A -> cycle -> boxed.
        let meta = metadata_with(vec![(
            ".pkg.A",
            vec![msg_field("next", singular_message(".pkg.A"))],
        )]);
        let boxed = boxed_fields(&meta, Runtime::Prost);
        assert!(boxed.contains(&(".pkg.A".to_string(), "next".to_string())));
        assert_eq!(boxed.len(), 1);
    }

    #[test]
    fn mutual_recursion_boxes_both_back_edges() {
        // A.b: B and B.a: A -> A<->B cycle -> both boxed.
        let meta = metadata_with(vec![
            (".pkg.A", vec![msg_field("b", singular_message(".pkg.B"))]),
            (".pkg.B", vec![msg_field("a", singular_message(".pkg.A"))]),
        ]);
        let boxed = boxed_fields(&meta, Runtime::Prost);
        assert!(boxed.contains(&(".pkg.A".to_string(), "b".to_string())));
        assert!(boxed.contains(&(".pkg.B".to_string(), "a".to_string())));
    }

    #[test]
    fn repeated_field_breaks_the_cycle_and_is_not_boxed() {
        // A.items: repeated A -> Vec breaks the size cycle -> never boxed.
        let meta = metadata_with(vec![(
            ".pkg.A",
            vec![msg_field("items", repeated_message(".pkg.A"))],
        )]);
        assert!(boxed_fields(&meta, Runtime::Prost).is_empty());
    }

    #[test]
    fn buffa_runtime_never_boxes() {
        let meta = metadata_with(vec![(
            ".pkg.A",
            vec![msg_field("next", singular_message(".pkg.A"))],
        )]);
        assert!(boxed_fields(&meta, Runtime::Buffa).is_empty());
    }
}