use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::ops::{Deref, DerefMut};

use bitflags::bitflags;
use ordermap::OrderSet;
use ruff_db::diagnostic::{Annotation, Diagnostic, Span, SubDiagnostic, SubDiagnosticSeverity};
use ruff_db::parsed::parsed_module;
use ruff_python_ast::Arguments;
use ruff_python_ast::{self as ast, AnyNodeRef, StmtClassDef, name::Name};
use ruff_text_size::Ranged;

use super::class::{ClassType, CodeGeneratorKind, Field};
use super::context::InferContext;
use super::diagnostic::{
    self, INVALID_ARGUMENT_TYPE, INVALID_ASSIGNMENT, report_invalid_key_on_typed_dict,
    report_missing_typed_dict_key,
};
use super::{ApplyTypeMappingVisitor, IntersectionBuilder, Type, TypeMapping, visitor};
use crate::Db;
use crate::semantic_index::definition::Definition;
use crate::types::TypeDefinition;
use crate::types::class::FieldKind;
use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension};
use crate::types::generics::InferableTypeVars;
use crate::types::relation::{
    HasRelationToVisitor, IsDisjointVisitor, IsEquivalentVisitor, TypeRelation,
};
use crate::types::{NormalizedVisitor, TypeContext};

bitflags! {
    /// Used for `TypedDict` class parameters.
    /// Keeps track of the keyword arguments that were passed-in during class definition.
    /// (see https://typing.python.org/en/latest/spec/typeddict.html)
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct TypedDictParams: u8 {
        /// Whether keys are required by default (`total=True`)
        const TOTAL = 1 << 0;
    }
}

impl get_size2::GetSize for TypedDictParams {}

impl Default for TypedDictParams {
    fn default() -> Self {
        Self::TOTAL
    }
}

/// Type that represents the set of all inhabitants (`dict` instances) that conform to
/// a given `TypedDict` schema.
///
/// # Ordering
/// Ordering is derived from the variant order (`Class` < `Synthesized`) and the inner types.
/// The Salsa IDs of inner types may change between runs or when the type was garbage collected
/// and recreated.
#[derive(
    Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, salsa::Update, Hash, get_size2::GetSize,
)]
pub enum TypedDictType<'db> {
    /// A reference to the class (inheriting from `typing.TypedDict`) that specifies the
    /// schema of this `TypedDict`.
    Class(ClassType<'db>),
    /// A `TypedDict` that doesn't correspond to a class definition, either because it's been
    /// `normalized`, or because it's been synthesized to represent constraints.
    Synthesized(SynthesizedTypedDictType<'db>),
}

impl<'db> TypedDictType<'db> {
    pub fn new(defining_class: ClassType<'db>) -> Self {
        Self::Class(defining_class)
    }

    pub fn defining_class(self) -> Option<ClassType<'db>> {
        match self {
            Self::Class(defining_class) => Some(defining_class),
            Self::Synthesized(_) => None,
        }
    }

    pub fn items(self, db: &'db dyn Db) -> &'db TypedDictSchema<'db> {
        #[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)]
        fn class_based_items<'db>(db: &'db dyn Db, class: ClassType<'db>) -> TypedDictSchema<'db> {
            let Some((class_literal, specialization)) = class.static_class_literal(db) else {
                return TypedDictSchema::default();
            };
            class_literal
                .fields(db, specialization, CodeGeneratorKind::TypedDict)
                .into_iter()
                .map(|(name, field)| {
                    let field = match field {
                        Field {
                            first_declaration,
                            declared_ty,
                            kind:
                                FieldKind::TypedDict {
                                    is_required,
                                    is_read_only,
                                },
                        } => TypedDictFieldBuilder::new(*declared_ty)
                            .required(*is_required)
                            .read_only(*is_read_only)
                            .first_declaration(*first_declaration)
                            .build(),
                        _ => unreachable!("TypedDict field expected"),
                    };
                    (name.clone(), field)
                })
                .collect()
        }

        match self {
            Self::Class(defining_class) => class_based_items(db, defining_class),
            Self::Synthesized(synthesized) => synthesized.items(db),
        }
    }

    pub fn apply_type_mapping_impl<'a>(
        self,
        db: &'db dyn Db,
        type_mapping: &TypeMapping<'a, 'db>,
        tcx: TypeContext<'db>,
        visitor: &ApplyTypeMappingVisitor<'db>,
    ) -> Self {
        // TODO: Materialization of gradual TypedDicts needs more logic
        match self {
            Self::Class(defining_class) => {
                Self::Class(defining_class.apply_type_mapping_impl(db, type_mapping, tcx, visitor))
            }
            Self::Synthesized(synthesized) => Self::Synthesized(
                synthesized.apply_type_mapping_impl(db, type_mapping, tcx, visitor),
            ),
        }
    }

    // Subtyping between `TypedDict`s follows the algorithm described at:
    // https://typing.python.org/en/latest/spec/typeddict.html#subtyping-between-typeddict-types
    pub fn has_relation_to_impl(
        self,
        db: &'db dyn Db,
        target: TypedDictType<'db>,
        inferable: InferableTypeVars<'_, 'db>,
        relation: TypeRelation<'db>,
        relation_visitor: &HasRelationToVisitor<'db>,
        disjointness_visitor: &IsDisjointVisitor<'db>,
    ) -> ConstraintSet<'db> {
        // First do a quick nominal check that (if it succeeds) means that we can avoid
        // materializing the full `TypedDict` schema for either `self` or `target`.
        // This should be cheaper in many cases, and also helps us avoid some cycles.
        if let Some(defining_class) = self.defining_class()
            && let Some(target_defining_class) = target.defining_class()
            && defining_class.is_subclass_of(db, target_defining_class)
        {
            return ConstraintSet::from(true);
        }

        let self_items = self.items(db);
        let target_items = target.items(db);
        // Many rules violations short-circuit with "never", but asking whether one field is
        // [relation] to/of another can produce more complicated constraints, and we collect those.
        let mut constraints = ConstraintSet::from(true);
        for (target_item_name, target_item_field) in target_items {
            let field_constraints = if target_item_field.is_required() {
                // required target fields
                let Some(self_item_field) = self_items.get(target_item_name) else {
                    // Self is missing a required field.
                    return ConstraintSet::from(false);
                };
                if !self_item_field.is_required() {
                    // A required field is not required in self.
                    return ConstraintSet::from(false);
                }
                if target_item_field.is_read_only() {
                    // For `ReadOnly[]` fields in the target, the corresponding fields in
                    // self need to have the same assignability/subtyping/etc relation
                    // individually that we're looking for overall between the
                    // `TypedDict`s.
                    self_item_field.declared_ty.has_relation_to_impl(
                        db,
                        target_item_field.declared_ty,
                        inferable,
                        relation,
                        relation_visitor,
                        disjointness_visitor,
                    )
                } else {
                    if self_item_field.is_read_only() {
                        // A read-only field can't be assigned to a mutable target.
                        return ConstraintSet::from(false);
                    }
                    // For mutable fields in the target, the relation needs to apply both
                    // ways, or else mutating the target could violate the structural
                    // invariants of self. For fully-static types, this is "equivalence".
                    // For gradual types, it depends on the relation, but mutual
                    // assignability is "consistency".
                    self_item_field
                        .declared_ty
                        .has_relation_to_impl(
                            db,
                            target_item_field.declared_ty,
                            inferable,
                            relation,
                            relation_visitor,
                            disjointness_visitor,
                        )
                        .and(db, || {
                            target_item_field.declared_ty.has_relation_to_impl(
                                db,
                                self_item_field.declared_ty,
                                inferable,
                                relation,
                                relation_visitor,
                                disjointness_visitor,
                            )
                        })
                }
            } else {
                // `NotRequired[]` target fields
                if target_item_field.is_read_only() {
                    // As above, for `NotRequired[]` + `ReadOnly[]` fields in the target. It's
                    // tempting to refactor things and unify some of these calls to
                    // `has_relation_to_impl`, but this branch will get more complicated when we
                    // add support for `closed` and `extra_items` (which is why the rules in the
                    // spec are structured like they are), and following the structure of the spec
                    // makes it easier to check the logic here.
                    if let Some(self_item_field) = self_items.get(target_item_name) {
                        self_item_field.declared_ty.has_relation_to_impl(
                            db,
                            target_item_field.declared_ty,
                            inferable,
                            relation,
                            relation_visitor,
                            disjointness_visitor,
                        )
                    } else {
                        // Self is missing this not-required, read-only item. However, since all
                        // `TypedDict`s by default are allowed to have "extra items" of any type
                        // (until we support `closed` and explicit `extra_items`), this key could
                        // actually turn out to have a value. To make sure this is type-safe, the
                        // not-required field in the target needs to be assignable from `object`.
                        // TODO: `closed` and `extra_items` support will go here.
                        Type::object().when_assignable_to(
                            db,
                            target_item_field.declared_ty,
                            inferable,
                        )
                    }
                } else {
                    // As above, for `NotRequired[]` mutable fields in the target. Again the logic
                    // is largely the same for now, but it will get more complicated with `closed`
                    // and `extra_items`.
                    if let Some(self_item_field) = self_items.get(target_item_name) {
                        if self_item_field.is_read_only() {
                            // A read-only field can't be assigned to a mutable target.
                            return ConstraintSet::from(false);
                        }
                        if self_item_field.is_required() {
                            // A required field can't be assigned to a not-required, mutable field
                            // in the target, because `del` is allowed on the target field.
                            return ConstraintSet::from(false);
                        }

                        // As above, for mutable fields in the target, the relation needs
                        // to apply both ways.
                        self_item_field
                            .declared_ty
                            .has_relation_to_impl(
                                db,
                                target_item_field.declared_ty,
                                inferable,
                                relation,
                                relation_visitor,
                                disjointness_visitor,
                            )
                            .and(db, || {
                                target_item_field.declared_ty.has_relation_to_impl(
                                    db,
                                    self_item_field.declared_ty,
                                    inferable,
                                    relation,
                                    relation_visitor,
                                    disjointness_visitor,
                                )
                            })
                    } else {
                        // Self is missing this not-required, mutable field. This isn't ok if self
                        // has read-only extra items, which all `TypedDict`s effectively do until
                        // we support `closed` and explicit `extra_items`. See "A subtle
                        // interaction between two structural assignability rules prevents
                        // unsoundness" in `typed_dict.md`.
                        // TODO: `closed` and `extra_items` support will go here.
                        ConstraintSet::from(false)
                    }
                }
            };
            constraints.intersect(db, field_constraints);
            if constraints.is_never_satisfied(db) {
                return constraints;
            }
        }
        constraints
    }

    pub fn definition(self, db: &'db dyn Db) -> Option<Definition<'db>> {
        match self {
            TypedDictType::Class(defining_class) => defining_class.definition(db),
            TypedDictType::Synthesized(_) => None,
        }
    }

    pub fn type_definition(self, db: &'db dyn Db) -> Option<TypeDefinition<'db>> {
        match self {
            TypedDictType::Class(defining_class) => defining_class.type_definition(db),
            TypedDictType::Synthesized(_) => None,
        }
    }

    pub fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
        match self {
            TypedDictType::Class(_) => {
                let synthesized = SynthesizedTypedDictType::new(db, self.items(db));
                TypedDictType::Synthesized(synthesized.normalized_impl(db, visitor))
            }
            TypedDictType::Synthesized(synthesized) => {
                TypedDictType::Synthesized(synthesized.normalized_impl(db, visitor))
            }
        }
    }

    pub fn is_equivalent_to_impl(
        self,
        db: &'db dyn Db,
        other: TypedDictType<'db>,
        inferable: InferableTypeVars<'_, 'db>,
        visitor: &IsEquivalentVisitor<'db>,
    ) -> ConstraintSet<'db> {
        // TODO: `closed` and `extra_items` support will go here. Until then we don't look at the
        // params at all, because `total` is already incorporated into `FieldKind`.

        // Since both sides' fields are pre-sorted into `BTreeMap`s, we can iterate over them in
        // sorted order instead of paying for a lookup for each field, as long as their lengths are
        // the same.
        if self.items(db).len() != other.items(db).len() {
            return ConstraintSet::from(false);
        }
        self.items(db).iter().zip(other.items(db)).when_all(
            db,
            |((name, field), (other_name, other_field))| {
                if name != other_name || field.flags != other_field.flags {
                    return ConstraintSet::from(false);
                }
                field.declared_ty.is_equivalent_to_impl(
                    db,
                    other_field.declared_ty,
                    inferable,
                    visitor,
                )
            },
        )
    }

    /// Two `TypedDict`s `A` and `B` are disjoint if it's impossible to come up with a third
    /// `TypedDict` `C` that's fully-static and assignable to both of them.
    ///
    /// `TypedDict` assignability is determined field-by-field, so we determine disjointness
    /// similarly. For any field that's only in `A`, it's always possible for our hypothetical `C`
    /// to copy/paste that field without losing assignability to `B` (and vice versa), so we only
    /// need to consider fields that are present in both `A` and `B`.
    ///
    /// There are three properties of each field to consider: the declared type, whether it's
    /// mutable ("mut" vs "imm" below), and whether it's required ("req" vs "opt" below). Here's a
    /// table summary of the restrictions on the declared type of a source field (for us that means
    /// in `C`, which we want to be assignable to both `A` and `B`) given a destination field (for
    /// us that means in either `A` or `B`). For completeness we'll also include the possibility
    /// that the source field is missing entirely, though we'll soon see that we can ignore that
    /// case. This table is essentially what `has_relation_to_impl` implements above. Here
    /// "equivalent" means the source and destination types must be equivalent/compatible,
    /// "assignable" means the source must be assignable to the destination, and "-" means the
    /// assignment is never allowed:
    ///
    /// | dest ↓ source →  | mut + req  | mut + opt  | imm + req  | imm + opt  |   \[missing]  |
    /// |------------------|------------|------------|------------|------------|---------------|
    /// |    mut + req     | equivalent |     -      |     -      |     -      |       -       |
    /// |    mut + opt     |     -      | equivalent |     -      |     -      |       -       |
    /// |    imm + req     | assignable |     -      | assignable |     -      |       -       |
    /// |    imm + opt     | assignable | assignable | assignable | assignable | \[dest is obj]|
    ///
    /// We can cut that table down substantially by noticing two things:
    ///
    /// - We don't need to consider the cases where the source field (in `C`) is `ReadOnly`/"imm",
    ///   because the mutable version of the same field is always "strictly more assignable". In
    ///   other words, nothing in the `TypedDict` assignability rules ever requires a source field
    ///   to be immutable.
    /// - We don't need to consider the special case where the source field is missing, because
    ///   that's only allowed when the destination is `ReadOnly[NotRequired[object]]`, which is
    ///   compatible with *any* choice of source field.
    ///
    /// The cases we actually need to reason about are this smaller table:
    ///
    /// | dest ↓ source →  | mut + req  | mut + opt  |
    /// |------------------|------------|------------|
    /// |    mut + req     | equivalent |     -      |
    /// |    mut + opt     |     -      | equivalent |
    /// |    imm + req     | assignable |     -      |
    /// |    imm + opt     | assignable | assignable |
    ///
    /// So, given a field name that's in both `A` and `B`, here are the conditions where it's
    /// *impossible* to choose a source field for `C` that's compatible with both destinations,
    /// which tells us that `A` and `B` are disjoint:
    ///
    /// 1. If one side is "mut+opt" (which forces the field in `C` to be "opt") and the other side
    ///    is "req" (which forces the field in `C` to be "req").
    /// 2. If both sides are mutable, and their types are not equivalent/compatible. (Because the
    ///    type in `C` must be compatible with both of them.)
    /// 3. If one sides is mutable, and its type is not assignable to the immutable side's type.
    ///    (Because the type in `C` must be compatible with the mutable side.)
    /// 4. If both sides are immutable, and their types are disjoint. (Because the type in `C` must
    ///    be assignable to both.)
    ///
    /// TODO: Adding support for `closed` and `extra_items` will complicate this.
    pub fn is_disjoint_from_impl(
        self,
        db: &'db dyn Db,
        other: TypedDictType<'db>,
        inferable: InferableTypeVars<'_, 'db>,
        disjointness_visitor: &IsDisjointVisitor<'db>,
        relation_visitor: &HasRelationToVisitor<'db>,
    ) -> ConstraintSet<'db> {
        let fields_in_common = btreemap_values_with_same_key(self.items(db), other.items(db));
        fields_in_common.when_any(db, |(self_field, other_field)| {
            // Condition 1 above.
            if self_field.is_required() || other_field.is_required() {
                if (!self_field.is_required() && !self_field.is_read_only())
                    || (!other_field.is_required() && !other_field.is_read_only())
                {
                    // One side demands a `Required` source field, while the other side demands a
                    // `NotRequired` one. They must be disjoint.
                    return ConstraintSet::from(true);
                }
            }
            if !self_field.is_read_only() && !other_field.is_read_only() {
                // Condition 2 above. This field is mutable on both sides, so the so the types must
                // be compatible, i.e. mutually assignable.
                self_field
                    .declared_ty
                    .has_relation_to_impl(
                        db,
                        other_field.declared_ty,
                        inferable,
                        TypeRelation::Assignability,
                        relation_visitor,
                        disjointness_visitor,
                    )
                    .and(db, || {
                        other_field.declared_ty.has_relation_to_impl(
                            db,
                            self_field.declared_ty,
                            inferable,
                            TypeRelation::Assignability,
                            relation_visitor,
                            disjointness_visitor,
                        )
                    })
                    .negate(db)
            } else if !self_field.is_read_only() {
                // Half of condition 3 above.
                self_field
                    .declared_ty
                    .has_relation_to_impl(
                        db,
                        other_field.declared_ty,
                        inferable,
                        TypeRelation::Assignability,
                        relation_visitor,
                        disjointness_visitor,
                    )
                    .negate(db)
            } else if !other_field.is_read_only() {
                // The other half of condition 3 above.
                other_field
                    .declared_ty
                    .has_relation_to_impl(
                        db,
                        self_field.declared_ty,
                        inferable,
                        TypeRelation::Assignability,
                        relation_visitor,
                        disjointness_visitor,
                    )
                    .negate(db)
            } else {
                // Condition 4 above.
                self_field.declared_ty.is_disjoint_from_impl(
                    db,
                    other_field.declared_ty,
                    inferable,
                    disjointness_visitor,
                    relation_visitor,
                )
            }
        })
    }
}

pub fn walk_typed_dict_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
    db: &'db dyn Db,
    typed_dict: TypedDictType<'db>,
    visitor: &V,
) {
    match typed_dict {
        TypedDictType::Class(defining_class) => {
            visitor.visit_type(db, defining_class.into());
        }
        TypedDictType::Synthesized(synthesized) => {
            for field in synthesized.items(db).values() {
                visitor.visit_type(db, field.declared_ty);
            }
        }
    }
}

pub fn typed_dict_params_from_class_def(class_stmt: &StmtClassDef) -> TypedDictParams {
    let mut typed_dict_params = TypedDictParams::default();

    // Check for `total` keyword argument in the class definition
    // Note that it is fine to only check for Boolean literals here
    // (https://typing.python.org/en/latest/spec/typeddict.html#totality)
    if let Some(arguments) = &class_stmt.arguments {
        for keyword in &arguments.keywords {
            if keyword.arg.as_deref() == Some("total")
                && matches!(
                    &keyword.value,
                    ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value: false, .. })
                )
            {
                typed_dict_params.remove(TypedDictParams::TOTAL);
            }
        }
    }

    typed_dict_params
}

#[derive(Debug, Clone, Copy)]
pub enum TypedDictAssignmentKind {
    /// For subscript assignments like `d["key"] = value`
    Subscript,
    /// For constructor arguments like `MyTypedDict(key=value)`
    Constructor,
}

impl TypedDictAssignmentKind {
    fn diagnostic_name(self) -> &'static str {
        match self {
            Self::Subscript => "assignment",
            Self::Constructor => "argument",
        }
    }

    fn diagnostic_type(self) -> &'static crate::lint::LintMetadata {
        match self {
            Self::Subscript => &INVALID_ASSIGNMENT,
            Self::Constructor => &INVALID_ARGUMENT_TYPE,
        }
    }

    const fn is_subscript(self) -> bool {
        matches!(self, Self::Subscript)
    }
}

/// A helper that validates assignments of a value to a specific key on a `TypedDict`.
pub struct TypedDictKeyAssignment<'a, 'db, 'ast> {
    pub context: &'a InferContext<'db, 'ast>,
    pub typed_dict: TypedDictType<'db>,
    pub full_object_ty: Option<Type<'db>>,
    pub key: &'a str,
    pub value_ty: Type<'db>,
    pub typed_dict_node: AnyNodeRef<'ast>,
    pub key_node: AnyNodeRef<'ast>,
    pub value_node: AnyNodeRef<'ast>,
    pub assignment_kind: TypedDictAssignmentKind,
    pub emit_diagnostic: bool,
}

impl<'db> TypedDictKeyAssignment<'_, 'db, '_> {
    pub fn validate(&self) -> bool {
        let db = self.context.db();
        let items = self.typed_dict.items(db);

        // Check if key exists in `TypedDict`
        let Some((_, item)) = items.iter().find(|(name, _)| *name == self.key) else {
            if self.emit_diagnostic {
                report_invalid_key_on_typed_dict(
                    self.context,
                    self.typed_dict_node,
                    self.key_node,
                    Type::TypedDict(self.typed_dict),
                    self.full_object_ty,
                    Type::string_literal(db, self.key),
                    items,
                );
            }

            return false;
        };

        if self.assignment_kind.is_subscript() && item.is_read_only() {
            if self.emit_diagnostic
                && let Some(builder) = self
                    .context
                    .report_lint(self.assignment_kind.diagnostic_type(), self.key_node)
            {
                let typed_dict_ty = Type::TypedDict(self.typed_dict);
                let typed_dict_d = typed_dict_ty.display(db);

                let mut diagnostic = builder.into_diagnostic(format_args!(
                    "Cannot assign to key \"{}\" on TypedDict `{typed_dict_d}`",
                    self.key,
                ));

                diagnostic.set_primary_message(format_args!("key is marked read-only"));
                self.add_object_type_annotation(db, &mut diagnostic);
                Self::add_item_definition_subdiagnostic(
                    db,
                    item,
                    &mut diagnostic,
                    "Read-only item declared here",
                );
            }

            return false;
        }

        // Key exists, check if value type is assignable to declared type
        if self.value_ty.is_assignable_to(db, item.declared_ty) {
            return true;
        }

        if diagnostic::is_invalid_typed_dict_literal(db, item.declared_ty, self.value_node) {
            return false;
        }

        // Invalid assignment - emit diagnostic
        if self.emit_diagnostic
            && let Some(builder) = self
                .context
                .report_lint(self.assignment_kind.diagnostic_type(), self.value_node)
        {
            let typed_dict_ty = Type::TypedDict(self.typed_dict);
            let typed_dict_d = typed_dict_ty.display(db);
            let value_d = self.value_ty.display(db);
            let item_type_d = item.declared_ty.display(db);

            let mut diagnostic = builder.into_diagnostic(format_args!(
                "Invalid {} to key \"{}\" with declared type `{item_type_d}` \
                on TypedDict `{typed_dict_d}`",
                self.assignment_kind.diagnostic_name(),
                self.key,
            ));

            diagnostic.set_primary_message(format_args!("value of type `{value_d}`"));

            diagnostic.annotate(
                self.context
                    .secondary(self.key_node)
                    .message(format_args!("key has declared type `{item_type_d}`")),
            );

            Self::add_item_definition_subdiagnostic(
                db,
                item,
                &mut diagnostic,
                "Item declared here",
            );
            self.add_object_type_annotation(db, &mut diagnostic);
        }

        false
    }

    fn add_object_type_annotation(&self, db: &'db dyn Db, diagnostic: &mut Diagnostic) {
        if let Some(full_object_ty) = self.full_object_ty {
            diagnostic.annotate(self.context.secondary(self.typed_dict_node).message(
                format_args!(
                    "TypedDict `{}` in {kind} type `{}`",
                    Type::TypedDict(self.typed_dict).display(db),
                    full_object_ty.display(db),
                    kind = if full_object_ty.is_union() {
                        "union"
                    } else {
                        "intersection"
                    },
                ),
            ));
        } else {
            diagnostic.annotate(self.context.secondary(self.typed_dict_node).message(
                format_args!(
                    "TypedDict `{}`",
                    Type::TypedDict(self.typed_dict).display(db)
                ),
            ));
        }
    }

    fn add_item_definition_subdiagnostic(
        db: &'db dyn Db,
        item: &TypedDictField<'db>,
        diagnostic: &mut Diagnostic,
        message: &str,
    ) {
        if let Some(declaration) = item.first_declaration() {
            let file = declaration.file(db);
            let module = parsed_module(db, file).load(db);

            let mut sub = SubDiagnostic::new(SubDiagnosticSeverity::Info, "Item declaration");
            sub.annotate(
                Annotation::secondary(
                    Span::from(file).with_range(declaration.full_range(db, &module).range()),
                )
                .message(message),
            );
            diagnostic.sub(sub);
        }
    }
}

/// Validates that all required keys are provided in a `TypedDict` construction.
///
/// Reports errors for any keys that are required but not provided.
///
/// Returns true if the assignment is valid, or false otherwise.
pub fn validate_typed_dict_required_keys<'db, 'ast>(
    context: &InferContext<'db, 'ast>,
    typed_dict: TypedDictType<'db>,
    provided_keys: &OrderSet<Name>,
    error_node: AnyNodeRef<'ast>,
) -> bool {
    let db = context.db();
    let items = typed_dict.items(db);

    let required_keys: OrderSet<Name> = items
        .iter()
        .filter_map(|(key_name, field)| field.is_required().then_some(key_name.clone()))
        .collect();

    let missing_keys = required_keys.difference(provided_keys);

    let mut has_missing_key = false;
    for missing_key in missing_keys {
        has_missing_key = true;

        report_missing_typed_dict_key(
            context,
            error_node,
            Type::TypedDict(typed_dict),
            missing_key.as_str(),
        );
    }

    !has_missing_key
}

/// Extracts `TypedDict` keys and their types from a type, resolving type aliases and handling
/// intersections.
///
/// For intersections, returns ALL keys from ALL `TypedDict` types (union of keys), because a
/// value of an intersection type must satisfy all `TypedDict`s and therefore has all their keys.
/// For keys that appear in multiple `TypedDict`s, the types are intersected.
fn extract_typed_dict_keys<'db>(
    db: &'db dyn Db,
    ty: Type<'db>,
) -> Option<BTreeMap<Name, Type<'db>>> {
    match ty {
        Type::TypedDict(td) => {
            let keys = td
                .items(db)
                .iter()
                .map(|(name, field)| (name.clone(), field.declared_ty))
                .collect();
            Some(keys)
        }
        Type::Intersection(intersection) => {
            // Collect key maps from all TypedDicts in the intersection
            let all_key_maps: Vec<_> = intersection
                .positive(db)
                .iter()
                .filter_map(|element| extract_typed_dict_keys(db, *element))
                .collect();

            if all_key_maps.is_empty() {
                return None;
            }

            // Union all keys from all TypedDicts, intersecting types for shared keys
            let mut result: BTreeMap<Name, Type<'db>> = BTreeMap::new();

            for key_map in all_key_maps {
                for (key, ty) in key_map {
                    result
                        .entry(key)
                        .and_modify(|existing_ty| {
                            // Key exists in multiple TypedDicts - intersect the types
                            *existing_ty = IntersectionBuilder::new(db)
                                .add_positive(*existing_ty)
                                .add_positive(ty)
                                .build();
                        })
                        .or_insert(ty);
                }
            }

            Some(result)
        }
        // TODO: handle unions by checking all TypedDict elements separately
        Type::Union(_) => None,
        Type::TypeAlias(alias) => extract_typed_dict_keys(db, alias.value_type(db)),
        // All other types cannot contain a TypedDict
        Type::Dynamic(_)
        | Type::Never
        | Type::FunctionLiteral(_)
        | Type::BoundMethod(_)
        | Type::KnownBoundMethod(_)
        | Type::WrapperDescriptor(_)
        | Type::DataclassDecorator(_)
        | Type::DataclassTransformer(_)
        | Type::Callable(_)
        | Type::ModuleLiteral(_)
        | Type::ClassLiteral(_)
        | Type::GenericAlias(_)
        | Type::SubclassOf(_)
        | Type::NominalInstance(_)
        | Type::ProtocolInstance(_)
        | Type::SpecialForm(_)
        | Type::KnownInstance(_)
        | Type::PropertyInstance(_)
        | Type::AlwaysTruthy
        | Type::AlwaysFalsy
        | Type::LiteralValue(_)
        | Type::TypeVar(_)
        | Type::BoundSuper(_)
        | Type::TypeIs(_)
        | Type::TypeGuard(_)
        | Type::NewTypeInstance(_) => None,
    }
}

pub fn validate_typed_dict_constructor<'db, 'ast>(
    context: &InferContext<'db, 'ast>,
    typed_dict: TypedDictType<'db>,
    arguments: &'ast Arguments,
    error_node: AnyNodeRef<'ast>,
    expression_type_fn: impl Fn(&ast::Expr) -> Type<'db>,
) {
    let db = context.db();

    // Check for a single positional argument that is a dict literal
    let has_positional_dict_literal = arguments.args.len() == 1 && arguments.args[0].is_dict_expr();

    // Check for a single positional argument (not a dict literal)
    let is_single_positional_arg =
        arguments.args.len() == 1 && arguments.keywords.is_empty() && !has_positional_dict_literal;

    if has_positional_dict_literal {
        let provided_keys = validate_from_dict_literal(
            context,
            typed_dict,
            arguments,
            error_node,
            &expression_type_fn,
        );
        validate_typed_dict_required_keys(context, typed_dict, &provided_keys, error_node);
    } else if is_single_positional_arg {
        // Single positional argument: check if assignable to the target TypedDict.
        // This handles TypedDict, intersections, unions, and type aliases correctly.
        // Assignability already checks for required keys and type compatibility,
        // so we don't need separate validation.
        let arg = &arguments.args[0];
        let arg_ty = expression_type_fn(arg);
        let target_ty = Type::TypedDict(typed_dict);

        if !arg_ty.is_assignable_to(db, target_ty) {
            if let Some(builder) = context.report_lint(&INVALID_ARGUMENT_TYPE, arg) {
                builder.into_diagnostic(format_args!(
                    "Argument of type `{}` is not assignable to `{}`",
                    arg_ty.display(db),
                    target_ty.display(db),
                ));
            }
        }
    } else {
        let provided_keys = validate_from_keywords(
            context,
            typed_dict,
            arguments,
            error_node,
            &expression_type_fn,
        );
        validate_typed_dict_required_keys(context, typed_dict, &provided_keys, error_node);
    }
}

/// Validates a `TypedDict` constructor call with a single positional dictionary argument
/// e.g. `Person({"name": "Alice", "age": 30})`
fn validate_from_dict_literal<'db, 'ast>(
    context: &InferContext<'db, 'ast>,
    typed_dict: TypedDictType<'db>,
    arguments: &'ast Arguments,
    typed_dict_node: AnyNodeRef<'ast>,
    expression_type_fn: &impl Fn(&ast::Expr) -> Type<'db>,
) -> OrderSet<Name> {
    let mut provided_keys = OrderSet::new();

    if let ast::Expr::Dict(dict_expr) = &arguments.args[0] {
        // Validate dict entries
        for dict_item in &dict_expr.items {
            if let Some(ref key_expr) = dict_item.key
                && let ast::Expr::StringLiteral(ast::ExprStringLiteral {
                    value: key_value, ..
                }) = key_expr
            {
                let key = key_value.to_str();
                provided_keys.insert(Name::new(key));

                // Get the already-inferred argument type
                let value_ty = expression_type_fn(&dict_item.value);
                TypedDictKeyAssignment {
                    context,
                    typed_dict,
                    full_object_ty: None,
                    key,
                    value_ty,
                    typed_dict_node,
                    key_node: key_expr.into(),
                    value_node: (&dict_item.value).into(),
                    assignment_kind: TypedDictAssignmentKind::Constructor,
                    emit_diagnostic: true,
                }
                .validate();
            }
        }
    }

    provided_keys
}

/// Validates a `TypedDict` constructor call with keywords
/// e.g. `Person(name="Alice", age=30)` or `Person(**other_typed_dict)`
fn validate_from_keywords<'db, 'ast>(
    context: &InferContext<'db, 'ast>,
    typed_dict: TypedDictType<'db>,
    arguments: &'ast Arguments,
    typed_dict_node: AnyNodeRef<'ast>,
    expression_type_fn: &impl Fn(&ast::Expr) -> Type<'db>,
) -> OrderSet<Name> {
    let db = context.db();

    // Collect keys from explicit keyword arguments
    let mut provided_keys: OrderSet<Name> = arguments
        .keywords
        .iter()
        .filter_map(|kw| kw.arg.as_ref().map(|arg| arg.id.clone()))
        .collect();

    // Validate that each key is assigned a type that is compatible with the key's value type
    for keyword in &arguments.keywords {
        if let Some(arg_name) = &keyword.arg {
            // Explicit keyword argument: e.g., `name="Alice"`
            let value_ty = expression_type_fn(&keyword.value);
            TypedDictKeyAssignment {
                context,
                typed_dict,
                full_object_ty: None,
                key: arg_name.as_str(),
                value_ty,
                typed_dict_node,
                key_node: keyword.into(),
                value_node: (&keyword.value).into(),
                assignment_kind: TypedDictAssignmentKind::Constructor,
                emit_diagnostic: true,
            }
            .validate();
        } else {
            // Keyword unpacking: e.g., `**other_typed_dict`
            // Unlike positional TypedDict arguments, unpacking passes all keys as explicit
            // keyword arguments, so extra keys should be flagged as errors (consistent with
            // explicitly providing those keys).
            let unpacked_type = expression_type_fn(&keyword.value);

            // Never and Dynamic types are special: they can have any keys, so we skip
            // validation and mark all required keys as provided.
            if unpacked_type.is_never() || unpacked_type.is_dynamic() {
                for (key_name, field) in typed_dict.items(db) {
                    if field.is_required() {
                        provided_keys.insert(key_name.clone());
                    }
                }
            } else if let Some(unpacked_keys) = extract_typed_dict_keys(db, unpacked_type) {
                for (key_name, value_ty) in &unpacked_keys {
                    provided_keys.insert(key_name.clone());
                    TypedDictKeyAssignment {
                        context,
                        typed_dict,
                        full_object_ty: None,
                        key: key_name.as_str(),
                        value_ty: *value_ty,
                        typed_dict_node,
                        key_node: keyword.into(),
                        value_node: (&keyword.value).into(),
                        assignment_kind: TypedDictAssignmentKind::Constructor,
                        emit_diagnostic: true,
                    }
                    .validate();
                }
            }
        }
    }

    provided_keys
}

/// Validates a `TypedDict` dictionary literal assignment,
/// e.g. `person: Person = {"name": "Alice", "age": 30}`
pub fn validate_typed_dict_dict_literal<'db>(
    context: &InferContext<'db, '_>,
    typed_dict: TypedDictType<'db>,
    dict_expr: &ast::ExprDict,
    typed_dict_node: AnyNodeRef,
    expression_type_fn: impl Fn(&ast::Expr) -> Type<'db>,
) -> Result<OrderSet<Name>, OrderSet<Name>> {
    let mut valid = true;
    let mut provided_keys = OrderSet::new();

    // Validate each key-value pair in the dictionary literal
    for item in &dict_expr.items {
        if let Some(key_expr) = &item.key
            && let Some(key_str) = expression_type_fn(key_expr).as_string_literal()
        {
            let key = key_str.value(context.db());
            provided_keys.insert(Name::new(key));

            let value_ty = expression_type_fn(&item.value);

            valid &= TypedDictKeyAssignment {
                context,
                typed_dict,
                full_object_ty: None,
                key,
                value_ty,
                typed_dict_node,
                key_node: key_expr.into(),
                value_node: (&item.value).into(),
                assignment_kind: TypedDictAssignmentKind::Constructor,
                emit_diagnostic: true,
            }
            .validate();
        }
    }

    valid &=
        validate_typed_dict_required_keys(context, typed_dict, &provided_keys, typed_dict_node);

    if valid {
        Ok(provided_keys)
    } else {
        Err(provided_keys)
    }
}

/// # Ordering
/// Ordering is based on the type's salsa-assigned id and not on its values.
/// The id may change between runs, or when the type was garbage collected and recreated.
#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
#[derive(PartialOrd, Ord)]
pub struct SynthesizedTypedDictType<'db> {
    #[returns(ref)]
    pub items: TypedDictSchema<'db>,
}

// The Salsa heap is tracked separately.
impl get_size2::GetSize for SynthesizedTypedDictType<'_> {}

impl<'db> SynthesizedTypedDictType<'db> {
    pub fn apply_type_mapping_impl<'a>(
        self,
        db: &'db dyn Db,
        type_mapping: &TypeMapping<'a, 'db>,
        tcx: TypeContext<'db>,
        visitor: &ApplyTypeMappingVisitor<'db>,
    ) -> Self {
        let items = self
            .items(db)
            .iter()
            .map(|(name, field)| {
                let field = field
                    .clone()
                    .apply_type_mapping_impl(db, type_mapping, tcx, visitor);

                (name.clone(), field)
            })
            .collect::<TypedDictSchema<'db>>();

        SynthesizedTypedDictType::new(db, items)
    }

    pub fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
        let items = self
            .items(db)
            .iter()
            .map(|(name, field)| {
                let field = field.clone().normalized_impl(db, visitor);
                (name.clone(), field)
            })
            .collect::<TypedDictSchema<'db>>();
        Self::new(db, items)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, get_size2::GetSize, salsa::Update)]
pub struct TypedDictSchema<'db>(BTreeMap<Name, TypedDictField<'db>>);

impl<'db> Deref for TypedDictSchema<'db> {
    type Target = BTreeMap<Name, TypedDictField<'db>>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for TypedDictSchema<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<'a> IntoIterator for &'a TypedDictSchema<'_> {
    type Item = (&'a Name, &'a TypedDictField<'a>);
    type IntoIter = std::collections::btree_map::Iter<'a, Name, TypedDictField<'a>>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<'db> FromIterator<(Name, TypedDictField<'db>)> for TypedDictSchema<'db> {
    fn from_iter<T: IntoIterator<Item = (Name, TypedDictField<'db>)>>(iter: T) -> Self {
        Self(iter.into_iter().collect())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, get_size2::GetSize, salsa::Update)]
pub struct TypedDictField<'db> {
    pub declared_ty: Type<'db>,
    flags: TypedDictFieldFlags,
    first_declaration: Option<Definition<'db>>,
}

impl<'db> TypedDictField<'db> {
    pub const fn is_required(&self) -> bool {
        self.flags.contains(TypedDictFieldFlags::REQUIRED)
    }

    pub const fn is_read_only(&self) -> bool {
        self.flags.contains(TypedDictFieldFlags::READ_ONLY)
    }

    pub const fn first_declaration(&self) -> Option<Definition<'db>> {
        self.first_declaration
    }

    pub fn apply_type_mapping_impl<'a>(
        self,
        db: &'db dyn Db,
        type_mapping: &TypeMapping<'a, 'db>,
        tcx: TypeContext<'db>,
        visitor: &ApplyTypeMappingVisitor<'db>,
    ) -> Self {
        Self {
            declared_ty: self
                .declared_ty
                .apply_type_mapping_impl(db, type_mapping, tcx, visitor),
            flags: self.flags,
            first_declaration: self.first_declaration,
        }
    }

    pub fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
        Self {
            declared_ty: self.declared_ty.normalized_impl(db, visitor),
            flags: self.flags,
            // A normalized typed-dict field does not hold onto the original declaration,
            // since a normalized typed-dict is an abstract type where equality does not depend
            // on the source-code definition.
            first_declaration: None,
        }
    }
}

pub struct TypedDictFieldBuilder<'db> {
    declared_ty: Type<'db>,
    flags: TypedDictFieldFlags,
    first_declaration: Option<Definition<'db>>,
}

impl<'db> TypedDictFieldBuilder<'db> {
    pub fn new(declared_ty: Type<'db>) -> Self {
        Self {
            declared_ty,
            flags: TypedDictFieldFlags::empty(),
            first_declaration: None,
        }
    }

    pub fn required(mut self, yes: bool) -> Self {
        self.flags.set(TypedDictFieldFlags::REQUIRED, yes);
        self
    }

    pub fn read_only(mut self, yes: bool) -> Self {
        self.flags.set(TypedDictFieldFlags::READ_ONLY, yes);
        self
    }

    pub fn first_declaration(mut self, definition: Option<Definition<'db>>) -> Self {
        self.first_declaration = definition;
        self
    }

    pub fn build(self) -> TypedDictField<'db> {
        TypedDictField {
            declared_ty: self.declared_ty,
            flags: self.flags,
            first_declaration: self.first_declaration,
        }
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, salsa::Update)]
    struct TypedDictFieldFlags: u8 {
        const REQUIRED = 1 << 0;
        const READ_ONLY = 1 << 1;
    }
}

impl get_size2::GetSize for TypedDictFieldFlags {}

/// Yield all the key/val pairs where the same key is present in both `BTreeMap`s. Take advantage
/// of the fact that keys are sorted to walk through each map once without doing any lookups. It
/// would be nice if `BTreeMap` had something like `BTreeSet::intersection` that did this for us,
/// but as far as I know we have to do it ourselves. Life is hard.
fn btreemap_values_with_same_key<'a, K, V1, V2>(
    left: &'a BTreeMap<K, V1>,
    right: &'a BTreeMap<K, V2>,
) -> impl Iterator<Item = (&'a V1, &'a V2)>
where
    K: Ord,
{
    let mut left_items = left.iter().peekable();
    let mut right_items = right.iter().peekable();
    std::iter::from_fn(move || {
        while let (Some((left_key, left_val)), Some((right_key, right_val))) =
            (left_items.peek().copied(), right_items.peek().copied())
        {
            match left_key.cmp(right_key) {
                Ordering::Equal => {
                    // Matching keys. Yield this pair of values and advance both iterators.
                    left_items.next();
                    right_items.next();
                    return Some((left_val, right_val));
                }
                Ordering::Less => {
                    // `left_items` is behind `right_items` in key order. Advance `left_items`.
                    left_items.next();
                }
                Ordering::Greater => {
                    // The opposite.
                    right_items.next();
                }
            }
        }
        // We've exhausted one or both of the maps, so there can be no more matching keys.
        None
    })
}

#[test]
fn test_btreemap_overlapping_items() {
    // A case with partial overlap and gaps.
    let left = BTreeMap::from_iter([("a", 1), ("b", 2), ("c", 3), ("d", 4), ("e", 5)]);
    let right = BTreeMap::from_iter([("b", 2.0), ("d", 4.0), ("f", 6.0)]);
    assert_eq!(
        btreemap_values_with_same_key(&left, &right).collect::<Vec<_>>(),
        vec![(&2, &2.0), (&4, &4.0)],
    );
    assert_eq!(
        btreemap_values_with_same_key(&right, &left).collect::<Vec<_>>(),
        vec![(&2.0, &2), (&4.0, &4)],
    );

    // A case where one side is empty.
    let left = BTreeMap::<i32, i32>::new();
    let right = BTreeMap::<i32, i32>::from_iter([(1, 1), (2, 2)]);
    assert!(
        btreemap_values_with_same_key(&left, &right)
            .next()
            .is_none()
    );
    assert!(
        btreemap_values_with_same_key(&right, &left)
            .next()
            .is_none()
    );
}
