use crate::preview::is_enumerate_for_loop_int_index_enabled;
use ruff_macros::{ViolationMetadata, derive_message_formats};
use ruff_python_ast::statement_visitor::{StatementVisitor, walk_stmt};
use ruff_python_ast::{self as ast, Expr, Int, Number, Operator, Stmt};
use ruff_python_semantic::analyze::type_inference::{NumberLike, PythonType, ResolvedPythonType};
use ruff_python_semantic::analyze::typing;
use ruff_text_size::Ranged;

use crate::Violation;
use crate::checkers::ast::Checker;

/// ## What it does
/// Checks for `for` loops with explicit loop-index variables that can be replaced
/// with `enumerate()`.
///
/// In [preview], this rule checks for index variables initialized with any integer rather than only
/// a literal zero.
///
/// ## Why is this bad?
/// When iterating over a sequence, it's often desirable to keep track of the
/// index of each element alongside the element itself. Prefer the `enumerate`
/// builtin over manually incrementing a counter variable within the loop, as
/// `enumerate` is more concise and idiomatic.
///
/// ## Example
/// ```python
/// fruits = ["apple", "banana", "cherry"]
/// i = 0
/// for fruit in fruits:
///     print(f"{i + 1}. {fruit}")
///     i += 1
/// ```
///
/// Use instead:
/// ```python
/// fruits = ["apple", "banana", "cherry"]
/// for i, fruit in enumerate(fruits):
///     print(f"{i + 1}. {fruit}")
/// ```
///
/// ## References
/// - [Python documentation: `enumerate`](https://docs.python.org/3/library/functions.html#enumerate)
///
/// [preview]: https://docs.astral.sh/ruff/preview/
#[derive(ViolationMetadata)]
#[violation_metadata(stable_since = "v0.2.0")]
pub(crate) struct EnumerateForLoop {
    index: String,
}

impl Violation for EnumerateForLoop {
    #[derive_message_formats]
    fn message(&self) -> String {
        let EnumerateForLoop { index } = self;
        format!("Use `enumerate()` for index variable `{index}` in `for` loop")
    }
}

/// SIM113
pub(crate) fn enumerate_for_loop(checker: &Checker, for_stmt: &ast::StmtFor) {
    // If the loop is async, abort.
    if for_stmt.is_async {
        return;
    }

    // If the loop contains a `continue`, abort.
    let mut visitor = LoopControlFlowVisitor::default();
    visitor.visit_body(&for_stmt.body);
    if visitor.has_continue {
        return;
    }

    for stmt in &for_stmt.body {
        // Find the augmented assignment expression (e.g., `i += 1`).
        if let Some(index) = match_index_increment(stmt) {
            // Find the binding corresponding to the initialization (e.g., `i = 1`).
            let Some(id) = checker.semantic().resolve_name(index) else {
                continue;
            };

            // If it's not an assignment (e.g., it's a function argument), ignore it.
            let initial_binding = checker.semantic().binding(id);
            if !initial_binding.kind.is_assignment() {
                continue;
            }

            // If the variable is global or nonlocal, ignore it.
            if initial_binding.is_global() || initial_binding.is_nonlocal() {
                continue;
            }

            // Ensure that the index variable was initialized to 0 (or instance of `int` if preview is enabled).
            let Some(value) = typing::find_binding_value(initial_binding, checker.semantic())
            else {
                continue;
            };
            if !(matches!(
                value,
                Expr::NumberLiteral(ast::ExprNumberLiteral {
                    value: Number::Int(Int::ZERO),
                    ..
                })
            ) || matches!(
                ResolvedPythonType::from(value),
                ResolvedPythonType::Atom(PythonType::Number(NumberLike::Integer))
            ) && is_enumerate_for_loop_int_index_enabled(checker.settings()))
            {
                continue;
            }

            // If the binding is not at the same level as the `for` loop (e.g., it's in an `if`),
            // ignore it.
            let Some(for_loop_id) = checker.semantic().current_statement_id() else {
                continue;
            };
            let Some(assignment_id) = initial_binding.source else {
                continue;
            };
            if checker.semantic().parent_statement_id(for_loop_id)
                != checker.semantic().parent_statement_id(assignment_id)
            {
                continue;
            }

            // Identify the binding created by the augmented assignment.
            // TODO(charlie): There should be a way to go from `ExprName` to `BindingId` (like
            // `resolve_name`, but for bindings rather than references).
            let increment_binding = {
                let mut bindings = checker
                    .semantic()
                    .current_scope()
                    .get_all(&index.id)
                    .map(|id| checker.semantic().binding(id))
                    .filter(|binding| for_stmt.range().contains_range(binding.range()));

                let Some(binding) = bindings.next() else {
                    continue;
                };

                // If there are multiple assignments to this variable _within_ the loop, ignore it.
                if bindings.next().is_some() {
                    continue;
                }

                binding
            };

            // Reassignments in the same scope inherit older references. Ignore anything that
            // predates the counter's initialization and only consider uses in its current lifetime.
            let initial_binding_start = initial_binding.range.start();
            if increment_binding.references().any(|id| {
                let reference = checker.semantic().reference(id);
                reference.start() >= initial_binding_start
                    && !for_stmt.range().contains_range(reference.range())
            }) {
                continue;
            }

            checker.report_diagnostic(
                EnumerateForLoop {
                    index: index.id.to_string(),
                },
                stmt.range(),
            );
        }
    }
}

/// If the statement is an index increment statement (e.g., `i += 1`), return
/// the name of the index variable.
fn match_index_increment(stmt: &Stmt) -> Option<&ast::ExprName> {
    let Stmt::AugAssign(ast::StmtAugAssign {
        target,
        op: Operator::Add,
        value,
        ..
    }) = stmt
    else {
        return None;
    };

    let name = target.as_name_expr()?;

    if matches!(
        value.as_ref(),
        Expr::NumberLiteral(ast::ExprNumberLiteral {
            value: Number::Int(Int::ONE),
            ..
        })
    ) {
        return Some(name);
    }

    None
}

#[derive(Debug, Default)]
struct LoopControlFlowVisitor {
    has_continue: bool,
}

impl StatementVisitor<'_> for LoopControlFlowVisitor {
    fn visit_stmt(&mut self, stmt: &Stmt) {
        match stmt {
            Stmt::Continue(_) => self.has_continue = true,
            Stmt::FunctionDef(_) | Stmt::ClassDef(_) => {
                // Don't recurse.
            }
            _ => walk_stmt(self, stmt),
        }
    }
}
