use crate::config::{CheckerConfig, UnknownCallBehavior};
use crate::model::{
    AssignKind, Expression, ExpressionKind, ExternalPurityAssessment, ExternalPurityCategory,
    ExternalPurityReason, FunctionAssessment, FunctionDef, ImplementationStyleAssessment,
    ImplementationStyleCategory, Module, SourceSpan, Statement, StatementKind, StyleMetrics,
};
use crate::rules::{
    ambient_dependency_call_path, is_direct_effectful_call, is_direct_io_call,
    is_in_place_mutation_call, is_non_deterministic_call, referential_transparency_reference_path,
};
use std::collections::{BTreeMap, BTreeSet};

pub fn assess_function(
    module: &Module,
    function: &FunctionDef,
    config: &CheckerConfig,
) -> FunctionAssessment {
    let context = FunctionAssessmentContext::new(module, function, config);

    FunctionAssessment {
        name: function.name.clone(),
        qualified_name: function.qualified_name.clone(),
        class_name: function.class_name.clone(),
        location: function.location.clone(),
        external_purity: config
            .purity
            .enable
            .then(|| assess_external_purity(function, &context)),
        implementation_style: config
            .style
            .enable
            .then(|| assess_implementation_style(function, &context)),
    }
}

struct FunctionAssessmentContext<'a> {
    config: &'a CheckerConfig,
    mutable_globals: BTreeSet<String>,
    module_globals: BTreeSet<String>,
    parameter_names: BTreeSet<String>,
    local_bindings: BTreeSet<String>,
    declared_globals: BTreeSet<String>,
    declared_nonlocals: BTreeSet<String>,
    closure_bindings: BTreeSet<String>,
}

impl<'a> FunctionAssessmentContext<'a> {
    fn new(module: &'a Module, function: &'a FunctionDef, config: &'a CheckerConfig) -> Self {
        Self {
            config,
            mutable_globals: module
                .globals
                .iter()
                .filter(|binding| binding.is_mutable)
                .map(|binding| binding.name.clone())
                .collect(),
            module_globals: module
                .globals
                .iter()
                .map(|binding| binding.name.clone())
                .collect(),
            parameter_names: function
                .parameters
                .iter()
                .map(|parameter| parameter.name.clone())
                .collect(),
            local_bindings: function.local_bindings.iter().cloned().collect(),
            declared_globals: collect_declared_globals(function),
            declared_nonlocals: collect_declared_nonlocals(function),
            closure_bindings: collect_closure_bindings(module, function),
        }
    }

    fn is_local_only_binding(&self, name: &str) -> bool {
        self.local_bindings.contains(name)
            && !self.parameter_names.contains(name)
            && !self.declared_globals.contains(name)
            && !self.declared_nonlocals.contains(name)
    }

    fn is_external_read_base(&self, name: &str) -> bool {
        self.closure_bindings.contains(name)
            || (self.mutable_globals.contains(name)
                && (!self.local_bindings.contains(name) || self.declared_globals.contains(name)))
    }

    fn is_external_write_base(&self, name: &str) -> bool {
        self.parameter_names.contains(name)
            || self.declared_globals.contains(name)
            || self.declared_nonlocals.contains(name)
            || self.closure_bindings.contains(name)
            || matches!(name, "self" | "cls")
            || self.mutable_globals.contains(name)
            || (self.module_globals.contains(name) && !self.is_local_only_binding(name))
    }

    fn is_configured_pure_call(&self, path: &str) -> bool {
        self.config
            .allowlists
            .pure_calls
            .iter()
            .any(|entry| entry == path)
    }

    fn is_configured_impure_call(&self, path: &str) -> bool {
        self.config
            .allowlists
            .impure_calls
            .iter()
            .any(|entry| entry == path)
    }
}

fn assess_external_purity(
    function: &FunctionDef,
    context: &FunctionAssessmentContext<'_>,
) -> ExternalPurityAssessment {
    let mut builder = ExternalPurityBuilder::default();
    visit_statements_for_purity(&function.body, context, &mut builder);

    let categories = if builder.categories.is_empty() {
        vec![ExternalPurityCategory::Pure]
    } else {
        builder.categories.iter().copied().collect()
    };

    let primary = primary_purity_category(&builder.categories);

    ExternalPurityAssessment {
        primary,
        severity: context.config.purity.severity.severity_for(primary),
        categories,
        reasons: builder.reasons,
    }
}

#[derive(Default)]
struct ExternalPurityBuilder {
    categories: BTreeSet<ExternalPurityCategory>,
    reasons: Vec<ExternalPurityReason>,
    seen: BTreeSet<(ExternalPurityCategory, usize, usize, String)>,
}

impl ExternalPurityBuilder {
    fn push(
        &mut self,
        category: ExternalPurityCategory,
        location: &SourceSpan,
        message: impl Into<String>,
    ) {
        let message = message.into();
        let key = (
            category,
            location.start.line,
            location.start.column,
            message.clone(),
        );
        if self.seen.insert(key) {
            self.categories.insert(category);
            self.reasons.push(ExternalPurityReason {
                category,
                message,
                location: location.clone(),
            });
        }
    }
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum ExpressionUsage {
    Normal,
    CallCallee,
}

fn visit_statements_for_purity(
    statements: &[Statement],
    context: &FunctionAssessmentContext<'_>,
    builder: &mut ExternalPurityBuilder,
) {
    for statement in statements {
        match &statement.kind {
            StatementKind::Return { value } => {
                if let Some(value) = value {
                    visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
                }
            }
            StatementKind::Delete { targets } => {
                for target in targets {
                    visit_target_for_purity(target, context, builder);
                }
            }
            StatementKind::Assign {
                targets,
                value,
                annotation,
                ..
            } => {
                for target in targets {
                    visit_target_for_purity(target, context, builder);
                }
                if let Some(value) = value {
                    visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
                }
                if let Some(annotation) = annotation {
                    visit_expression_for_purity(
                        annotation,
                        context,
                        ExpressionUsage::Normal,
                        builder,
                    );
                }
            }
            StatementKind::For {
                iter, body, orelse, ..
            } => {
                visit_expression_for_purity(iter, context, ExpressionUsage::Normal, builder);
                visit_statements_for_purity(body, context, builder);
                visit_statements_for_purity(orelse, context, builder);
            }
            StatementKind::While { test, body, orelse }
            | StatementKind::If { test, body, orelse } => {
                visit_expression_for_purity(test, context, ExpressionUsage::Normal, builder);
                visit_statements_for_purity(body, context, builder);
                visit_statements_for_purity(orelse, context, builder);
            }
            StatementKind::With { items, body, .. } => {
                for item in items {
                    visit_expression_for_purity(
                        &item.context_expr,
                        context,
                        ExpressionUsage::Normal,
                        builder,
                    );
                }
                visit_statements_for_purity(body, context, builder);
            }
            StatementKind::Try {
                body,
                handlers,
                orelse,
                finalbody,
            } => {
                if !handlers.is_empty() && block_contains_raise(body) {
                    builder.push(
                        ExternalPurityCategory::ImpureExceptionalControlDependency,
                        &statement.location,
                        "Uses exception-based control flow that changes behavior through raises and handlers.",
                    );
                }
                visit_statements_for_purity(body, context, builder);
                for handler in handlers {
                    if let Some(type_expr) = &handler.type_expr {
                        visit_expression_for_purity(
                            type_expr,
                            context,
                            ExpressionUsage::Normal,
                            builder,
                        );
                    }
                    visit_statements_for_purity(&handler.body, context, builder);
                }
                visit_statements_for_purity(orelse, context, builder);
                visit_statements_for_purity(finalbody, context, builder);
            }
            StatementKind::Raise { exc, cause } => {
                if let Some(exc) = exc {
                    visit_expression_for_purity(exc, context, ExpressionUsage::Normal, builder);
                }
                if let Some(cause) = cause {
                    visit_expression_for_purity(cause, context, ExpressionUsage::Normal, builder);
                }
            }
            StatementKind::Assert { test, message } => {
                visit_expression_for_purity(test, context, ExpressionUsage::Normal, builder);
                if let Some(message) = message {
                    visit_expression_for_purity(message, context, ExpressionUsage::Normal, builder);
                }
            }
            StatementKind::Global { names } | StatementKind::Nonlocal { names } => {
                builder.push(
                    ExternalPurityCategory::ImpureWritesExternalState,
                    &statement.location,
                    format!(
                        "Declares external state access through `{}`.",
                        names.join(", ")
                    ),
                );
            }
            StatementKind::Expression { expression } => {
                visit_expression_for_purity(expression, context, ExpressionUsage::Normal, builder);
            }
            StatementKind::FunctionDef(_)
            | StatementKind::ClassDef { .. }
            | StatementKind::Import { .. }
            | StatementKind::ControlFlow { .. }
            | StatementKind::Other { .. } => {}
        }
    }
}

fn visit_target_for_purity(
    target: &Expression,
    context: &FunctionAssessmentContext<'_>,
    builder: &mut ExternalPurityBuilder,
) {
    match &target.kind {
        ExpressionKind::Name { id } => {
            if context.declared_globals.contains(id) {
                builder.push(
                    ExternalPurityCategory::ImpureWritesExternalState,
                    &target.location,
                    format!("Writes global state `{id}`."),
                );
            }
        }
        ExpressionKind::Attribute { value, .. } | ExpressionKind::Subscript { value, .. } => {
            if let Some(base_name) = value.base_name() {
                if context.is_external_write_base(&base_name)
                    || (!context.is_local_only_binding(&base_name)
                        && context.config.purity.treat_local_mutation_as_impurity)
                {
                    builder.push(
                        ExternalPurityCategory::ImpureWritesExternalState,
                        &target.location,
                        format!("Mutates externally visible state `{base_name}`."),
                    );
                }
            }
        }
        _ => {}
    }
}

fn visit_expression_for_purity(
    expression: &Expression,
    context: &FunctionAssessmentContext<'_>,
    usage: ExpressionUsage,
    builder: &mut ExternalPurityBuilder,
) {
    if usage == ExpressionUsage::Normal {
        if let Some(reference) = referential_transparency_reference_path(expression) {
            builder.push(
                ExternalPurityCategory::ImpureNondeterministic,
                &expression.location,
                format!("Reads ambient runtime state through `{reference}`."),
            );
        }

        if let Some(base_name) = expression.base_name() {
            if context.is_external_read_base(&base_name) {
                builder.push(
                    ExternalPurityCategory::ImpureReadsExternalState,
                    &expression.location,
                    format!("Reads mutable external state `{base_name}`."),
                );
            } else if context.config.purity.self_attribute_reads_as_external_state
                && matches!(base_name.as_str(), "self" | "cls")
                && matches!(
                    expression.kind,
                    ExpressionKind::Attribute { .. } | ExpressionKind::Subscript { .. }
                )
            {
                builder.push(
                    ExternalPurityCategory::ImpureReadsExternalState,
                    &expression.location,
                    format!("Reads instance or class state through `{base_name}`."),
                );
            }
        }
    }

    match &expression.kind {
        ExpressionKind::Attribute { value, .. }
        | ExpressionKind::Unary { operand: value, .. }
        | ExpressionKind::Await { value }
        | ExpressionKind::YieldFrom { value }
        | ExpressionKind::Starred { value } => {
            visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
        }
        ExpressionKind::Call {
            function,
            args,
            keywords,
        } => {
            if let Some(path) = function.path() {
                record_call_purity(&path, expression, function, context, builder);
            } else if let Some(path) = ambient_dependency_call_path(function) {
                builder.push(
                    ExternalPurityCategory::ImpureNondeterministic,
                    &expression.location,
                    format!("Depends on ambient runtime API `{path}`."),
                );
            }
            visit_expression_for_purity(function, context, ExpressionUsage::CallCallee, builder);
            for argument in args {
                visit_expression_for_purity(argument, context, ExpressionUsage::Normal, builder);
            }
            for keyword in keywords {
                visit_expression_for_purity(
                    &keyword.value,
                    context,
                    ExpressionUsage::Normal,
                    builder,
                );
            }
        }
        ExpressionKind::List { elements }
        | ExpressionKind::Tuple { elements }
        | ExpressionKind::Set { elements }
        | ExpressionKind::BoolOp {
            values: elements, ..
        } => {
            for element in elements {
                visit_expression_for_purity(element, context, ExpressionUsage::Normal, builder);
            }
        }
        ExpressionKind::Dict { entries } => {
            for entry in entries {
                if let Some(key) = &entry.key {
                    visit_expression_for_purity(key, context, ExpressionUsage::Normal, builder);
                }
                visit_expression_for_purity(
                    &entry.value,
                    context,
                    ExpressionUsage::Normal,
                    builder,
                );
            }
        }
        ExpressionKind::Binary { left, right, .. } => {
            visit_expression_for_purity(left, context, ExpressionUsage::Normal, builder);
            visit_expression_for_purity(right, context, ExpressionUsage::Normal, builder);
        }
        ExpressionKind::Compare {
            left, comparators, ..
        } => {
            visit_expression_for_purity(left, context, ExpressionUsage::Normal, builder);
            for comparator in comparators {
                visit_expression_for_purity(comparator, context, ExpressionUsage::Normal, builder);
            }
        }
        ExpressionKind::IfElse { test, body, orelse } => {
            visit_expression_for_purity(test, context, ExpressionUsage::Normal, builder);
            visit_expression_for_purity(body, context, ExpressionUsage::Normal, builder);
            visit_expression_for_purity(orelse, context, ExpressionUsage::Normal, builder);
        }
        ExpressionKind::Subscript { value, slice } => {
            visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
            visit_expression_for_purity(slice, context, ExpressionUsage::Normal, builder);
        }
        ExpressionKind::Yield { value } => {
            if let Some(value) = value {
                visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
            }
        }
        ExpressionKind::Named { target, value } => {
            visit_expression_for_purity(target, context, ExpressionUsage::Normal, builder);
            visit_expression_for_purity(value, context, ExpressionUsage::Normal, builder);
        }
        ExpressionKind::Slice { lower, upper, step } => {
            if let Some(lower) = lower {
                visit_expression_for_purity(lower, context, ExpressionUsage::Normal, builder);
            }
            if let Some(upper) = upper {
                visit_expression_for_purity(upper, context, ExpressionUsage::Normal, builder);
            }
            if let Some(step) = step {
                visit_expression_for_purity(step, context, ExpressionUsage::Normal, builder);
            }
        }
        ExpressionKind::Name { .. }
        | ExpressionKind::Constant { .. }
        | ExpressionKind::Lambda
        | ExpressionKind::Other { .. } => {}
    }
}

fn record_call_purity(
    path: &str,
    expression: &Expression,
    function: &Expression,
    context: &FunctionAssessmentContext<'_>,
    builder: &mut ExternalPurityBuilder,
) {
    if context.is_configured_impure_call(path) {
        builder.push(
            ExternalPurityCategory::ImpureUnknownCall,
            &expression.location,
            format!("Configured impure call `{path}` was invoked."),
        );
        return;
    }

    if is_direct_io_call(path) || is_direct_effectful_call(path) {
        builder.push(
            ExternalPurityCategory::ImpureDirectIo,
            &expression.location,
            format!("Calls effectful boundary API `{path}`."),
        );
        return;
    }

    if is_non_deterministic_call(path) {
        builder.push(
            ExternalPurityCategory::ImpureNondeterministic,
            &expression.location,
            format!("Depends on non-deterministic API `{path}`."),
        );
        return;
    }

    if let Some(ambient_path) = ambient_dependency_call_path(function) {
        builder.push(
            ExternalPurityCategory::ImpureNondeterministic,
            &expression.location,
            format!("Depends on ambient runtime API `{ambient_path}`."),
        );
        return;
    }

    if is_in_place_mutation_call(path) {
        if let Some(base_name) = function.base_name() {
            if context.is_external_write_base(&base_name)
                || (!context.is_local_only_binding(&base_name)
                    && context.config.purity.treat_local_mutation_as_impurity)
            {
                builder.push(
                    ExternalPurityCategory::ImpureWritesExternalState,
                    &expression.location,
                    format!("Mutates externally visible object `{base_name}` via `{path}`."),
                );
            }
        }
        return;
    }

    if context.is_configured_pure_call(path) || is_known_pure_call(path) {
        return;
    }

    match context.config.purity.treat_unknown_calls_as {
        UnknownCallBehavior::AssumePure => {}
        UnknownCallBehavior::AssumeImpure | UnknownCallBehavior::WarnUnknown => builder.push(
            ExternalPurityCategory::ImpureUnknownCall,
            &expression.location,
            format!("Calls `{path}` without purity metadata."),
        ),
        UnknownCallBehavior::RequireAnnotation => builder.push(
            ExternalPurityCategory::ImpureUnknownCall,
            &expression.location,
            format!("Calls `{path}` without the required purity annotation."),
        ),
    }
}

fn is_known_pure_call(path: &str) -> bool {
    matches!(
        path,
        "abs"
            | "all"
            | "any"
            | "bool"
            | "dict"
            | "enumerate"
            | "float"
            | "frozenset"
            | "int"
            | "len"
            | "list"
            | "max"
            | "min"
            | "range"
            | "reversed"
            | "round"
            | "set"
            | "sorted"
            | "str"
            | "sum"
            | "tuple"
            | "zip"
    ) || path.starts_with("math.")
        || path.starts_with("operator.")
}

fn primary_purity_category(
    categories: &BTreeSet<ExternalPurityCategory>,
) -> ExternalPurityCategory {
    if categories.contains(&ExternalPurityCategory::ImpureWritesExternalState) {
        ExternalPurityCategory::ImpureWritesExternalState
    } else if categories.contains(&ExternalPurityCategory::ImpureDirectIo) {
        ExternalPurityCategory::ImpureDirectIo
    } else if categories.contains(&ExternalPurityCategory::ImpureNondeterministic) {
        ExternalPurityCategory::ImpureNondeterministic
    } else if categories.contains(&ExternalPurityCategory::ImpureReadsExternalState) {
        ExternalPurityCategory::ImpureReadsExternalState
    } else if categories.contains(&ExternalPurityCategory::ImpureUnknownCall) {
        ExternalPurityCategory::ImpureUnknownCall
    } else if categories.contains(&ExternalPurityCategory::ImpureExceptionalControlDependency) {
        ExternalPurityCategory::ImpureExceptionalControlDependency
    } else {
        ExternalPurityCategory::Pure
    }
}

fn assess_implementation_style(
    function: &FunctionDef,
    context: &FunctionAssessmentContext<'_>,
) -> ImplementationStyleAssessment {
    let mut metrics = StyleMetricsBuilder::default();
    for name in function
        .parameters
        .iter()
        .map(|parameter| parameter.name.as_str())
    {
        metrics.assignment_counts.insert(name.to_owned(), 1);
    }
    for name in context
        .declared_globals
        .iter()
        .chain(context.declared_nonlocals.iter())
        .chain(context.closure_bindings.iter())
    {
        metrics.assignment_counts.insert(name.clone(), 1);
    }
    visit_statements_for_style(&function.body, context, 0, &mut metrics);

    let style_metrics = metrics.finish();
    let mut categories = Vec::new();

    if style_metrics.reassignment_count == 0 {
        categories.push(ImplementationStyleCategory::NoReassignment);
    } else {
        categories.push(ImplementationStyleCategory::HasReassignment);
        if style_metrics.reassignment_count >= context.config.style.reassignment_threshold {
            categories.push(ImplementationStyleCategory::ReassignmentHeavy);
        }
    }

    if style_metrics.mutation_count > 0 {
        categories.push(ImplementationStyleCategory::HasMutation);
        if style_metrics.mutation_count >= context.config.style.mutation_threshold {
            categories.push(ImplementationStyleCategory::MutationHeavy);
        }
    }

    if style_metrics.imperative_score >= context.config.style.imperative_style_threshold {
        categories.push(ImplementationStyleCategory::ImperativeStyleHigh);
    }

    let primary = if categories.contains(&ImplementationStyleCategory::ImperativeStyleHigh) {
        ImplementationStyleCategory::ImperativeStyleHigh
    } else if categories.contains(&ImplementationStyleCategory::MutationHeavy) {
        ImplementationStyleCategory::MutationHeavy
    } else if categories.contains(&ImplementationStyleCategory::ReassignmentHeavy) {
        ImplementationStyleCategory::ReassignmentHeavy
    } else if categories.contains(&ImplementationStyleCategory::HasMutation) {
        ImplementationStyleCategory::HasMutation
    } else if categories.contains(&ImplementationStyleCategory::HasReassignment) {
        ImplementationStyleCategory::HasReassignment
    } else {
        ImplementationStyleCategory::NoReassignment
    };

    ImplementationStyleAssessment {
        primary,
        severity: context.config.style.severity.severity_for(primary),
        categories,
        metrics: style_metrics,
    }
}

#[derive(Default)]
struct StyleMetricsBuilder {
    assignment_counts: BTreeMap<String, usize>,
    intermediate_names: BTreeSet<String>,
    reassignment_count: usize,
    cumulative_reassignment_count: usize,
    mutation_count: usize,
    loop_count: usize,
    branch_count: usize,
    max_nesting_depth: usize,
}

impl StyleMetricsBuilder {
    fn record_assignment_name(
        &mut self,
        name: &str,
        context: &FunctionAssessmentContext<'_>,
        is_cumulative: bool,
    ) {
        let count = self.assignment_counts.entry(name.to_owned()).or_insert(0);
        let previous_count = *count;
        *count += 1;
        if previous_count > 0 || is_cumulative {
            self.reassignment_count += 1;
        }
        if is_cumulative {
            self.cumulative_reassignment_count += 1;
        }
        if !context.parameter_names.contains(name) {
            self.intermediate_names.insert(name.to_owned());
        }
    }

    fn finish(self) -> StyleMetrics {
        let imperative_score = self.reassignment_count
            + self.cumulative_reassignment_count
            + self.mutation_count
            + self.loop_count
            + self.branch_count
            + self.max_nesting_depth
            + self.intermediate_names.len();

        StyleMetrics {
            reassignment_count: self.reassignment_count,
            cumulative_reassignment_count: self.cumulative_reassignment_count,
            mutation_count: self.mutation_count,
            loop_count: self.loop_count,
            branch_count: self.branch_count,
            max_nesting_depth: self.max_nesting_depth,
            intermediate_variable_count: self.intermediate_names.len(),
            imperative_score,
        }
    }
}

fn visit_statements_for_style(
    statements: &[Statement],
    context: &FunctionAssessmentContext<'_>,
    nesting_depth: usize,
    metrics: &mut StyleMetricsBuilder,
) {
    metrics.max_nesting_depth = metrics.max_nesting_depth.max(nesting_depth);

    for statement in statements {
        match &statement.kind {
            StatementKind::Assign {
                targets,
                value,
                kind,
                ..
            } => {
                for target in targets {
                    for name in extract_target_names(target) {
                        metrics.record_assignment_name(
                            &name,
                            context,
                            is_cumulative_assignment(
                                &name,
                                matches!(kind, AssignKind::Augmented),
                                value.as_ref(),
                                context,
                            ),
                        );
                    }
                    if matches!(
                        target.kind,
                        ExpressionKind::Attribute { .. } | ExpressionKind::Subscript { .. }
                    ) {
                        metrics.mutation_count += 1;
                    }
                }
                if let Some(value) = value {
                    visit_expression_for_style(value, metrics);
                }
            }
            StatementKind::For {
                iter, body, orelse, ..
            } => {
                metrics.loop_count += 1;
                visit_expression_for_style(iter, metrics);
                visit_statements_for_style(body, context, nesting_depth + 1, metrics);
                visit_statements_for_style(orelse, context, nesting_depth + 1, metrics);
            }
            StatementKind::While { test, body, orelse } => {
                metrics.loop_count += 1;
                visit_expression_for_style(test, metrics);
                visit_statements_for_style(body, context, nesting_depth + 1, metrics);
                visit_statements_for_style(orelse, context, nesting_depth + 1, metrics);
            }
            StatementKind::If { test, body, orelse } => {
                metrics.branch_count += 1;
                visit_expression_for_style(test, metrics);
                visit_statements_for_style(body, context, nesting_depth + 1, metrics);
                visit_statements_for_style(orelse, context, nesting_depth + 1, metrics);
            }
            StatementKind::With { items, body, .. } => {
                for item in items {
                    visit_expression_for_style(&item.context_expr, metrics);
                }
                visit_statements_for_style(body, context, nesting_depth + 1, metrics);
            }
            StatementKind::Try {
                body,
                handlers,
                orelse,
                finalbody,
            } => {
                visit_statements_for_style(body, context, nesting_depth + 1, metrics);
                for handler in handlers {
                    visit_statements_for_style(&handler.body, context, nesting_depth + 1, metrics);
                }
                visit_statements_for_style(orelse, context, nesting_depth + 1, metrics);
                visit_statements_for_style(finalbody, context, nesting_depth + 1, metrics);
            }
            StatementKind::Return { value } => {
                if let Some(value) = value {
                    visit_expression_for_style(value, metrics);
                }
            }
            StatementKind::Delete { targets } => {
                for target in targets {
                    visit_expression_for_style(target, metrics);
                }
            }
            StatementKind::Raise { exc, cause } => {
                if let Some(exc) = exc {
                    visit_expression_for_style(exc, metrics);
                }
                if let Some(cause) = cause {
                    visit_expression_for_style(cause, metrics);
                }
            }
            StatementKind::Assert { test, message } => {
                visit_expression_for_style(test, metrics);
                if let Some(message) = message {
                    visit_expression_for_style(message, metrics);
                }
            }
            StatementKind::Expression { expression } => {
                visit_expression_for_style(expression, metrics);
            }
            StatementKind::FunctionDef(_)
            | StatementKind::ClassDef { .. }
            | StatementKind::Import { .. }
            | StatementKind::Global { .. }
            | StatementKind::Nonlocal { .. }
            | StatementKind::ControlFlow { .. }
            | StatementKind::Other { .. } => {}
        }
    }
}

fn visit_expression_for_style(expression: &Expression, metrics: &mut StyleMetricsBuilder) {
    match &expression.kind {
        ExpressionKind::Call {
            function,
            args,
            keywords,
        } => {
            if function
                .path()
                .is_some_and(|path| is_in_place_mutation_call(&path))
            {
                metrics.mutation_count += 1;
            }
            visit_expression_for_style(function, metrics);
            for argument in args {
                visit_expression_for_style(argument, metrics);
            }
            for keyword in keywords {
                visit_expression_for_style(&keyword.value, metrics);
            }
        }
        ExpressionKind::Attribute { value, .. }
        | ExpressionKind::Unary { operand: value, .. }
        | ExpressionKind::Await { value }
        | ExpressionKind::YieldFrom { value }
        | ExpressionKind::Starred { value } => visit_expression_for_style(value, metrics),
        ExpressionKind::List { elements }
        | ExpressionKind::Tuple { elements }
        | ExpressionKind::Set { elements }
        | ExpressionKind::BoolOp {
            values: elements, ..
        } => {
            for element in elements {
                visit_expression_for_style(element, metrics);
            }
        }
        ExpressionKind::Dict { entries } => {
            for entry in entries {
                if let Some(key) = &entry.key {
                    visit_expression_for_style(key, metrics);
                }
                visit_expression_for_style(&entry.value, metrics);
            }
        }
        ExpressionKind::Binary { left, right, .. } => {
            visit_expression_for_style(left, metrics);
            visit_expression_for_style(right, metrics);
        }
        ExpressionKind::Compare {
            left, comparators, ..
        } => {
            visit_expression_for_style(left, metrics);
            for comparator in comparators {
                visit_expression_for_style(comparator, metrics);
            }
        }
        ExpressionKind::IfElse { test, body, orelse } => {
            visit_expression_for_style(test, metrics);
            visit_expression_for_style(body, metrics);
            visit_expression_for_style(orelse, metrics);
        }
        ExpressionKind::Subscript { value, slice } => {
            visit_expression_for_style(value, metrics);
            visit_expression_for_style(slice, metrics);
        }
        ExpressionKind::Yield { value } => {
            if let Some(value) = value {
                visit_expression_for_style(value, metrics);
            }
        }
        ExpressionKind::Named { target, value } => {
            visit_expression_for_style(target, metrics);
            visit_expression_for_style(value, metrics);
        }
        ExpressionKind::Slice { lower, upper, step } => {
            if let Some(lower) = lower {
                visit_expression_for_style(lower, metrics);
            }
            if let Some(upper) = upper {
                visit_expression_for_style(upper, metrics);
            }
            if let Some(step) = step {
                visit_expression_for_style(step, metrics);
            }
        }
        ExpressionKind::Name { .. }
        | ExpressionKind::Constant { .. }
        | ExpressionKind::Lambda
        | ExpressionKind::Other { .. } => {}
    }
}

fn block_contains_raise(statements: &[Statement]) -> bool {
    statements.iter().any(|statement| match &statement.kind {
        StatementKind::Raise { .. } => true,
        StatementKind::If { body, orelse, .. }
        | StatementKind::While { body, orelse, .. }
        | StatementKind::For { body, orelse, .. } => {
            block_contains_raise(body) || block_contains_raise(orelse)
        }
        StatementKind::With { body, .. } => block_contains_raise(body),
        StatementKind::Try {
            body,
            handlers,
            orelse,
            finalbody,
        } => {
            block_contains_raise(body)
                || handlers
                    .iter()
                    .any(|handler| block_contains_raise(&handler.body))
                || block_contains_raise(orelse)
                || block_contains_raise(finalbody)
        }
        _ => false,
    })
}

fn collect_declared_globals(function: &FunctionDef) -> BTreeSet<String> {
    let mut names = BTreeSet::new();
    collect_declared_names_in_block(&function.body, &mut names, false);
    names
}

fn collect_declared_nonlocals(function: &FunctionDef) -> BTreeSet<String> {
    let mut names = BTreeSet::new();
    collect_declared_names_in_block(&function.body, &mut names, true);
    names
}

fn collect_declared_names_in_block(
    statements: &[Statement],
    names: &mut BTreeSet<String>,
    collect_nonlocal: bool,
) {
    for statement in statements {
        match &statement.kind {
            StatementKind::Global { names: declared } if !collect_nonlocal => {
                names.extend(declared.iter().cloned());
            }
            StatementKind::Nonlocal { names: declared } if collect_nonlocal => {
                names.extend(declared.iter().cloned());
            }
            StatementKind::For { body, orelse, .. }
            | StatementKind::While { body, orelse, .. }
            | StatementKind::If { body, orelse, .. } => {
                collect_declared_names_in_block(body, names, collect_nonlocal);
                collect_declared_names_in_block(orelse, names, collect_nonlocal);
            }
            StatementKind::With { body, .. } => {
                collect_declared_names_in_block(body, names, collect_nonlocal);
            }
            StatementKind::Try {
                body,
                handlers,
                orelse,
                finalbody,
            } => {
                collect_declared_names_in_block(body, names, collect_nonlocal);
                for handler in handlers {
                    collect_declared_names_in_block(&handler.body, names, collect_nonlocal);
                }
                collect_declared_names_in_block(orelse, names, collect_nonlocal);
                collect_declared_names_in_block(finalbody, names, collect_nonlocal);
            }
            _ => {}
        }
    }
}

fn collect_closure_bindings(module: &Module, function: &FunctionDef) -> BTreeSet<String> {
    let mut bindings = BTreeSet::new();
    let mut prefix = function.qualified_name.as_str();

    while let Some((parent, _)) = prefix.rsplit_once('.') {
        if let Some(parent_function) = module
            .functions
            .iter()
            .find(|candidate| candidate.qualified_name == parent)
        {
            bindings.extend(parent_function.local_bindings.iter().cloned());
        }
        prefix = parent;
    }

    bindings
}

fn is_cumulative_assignment(
    name: &str,
    is_augmented: bool,
    value: Option<&Expression>,
    context: &FunctionAssessmentContext<'_>,
) -> bool {
    is_augmented
        || value.is_some_and(|value| {
            expression_mentions_name(value, name)
                && (context.parameter_names.contains(name)
                    || context.declared_globals.contains(name)
                    || context.declared_nonlocals.contains(name)
                    || context.closure_bindings.contains(name)
                    || context.local_bindings.contains(name))
        })
}

fn expression_mentions_name(expression: &Expression, name: &str) -> bool {
    match &expression.kind {
        ExpressionKind::Name { id } => id == name,
        ExpressionKind::Attribute { value, .. }
        | ExpressionKind::Unary { operand: value, .. }
        | ExpressionKind::Await { value }
        | ExpressionKind::YieldFrom { value }
        | ExpressionKind::Starred { value } => expression_mentions_name(value, name),
        ExpressionKind::Call {
            function,
            args,
            keywords,
        } => {
            expression_mentions_name(function, name)
                || args
                    .iter()
                    .any(|argument| expression_mentions_name(argument, name))
                || keywords
                    .iter()
                    .any(|keyword| expression_mentions_name(&keyword.value, name))
        }
        ExpressionKind::List { elements }
        | ExpressionKind::Tuple { elements }
        | ExpressionKind::Set { elements }
        | ExpressionKind::BoolOp {
            values: elements, ..
        } => elements
            .iter()
            .any(|element| expression_mentions_name(element, name)),
        ExpressionKind::Dict { entries } => entries.iter().any(|entry| {
            entry
                .key
                .as_ref()
                .is_some_and(|key| expression_mentions_name(key, name))
                || expression_mentions_name(&entry.value, name)
        }),
        ExpressionKind::Binary { left, right, .. } => {
            expression_mentions_name(left, name) || expression_mentions_name(right, name)
        }
        ExpressionKind::Compare {
            left, comparators, ..
        } => {
            expression_mentions_name(left, name)
                || comparators
                    .iter()
                    .any(|comparator| expression_mentions_name(comparator, name))
        }
        ExpressionKind::IfElse { test, body, orelse } => {
            expression_mentions_name(test, name)
                || expression_mentions_name(body, name)
                || expression_mentions_name(orelse, name)
        }
        ExpressionKind::Subscript { value, slice } => {
            expression_mentions_name(value, name) || expression_mentions_name(slice, name)
        }
        ExpressionKind::Yield { value } => value
            .as_ref()
            .is_some_and(|inner| expression_mentions_name(inner, name)),
        ExpressionKind::Named { target, value } => {
            expression_mentions_name(target, name) || expression_mentions_name(value, name)
        }
        ExpressionKind::Slice { lower, upper, step } => {
            lower
                .as_ref()
                .is_some_and(|lower| expression_mentions_name(lower, name))
                || upper
                    .as_ref()
                    .is_some_and(|upper| expression_mentions_name(upper, name))
                || step
                    .as_ref()
                    .is_some_and(|step| expression_mentions_name(step, name))
        }
        ExpressionKind::Constant { .. } | ExpressionKind::Lambda | ExpressionKind::Other { .. } => {
            false
        }
    }
}

fn extract_target_names(target: &Expression) -> Vec<String> {
    let mut names = Vec::new();
    collect_names_from_target(target, &mut names);
    names
}

fn collect_names_from_target(target: &Expression, names: &mut Vec<String>) {
    match &target.kind {
        ExpressionKind::Name { id } => names.push(id.clone()),
        ExpressionKind::Tuple { elements } | ExpressionKind::List { elements } => {
            for element in elements {
                collect_names_from_target(element, names);
            }
        }
        ExpressionKind::Starred { value } => collect_names_from_target(value, names),
        _ => {}
    }
}
