use ruff_db::parsed::parsed_module;
use ruff_python_ast::{Expr, ExprCall, ExprDict, Keyword, name::Name};
use rustc_hash::FxHashSet;
use ty_module_resolver::{KnownModule, file_to_module};
use ty_python_core::definition::{Definition, DefinitionKind};

use crate::place::{DefinedPlace, Definedness, Place, Provenance, known_module_symbol};
use crate::types::class::CodeGeneratorKind;
use crate::types::ide_support::{ImportAliasResolution, definitions_for_name};
use crate::types::known_instance::FieldInstance;
use crate::types::member::class_member;
use crate::types::special_form::SpecialFormType;
use crate::types::{
    ClassBase, DataclassTransformerParams, FunctionType, KnownClass, KnownFunction,
    KnownInstanceType, KnownUnion, Parameter, Specialization, StaticClassLiteral, Type, UnionType,
    definition_expression_type,
};
use crate::{Db, SemanticModel};

/// Metadata that controls Pydantic-specific model synthesis.
#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
pub(crate) struct ModelMetadata<'db> {
    #[returns(deref)]
    pub(in crate::types) field_specifiers: Box<[Type<'db>]>,
    #[returns(copy)]
    config: ModelConfig,
}

impl get_size2::GetSize for ModelMetadata<'_> {}

impl<'db> ModelMetadata<'db> {
    pub(in crate::types) fn from_class(
        db: &'db dyn Db,
        class: StaticClassLiteral<'db>,
        transformer_params: DataclassTransformerParams<'db>,
    ) -> Self {
        Self::new(
            db,
            transformer_params.field_specifiers(db),
            model_config(db, class),
        )
    }

    fn accepts_extra(self, db: &'db dyn Db) -> bool {
        !matches!(self.config(db).extra, Some(ExtraBehavior::Forbid))
    }

    pub(in crate::types) fn validates_by_alias(self, db: &'db dyn Db) -> bool {
        self.config(db).validate_by_alias.enabled_or(true)
    }

    pub(in crate::types) fn validates_by_name(self, db: &'db dyn Db) -> bool {
        let config = self.config(db);
        let validate_by_name = config.validate_by_name;
        // If `validate_by_alias=False` is set without specifying `validate_by_name`, Pydantic
        // implicitly enables validation by name.
        if matches!(validate_by_name, ConfigBoolean::Unspecified)
            && matches!(config.validate_by_alias, ConfigBoolean::Disabled)
        {
            true
        } else {
            validate_by_name.enabled_or(false)
        }
    }

    pub(in crate::types) fn is_frozen(self, db: &'db dyn Db) -> bool {
        matches!(self.config(db).frozen, ConfigBoolean::Enabled)
    }
}

/// Pydantic-specific metadata resolved from a field's annotation (via `Annotation`)
/// and right-hand side `Field(...)` specifier.
///
/// For example:
/// ```py
/// class Model(BaseModel):
///     value: Annotated[int, Strict()] = Field(default=0)
/// ```
pub(in crate::types) struct FieldMetadata<'db> {
    pub(in crate::types) default_ty: Option<Type<'db>>,
    pub(in crate::types) init: bool,
    pub(in crate::types) alias: Option<Box<str>>,
    pub(in crate::types) strict: ConfigBoolean,
}

impl Default for FieldMetadata<'_> {
    fn default() -> Self {
        Self {
            default_ty: None,
            init: true,
            alias: None,
            strict: ConfigBoolean::Unspecified,
        }
    }
}

impl<'db> FieldMetadata<'db> {
    /// Collect Pydantic field metadata from the right-hand side of a field's assignment.
    ///
    /// For example, collect the default value and alias from the following field assignment:
    /// ```py
    /// field: int = Field(default=0, alias="field_alias")
    /// ```
    fn collect_from_rhs_type(
        &mut self,
        db: &'db dyn Db,
        rhs_type: Option<Type<'db>>,
        specialization: Option<Specialization<'db>>,
    ) {
        match rhs_type {
            Some(Type::KnownInstance(KnownInstanceType::Field(field))) => {
                self.merge_field(db, field, specialization);
            }
            Some(rhs_type) => self.default_ty = Some(rhs_type),
            None => {}
        }
    }

    /// Collect Pydantic field metadata from a field's annotation.
    ///
    /// For example, collect the strictness metadata from the following field annotation:
    /// ```py
    /// field: Annotated[int, Strict()]
    /// ```
    ///
    /// This method also handles the case where the annotation is an alias to an `Annotated` type, such as:
    /// ```py
    /// field: StrictInt
    /// ```
    /// where `StrictInt` is defined as `StrictInt = Annotated[int, Strict()]`.
    fn collect_from_annotation(
        &mut self,
        db: &'db dyn Db,
        definition: Definition<'db>,
        specialization: Option<Specialization<'db>>,
    ) {
        let module = parsed_module(db, definition.file(db)).load(db);
        let DefinitionKind::AnnotatedAssignment(assignment) = definition.kind(db) else {
            return;
        };
        let annotation = assignment.annotation(&module);

        if self.collect_from_annotated(db, definition, annotation, specialization) {
            return;
        }

        let Expr::Name(name) = annotation else {
            return;
        };

        // The following part is unfortunate. Pydantic defines `StrictInt` and the other aliases
        // using `StrictInt = Annotated[int, Strict()]`. Since we don't retain the `Annotated`
        // metadata, we need to follow the alias back to its definition and parse the metadata
        // from there.
        let model = SemanticModel::new(db, definition.file(db));
        let Some(alias_definition) = definitions_for_name(
            &model,
            name.id.as_str(),
            name.into(),
            ImportAliasResolution::ResolveAliases,
        )
        .into_iter()
        .find_map(|resolved| resolved.definition()) else {
            return;
        };

        let module = parsed_module(db, alias_definition.file(db)).load(db);
        let kind = alias_definition.kind(db);
        let value = match &kind {
            DefinitionKind::Assignment(assignment) => assignment.value(&module),
            DefinitionKind::AnnotatedAssignment(assignment) => {
                let Some(value) = assignment.value(&module) else {
                    return;
                };
                value
            }
            _ => return,
        };

        self.collect_from_annotated(db, alias_definition, value, specialization);
    }

    /// Collect Pydantic field metadata from the `Annotated` part of a field's annotation.
    fn collect_from_annotated(
        &mut self,
        db: &'db dyn Db,
        definition: Definition<'db>,
        annotation: &Expr,
        specialization: Option<Specialization<'db>>,
    ) -> bool {
        let Some(subscript) = annotation.as_subscript_expr() else {
            return false;
        };
        if definition_expression_type(db, definition, &subscript.value)
            != Type::SpecialForm(SpecialFormType::Annotated)
        {
            return false;
        }
        let Some(arguments) = subscript
            .slice
            .as_tuple_expr()
            .and_then(|tuple| tuple.elts.get(1..))
        else {
            return false;
        };

        for metadata in arguments {
            let Some(call) = metadata.as_call_expr() else {
                continue;
            };
            let callee = definition_expression_type(db, definition, &call.func);

            if callee
                .as_class_literal()
                .is_some_and(|class| class.is_known(db, KnownClass::PydanticStrict))
            {
                let strict = call.arguments.find_argument_value("strict", 0).map_or(
                    ConfigBoolean::Enabled,
                    |strict| {
                        ConfigBoolean::from_type(definition_expression_type(db, definition, strict))
                    },
                );
                self.merge_strict(strict);
            } else if matches!(
                callee,
                Type::FunctionLiteral(function)
                    if function.is_known(db, KnownFunction::PydanticField)
            ) {
                let field_type = definition_expression_type(db, definition, metadata);
                if let Type::KnownInstance(KnownInstanceType::Field(field)) = field_type {
                    self.merge_field(db, field, specialization);
                } else {
                    self.merge_field_call(db, definition, call, field_type, specialization);
                }
            }
        }

        true
    }

    fn merge_strict(&mut self, strict: ConfigBoolean) {
        self.strict = strict;
    }

    fn merge_field(
        &mut self,
        db: &'db dyn Db,
        field: FieldInstance<'db>,
        specialization: Option<Specialization<'db>>,
    ) {
        if let Some(default_type) = field.default_type(db) {
            self.default_ty = Some(default_type.apply_optional_specialization(db, specialization));
        }
        self.init &= field.init(db);
        if let Some(alias) = field.alias(db) {
            self.alias = Some(alias.clone());
        }
        if !matches!(field.strict(db), ConfigBoolean::Unspecified) {
            self.strict = field.strict(db);
        }
    }

    fn merge_field_call(
        &mut self,
        db: &'db dyn Db,
        definition: Definition<'db>,
        call: &ExprCall,
        call_type: Type<'db>,
        specialization: Option<Specialization<'db>>,
    ) {
        if let Some(default) = call.arguments.find_argument_value("default", 0) {
            let default_type = definition_expression_type(db, definition, default);
            if !default_type.is_instance_of(db, KnownClass::EllipsisType) {
                self.default_ty =
                    Some(default_type.apply_optional_specialization(db, specialization));
            }
        } else if call.arguments.find_keyword("default_factory").is_some() {
            self.default_ty = Some(call_type.apply_optional_specialization(db, specialization));
        }

        if let Some(init) = call.arguments.find_keyword("init") {
            let init = definition_expression_type(db, definition, &init.value);
            self.init &= !init.bool(db).is_always_false();
        }

        if let Some(alias) = call
            .arguments
            .find_keyword("validation_alias")
            .or_else(|| call.arguments.find_keyword("alias"))
        {
            self.alias = definition_expression_type(db, definition, &alias.value)
                .as_string_literal()
                .map(|literal| Box::from(literal.value(db)));
        }

        if let Some(strict) = call.arguments.find_keyword("strict") {
            let strict = definition_expression_type(db, definition, &strict.value);
            if !strict.is_none(db) {
                self.merge_strict(ConfigBoolean::from_type(strict));
            }
        }
    }
}

/// Resolve a Pydantic field's metadata from its annotation and right-hand side.
pub(in crate::types) fn field_metadata<'db>(
    db: &'db dyn Db,
    definition: Option<Definition<'db>>,
    rhs_type: Option<Type<'db>>,
    specialization: Option<Specialization<'db>>,
) -> FieldMetadata<'db> {
    let mut metadata = FieldMetadata::default();
    if let Some(definition) = definition {
        metadata.collect_from_annotation(db, definition, specialization);
    }
    metadata.collect_from_rhs_type(db, rhs_type, specialization);
    metadata
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, get_size2::GetSize)]
pub(crate) struct ModelConfig {
    /// The `extra` configuration controls whether the synthesized constructor accepts keyword
    /// arguments that do not correspond to declared model fields.
    extra: Option<ExtraBehavior>,
    /// The `strict` configuration controls whether constructor parameters accept values that
    /// Pydantic can coerce to the declared field type.
    strict: ConfigBoolean,
    /// Whether assignments to fields on model instances are forbidden.
    frozen: ConfigBoolean,
    /// Whether fields with aliases can be initialized by their alias.
    validate_by_alias: ConfigBoolean,
    /// Whether fields with aliases can be initialized by their field name.
    validate_by_name: ConfigBoolean,
}

impl ModelConfig {
    const fn unknown() -> Self {
        Self {
            extra: Some(ExtraBehavior::Unknown),
            strict: ConfigBoolean::Unknown,
            frozen: ConfigBoolean::Unknown,
            validate_by_alias: ConfigBoolean::Unknown,
            validate_by_name: ConfigBoolean::Unknown,
        }
    }

    /// Merge `other` into this config, with values from `other` taking precedence.
    fn merge(&mut self, other: Self) {
        self.extra = other.extra.or(self.extra);
        self.strict = other.strict.or(self.strict);
        self.frozen = other.frozen.or(self.frozen);
        self.validate_by_alias = other.validate_by_alias.or(self.validate_by_alias);
        self.validate_by_name = other.validate_by_name.or(self.validate_by_name);
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, get_size2::GetSize)]
enum ExtraBehavior {
    Allow,
    Forbid,
    Ignore,
    Unknown,
}

impl ExtraBehavior {
    fn from_value(value: Option<&str>) -> Self {
        match value {
            Some("allow") => Self::Allow,
            Some("forbid") => Self::Forbid,
            Some("ignore") => Self::Ignore,
            _ => Self::Unknown,
        }
    }
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, get_size2::GetSize)]
pub enum ConfigBoolean {
    /// No value was specified at this precedence level, so a lower-precedence value can apply.
    #[default]
    Unspecified,
    /// The setting was explicitly disabled.
    Disabled,
    /// The setting was explicitly enabled.
    Enabled,
    /// A value was specified, but it could not be resolved statically.
    Unknown,
}

impl ConfigBoolean {
    const fn or(self, other: Self) -> Self {
        if matches!(self, Self::Unspecified) {
            other
        } else {
            self
        }
    }

    /// Resolve a boolean configuration value from its inferred type.
    pub(in crate::types) fn from_type(value: Type<'_>) -> Self {
        if value == Type::bool_literal(true) {
            Self::Enabled
        } else if value == Type::bool_literal(false) {
            Self::Disabled
        } else {
            Self::Unknown
        }
    }

    /// Resolve this value to a boolean, using `default` when it is unspecified.
    ///
    /// An unknown value resolves to `true`.
    const fn enabled_or(self, default: bool) -> bool {
        match self {
            Self::Unspecified => default,
            Self::Disabled => false,
            Self::Enabled | Self::Unknown => true,
        }
    }
}

fn config_boolean(
    db: &dyn Db,
    definition: Definition<'_>,
    keyword: Option<&Keyword>,
) -> ConfigBoolean {
    keyword.map_or(ConfigBoolean::Unspecified, |keyword| {
        ConfigBoolean::from_type(definition_expression_type(db, definition, &keyword.value))
    })
}

pub(in crate::types) fn is_model(db: &dyn Db, class: StaticClassLiteral<'_>) -> bool {
    class
        .iter_mro(db, None)
        .filter_map(ClassBase::into_class)
        .any(|base| base.is_known(db, KnownClass::PydanticBaseModel))
}

/// Return whether a field specifier's `default` argument provides a default value.
///
/// Pydantic's `Field(...)` uses the ellipsis as a required-field sentinel, so it does not provide
/// a default value.
pub(in crate::types) fn field_provides_default(
    db: &dyn Db,
    function: FunctionType<'_>,
    default: Type<'_>,
) -> bool {
    !default.is_instance_of(db, KnownClass::EllipsisType)
        || !function.is_known(db, KnownFunction::PydanticField)
}

/// Return `true` if fields in `class` are keyword-only constructor parameters.
///
/// Pydantic model fields are generally keyword-only, but a root model's `root` field can also be
/// passed positionally.
pub(in crate::types) fn constructor_fields_are_keyword_only(
    db: &dyn Db,
    class: StaticClassLiteral<'_>,
) -> bool {
    !is_root_model(db, class)
}

#[salsa::tracked(returns(copy), heap_size=ruff_memory_usage::heap_size)]
fn is_root_model<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> bool {
    class
        .iter_mro(db, None)
        .filter_map(ClassBase::into_class)
        .any(|base| base.is_known(db, KnownClass::PydanticRootModel))
}

/// Return `true` if fields are optional constructor parameters for `class`.
///
/// A settings model can populate any field from environment variables or another configured
/// settings source, so no field value is necessarily required at the call site.
pub(in crate::types) fn constructor_fields_are_optional(
    db: &dyn Db,
    class: StaticClassLiteral<'_>,
) -> bool {
    class
        .iter_mro(db, None)
        .filter_map(ClassBase::into_class)
        .any(|base| base.is_known(db, KnownClass::PydanticBaseSettings))
}

#[salsa::tracked(
    returns(copy),
    cycle_initial=|_, _, _| ModelConfig::unknown(),
    heap_size=ruff_memory_usage::heap_size,
)]
fn model_config<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> ModelConfig {
    let mut config = ModelConfig::default();

    // Pydantic merges the effective config from each direct base from left to right. A later base
    // therefore takes precedence over an earlier base.
    for base in class.explicit_bases(db) {
        let Some(base) = base.to_class_type(db) else {
            config = ModelConfig::unknown();
            continue;
        };
        let base_literal = base.class_literal(db);
        let Some(base) = base_literal.as_static() else {
            config.merge(ModelConfig::unknown());
            continue;
        };

        if is_model(db, base) {
            config.merge(model_config(db, base));
        } else if let Some(base_config) = inherited_model_config(db, base) {
            config.merge(base_config);
        }
    }

    if let Some(own_config) = own_model_config(db, class) {
        config.merge(own_config);
    }
    config.merge(class_keyword_config(db, class));
    config
}

fn inherited_model_config(db: &dyn Db, class: StaticClassLiteral<'_>) -> Option<ModelConfig> {
    for base in class.iter_mro(db, None).filter_map(ClassBase::into_class) {
        let Some((base, _)) = base.static_class_literal(db) else {
            return Some(ModelConfig::unknown());
        };
        if let Some(config) = own_model_config(db, base) {
            return Some(config);
        }
    }
    None
}

fn own_model_config(db: &dyn Db, class: StaticClassLiteral<'_>) -> Option<ModelConfig> {
    let model_config = class_member(db, class.body_scope(db), "model_config")
        .inner
        .place;
    let Place::Defined(DefinedPlace {
        definedness: Definedness::AlwaysDefined,
        provenance: Provenance::SingleDefinition(definition),
        ..
    }) = model_config
    else {
        return if model_config.is_undefined() {
            None
        } else {
            Some(ModelConfig::unknown())
        };
    };

    let module = parsed_module(db, class.file(db)).load(db);
    let kind = definition.kind(db);
    let value = match &kind {
        DefinitionKind::Assignment(assignment) => assignment.value(&module),
        DefinitionKind::AnnotatedAssignment(assignment) => {
            let Some(value) = assignment.value(&module) else {
                return Some(ModelConfig::default());
            };
            value
        }
        _ => {
            return Some(ModelConfig::unknown());
        }
    };

    if let Some(dict) = value.as_dict_expr() {
        return Some(model_config_from_dict(db, definition, dict));
    }

    let Some(call) = value.as_call_expr() else {
        return Some(ModelConfig::unknown());
    };
    let callee = definition_expression_type(db, definition, &call.func);
    if !callee.as_class_literal().is_some_and(|class| {
        class.is_known(db, KnownClass::PydanticConfigDict) || class.is_known(db, KnownClass::Dict)
    }) {
        return Some(ModelConfig::unknown());
    }

    if call
        .arguments
        .keywords
        .iter()
        .any(|keyword| keyword.arg.is_none())
    {
        return Some(ModelConfig::unknown());
    }

    // Keep this list of recognized options in sync with `model_config_from_dict` and
    // `class_keyword_config`.
    let extra = call.arguments.find_keyword("extra").map(|extra| {
        let extra = definition_expression_type(db, definition, &extra.value)
            .as_string_literal()
            .map(|literal| literal.value(db));
        ExtraBehavior::from_value(extra)
    });
    let strict = config_boolean(db, definition, call.arguments.find_keyword("strict"));
    let frozen = config_boolean(db, definition, call.arguments.find_keyword("frozen"));
    let validate_by_alias = config_boolean(
        db,
        definition,
        call.arguments.find_keyword("validate_by_alias"),
    );
    let validate_by_name = config_boolean(
        db,
        definition,
        call.arguments.find_keyword("validate_by_name"),
    );

    Some(ModelConfig {
        extra,
        strict,
        frozen,
        validate_by_alias,
        validate_by_name,
    })
}

fn model_config_from_dict(db: &dyn Db, definition: Definition<'_>, dict: &ExprDict) -> ModelConfig {
    let mut config = ModelConfig::default();

    for item in dict {
        let Some(key) = item
            .key
            .as_ref()
            .and_then(|key| key.as_string_literal_expr())
        else {
            // A dynamic key or dictionary unpacking can override any recognized option. Ignoring
            // it could incorrectly preserve a lower-precedence configuration value.
            return ModelConfig::unknown();
        };
        let value = definition_expression_type(db, definition, &item.value);

        // Keep this match in sync with the options recognized for `ConfigDict` calls in
        // `own_model_config` and for class keywords in `class_keyword_config`.
        match key.value.to_str() {
            "extra" => {
                config.extra = Some(ExtraBehavior::from_value(
                    value.as_string_literal().map(|literal| literal.value(db)),
                ));
            }
            "strict" => config.strict = ConfigBoolean::from_type(value),
            "frozen" => config.frozen = ConfigBoolean::from_type(value),
            "validate_by_alias" => config.validate_by_alias = ConfigBoolean::from_type(value),
            "validate_by_name" => config.validate_by_name = ConfigBoolean::from_type(value),
            _ => {}
        }
    }

    config
}

fn class_keyword_config(db: &dyn Db, class: StaticClassLiteral<'_>) -> ModelConfig {
    let definition = class.definition(db);
    let module = parsed_module(db, class.file(db)).load(db);
    let kind = definition.kind(db);
    let Some(class) = kind.as_class() else {
        return ModelConfig::default();
    };
    let class_node = class.node(&module);
    let Some(arguments) = class_node.arguments.as_ref() else {
        return ModelConfig::default();
    };
    if arguments
        .keywords
        .iter()
        .any(|keyword| keyword.arg.is_none())
    {
        return ModelConfig::unknown();
    }
    // Keep this list of recognized options in sync with `own_model_config` and
    // `model_config_from_dict`.
    let extra = arguments.find_keyword("extra").map(|extra| {
        let extra = definition_expression_type(db, definition, &extra.value)
            .as_string_literal()
            .map(|literal| literal.value(db));
        ExtraBehavior::from_value(extra)
    });
    let strict = config_boolean(db, definition, arguments.find_keyword("strict"));
    let frozen = config_boolean(db, definition, arguments.find_keyword("frozen"));
    let validate_by_alias =
        config_boolean(db, definition, arguments.find_keyword("validate_by_alias"));
    let validate_by_name =
        config_boolean(db, definition, arguments.find_keyword("validate_by_name"));

    ModelConfig {
        extra,
        strict,
        frozen,
        validate_by_alias,
        validate_by_name,
    }
}

/// Return the input type accepted by a Pydantic field's synthesized constructor parameter.
pub(in crate::types) fn constructor_parameter_type<'db>(
    db: &'db dyn Db,
    field_type: Type<'db>,
    field_strict: ConfigBoolean,
    metadata: ModelMetadata<'db>,
) -> Type<'db> {
    if field_strict.or(metadata.config(db).strict) == ConfigBoolean::Enabled {
        return field_type;
    }

    lax_input_type(db, field_type)
}

/// Return the documented Python input type accepted by Pydantic for `field_type` in lax mode.
fn lax_input_type<'db>(db: &'db dyn Db, field_type: Type<'db>) -> Type<'db> {
    lax_input_type_impl(db, field_type, &mut FxHashSet::default())
}

fn lax_input_type_impl<'db>(
    db: &'db dyn Db,
    field_type: Type<'db>,
    expanding_types: &mut FxHashSet<Type<'db>>,
) -> Type<'db> {
    if field_type.is_none(db) || matches!(field_type, Type::LiteralValue(_) | Type::SubclassOf(_)) {
        return field_type;
    }

    if let Type::TypeAlias(alias) = field_type {
        if !expanding_types.insert(field_type) {
            return Type::any();
        }
        let result = lax_input_type_impl(db, alias.value_type(db), expanding_types);
        expanding_types.remove(&field_type);
        return result;
    }

    if field_type.as_union().and_then(|union| union.known(db)) == Some(KnownUnion::Float) {
        return lax_alias(db, "LaxFloat");
    }

    if let Type::Union(union) = field_type {
        return UnionType::from_elements_leave_aliases(
            db,
            union
                .elements(db)
                .iter()
                .map(|element| lax_input_type_impl(db, *element, expanding_types)),
        );
    }

    if let Some(input_type) = root_model_input_type(db, field_type, expanding_types) {
        return input_type;
    }

    let known_class = field_type
        .nominal_class(db)
        .and_then(|class| class.known(db));

    if matches!(
        known_class,
        Some(
            KnownClass::List
                | KnownClass::Set
                | KnownClass::FrozenSet
                | KnownClass::Deque
                | KnownClass::Sequence
                | KnownClass::Iterable
                | KnownClass::Tuple
        )
    ) {
        let Ok(elements) = field_type.try_iterate(db) else {
            return Type::any();
        };
        let element_type =
            lax_input_type_impl(db, elements.homogeneous_element_type(db), expanding_types);
        return KnownClass::Iterable.to_specialized_instance(db, &[element_type]);
    }

    if matches!(known_class, Some(KnownClass::Dict | KnownClass::Mapping)) {
        let Some(specialization) =
            known_class.and_then(|known_class| field_type.known_specialization(db, known_class))
        else {
            return Type::any();
        };
        let [key_type, value_type] = specialization.types(db) else {
            return Type::any();
        };
        let value_type = lax_input_type_impl(db, *value_type, expanding_types);
        return KnownClass::Mapping.to_specialized_instance(db, &[*key_type, value_type]);
    }

    let builtin_alias = match known_class {
        Some(KnownClass::Bool) => Some("LaxBool"),
        Some(KnownClass::Bytes) => Some("LaxBytes"),
        Some(KnownClass::Float) => Some("LaxFloat"),
        Some(KnownClass::Int) => Some("LaxInt"),
        Some(KnownClass::Str) => Some("LaxStr"),
        Some(KnownClass::Path) => Some("LaxPath"),
        _ => None,
    };
    if let Some(alias) = builtin_alias {
        return lax_alias(db, alias);
    }

    let Some((module, symbol, class)) = instance_symbol(db, field_type) else {
        return Type::any();
    };
    let symbol_alias = match (module, symbol) {
        (KnownModule::Datetime, "date") => Some("LaxDate"),
        (KnownModule::Datetime, "datetime") => Some("LaxDatetime"),
        (KnownModule::Datetime, "time") => Some("LaxTime"),
        (KnownModule::Datetime, "timedelta") => Some("LaxTimedelta"),
        (KnownModule::Decimal, "Decimal") => Some("LaxDecimal"),
        (KnownModule::Uuid, "UUID") => Some("LaxUUID"),
        (KnownModule::Ipaddress, "IPv4Address") => Some("LaxIPv4Address"),
        (KnownModule::Ipaddress, "IPv4Interface") => Some("LaxIPv4Interface"),
        (KnownModule::Ipaddress, "IPv4Network") => Some("LaxIPv4Network"),
        (KnownModule::Ipaddress, "IPv6Address") => Some("LaxIPv6Address"),
        (KnownModule::Ipaddress, "IPv6Interface") => Some("LaxIPv6Interface"),
        (KnownModule::Ipaddress, "IPv6Network") => Some("LaxIPv6Network"),
        (KnownModule::PydanticTypes, "ByteSize") => Some("LaxByteSize"),
        _ => None,
    };
    if let Some(alias) = symbol_alias {
        return lax_alias(db, alias);
    }

    let alias = if (module, symbol) == (KnownModule::Re, "Pattern") {
        let Some(specialization) = field_type.specialization_of(db, class) else {
            return Type::any();
        };
        let [pattern_type] = specialization.types(db) else {
            return Type::any();
        };
        if pattern_type
            .nominal_class(db)
            .is_some_and(|class| class.is_known(db, KnownClass::Str))
        {
            "LaxStrPattern"
        } else if pattern_type
            .nominal_class(db)
            .is_some_and(|class| class.is_known(db, KnownClass::Bytes))
        {
            "LaxBytesPattern"
        } else {
            return Type::any();
        }
    } else {
        return Type::any();
    };

    lax_alias(db, alias)
}

/// Return the input type accepted for a Pydantic root model field.
///
/// This builds a union of the root model instance and its transformed raw root type. For example,
/// a field annotated with an `IntList` derived from `RootModel[list[int]]` accepts both an
/// `IntList` instance and an `Iterable[LaxInt]`.
fn root_model_input_type<'db>(
    db: &'db dyn Db,
    field_type: Type<'db>,
    expanding_types: &mut FxHashSet<Type<'db>>,
) -> Option<Type<'db>> {
    let (class, specialization) = field_type.nominal_class(db)?.static_class_literal(db)?;
    if !is_root_model(db, class) {
        return None;
    }

    let Some(field_policy @ CodeGeneratorKind::Pydantic(_)) =
        CodeGeneratorKind::from_class(db, class.into())
    else {
        return Some(Type::any());
    };
    let Some(root_field) = class.fields(db, specialization, field_policy).get("root") else {
        return Some(Type::any());
    };

    if !expanding_types.insert(field_type) {
        return Some(Type::any());
    }
    let root_input_type = lax_input_type_impl(db, root_field.declared_ty, expanding_types);

    expanding_types.remove(&field_type);
    Some(UnionType::from_two_elements(
        db,
        field_type,
        root_input_type,
    ))
}

/// Return the known module, name, and class literal for an instance's nominal class.
fn instance_symbol<'db>(
    db: &'db dyn Db,
    ty: Type<'db>,
) -> Option<(KnownModule, &'db str, StaticClassLiteral<'db>)> {
    let class = ty.nominal_class(db)?.class_literal(db).as_static()?;
    let module = file_to_module(db, class.file(db))?.known(db)?;
    Some((module, class.name(db).as_str(), class))
}

/// Return a lax-input alias like `LaxInt` from `ty_extensions.pydantic`.
fn lax_alias<'db>(db: &'db dyn Db, name: &str) -> Type<'db> {
    match known_module_symbol(db, KnownModule::TyExtensionsPydantic, name)
        .place
        .ignore_possibly_undefined()
    {
        Some(Type::KnownInstance(KnownInstanceType::TypeAliasType(alias))) => {
            Type::TypeAlias(alias)
        }
        _ => Type::any(),
    }
}

/// Return `true` if `class` should accept extra keywords in its synthesized constructor.
pub(in crate::types) fn model_init_accepts_extra(
    db: &dyn Db,
    class: StaticClassLiteral<'_>,
    metadata: ModelMetadata<'_>,
) -> bool {
    if !metadata.accepts_extra(db) {
        return false;
    }

    for base in class
        .iter_mro(db, None)
        .skip(1)
        .filter_map(ClassBase::into_class)
        .filter_map(|base| base.static_class_literal(db))
        .map(|(base, _)| base)
    {
        if base.is_known(db, KnownClass::PydanticBaseModel) {
            return true;
        }

        // These constructors use variadic keywords for specialized inputs, not arbitrary extras.
        if base.is_known(db, KnownClass::PydanticRootModel)
            || base.is_known(db, KnownClass::PydanticBaseSettings)
        {
            return false;
        }

        let init = class_member(db, base.body_scope(db), "__init__");
        if !init.is_undefined() {
            return init
                .ignore_possibly_undefined()
                .and_then(Type::as_function_literal)
                .is_some_and(|init| {
                    init.signature(db)
                        .iter()
                        .any(|signature| signature.parameters().keyword_variadic().is_some())
                });
        }
    }

    false
}

/// Create the catch-all keyword parameter for a Pydantic model constructor.
///
/// Start with `extra` and append underscores until the name does not collide with a model field.
pub(in crate::types) fn extra_parameter<'db>(parameters: &[Parameter<'db>]) -> Parameter<'db> {
    let mut name = String::from("extra");

    while parameters
        .iter()
        .filter_map(Parameter::name)
        .any(|parameter_name| parameter_name.as_str() == name)
    {
        name.push('_');
    }

    Parameter::keyword_variadic(Name::new(name)).with_annotated_type(Type::any())
}
