use itertools::Itertools;

use crate::autofix::edits::pad;
use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::{self as ast, Arguments, Constant, Expr};
use ruff_text_size::{Ranged, TextRange};

use crate::autofix::snippet::SourceCodeSnippet;
use crate::checkers::ast::Checker;
use crate::registry::AsRule;
use crate::rules::flynt::helpers;

/// ## What it does
/// Checks for `str.join` calls that can be replaced with f-strings.
///
/// ## Why is this bad?
/// f-strings are more readable and generally preferred over `str.join` calls.
///
/// ## Example
/// ```python
/// " ".join((foo, bar))
/// ```
///
/// Use instead:
/// ```python
/// f"{foo} {bar}"
/// ```
///
/// ## References
/// - [Python documentation: f-strings](https://docs.python.org/3/reference/lexical_analysis.html#f-strings)
#[violation]
pub struct StaticJoinToFString {
    expression: SourceCodeSnippet,
}

impl AlwaysAutofixableViolation for StaticJoinToFString {
    #[derive_message_formats]
    fn message(&self) -> String {
        let StaticJoinToFString { expression } = self;
        if let Some(expression) = expression.full_display() {
            format!("Consider `{expression}` instead of string join")
        } else {
            format!("Consider f-string instead of string join")
        }
    }

    fn autofix_title(&self) -> String {
        let StaticJoinToFString { expression } = self;
        if let Some(expression) = expression.full_display() {
            format!("Replace with `{expression}`")
        } else {
            format!("Replace with f-string")
        }
    }
}

fn is_static_length(elts: &[Expr]) -> bool {
    elts.iter().all(|e| !e.is_starred_expr())
}

fn build_fstring(joiner: &str, joinees: &[Expr]) -> Option<Expr> {
    // If all elements are string constants, join them into a single string.
    if joinees.iter().all(|expr| {
        matches!(
            expr,
            Expr::Constant(ast::ExprConstant {
                value: Constant::Str(_),
                ..
            })
        )
    }) {
        let node = ast::ExprConstant {
            value: joinees
                .iter()
                .filter_map(|expr| {
                    if let Expr::Constant(ast::ExprConstant {
                        value: Constant::Str(ast::StringConstant { value, .. }),
                        ..
                    }) = expr
                    {
                        Some(value.as_str())
                    } else {
                        None
                    }
                })
                .join(joiner)
                .into(),
            range: TextRange::default(),
        };
        return Some(node.into());
    }

    let mut fstring_elems = Vec::with_capacity(joinees.len() * 2);
    let mut first = true;

    for expr in joinees {
        if expr.is_f_string_expr() {
            // Oops, already an f-string. We don't know how to handle those
            // gracefully right now.
            return None;
        }
        if !std::mem::take(&mut first) {
            fstring_elems.push(helpers::to_constant_string(joiner));
        }
        fstring_elems.push(helpers::to_f_string_element(expr)?);
    }

    let node = ast::ExprFString {
        values: fstring_elems,
        implicit_concatenated: false,
        range: TextRange::default(),
    };
    Some(node.into())
}

/// FLY002
pub(crate) fn static_join_to_fstring(checker: &mut Checker, expr: &Expr, joiner: &str) {
    let Expr::Call(ast::ExprCall {
        arguments: Arguments { args, keywords, .. },
        ..
    }) = expr
    else {
        return;
    };

    // If there are kwargs or more than one argument, this is some non-standard
    // string join call.
    if !keywords.is_empty() {
        return;
    }
    let [arg] = args.as_slice() else {
        return;
    };

    // Get the elements to join; skip (e.g.) generators, sets, etc.
    let joinees = match &arg {
        Expr::List(ast::ExprList { elts, .. }) if is_static_length(elts) => elts,
        Expr::Tuple(ast::ExprTuple { elts, .. }) if is_static_length(elts) => elts,
        _ => return,
    };

    // Try to build the fstring (internally checks whether e.g. the elements are
    // convertible to f-string parts).
    let Some(new_expr) = build_fstring(joiner, joinees) else {
        return;
    };

    let contents = checker.generator().expr(&new_expr);

    let mut diagnostic = Diagnostic::new(
        StaticJoinToFString {
            expression: SourceCodeSnippet::new(contents.clone()),
        },
        expr.range(),
    );
    if checker.patch(diagnostic.kind.rule()) {
        diagnostic.set_fix(Fix::suggested(Edit::range_replacement(
            pad(contents, expr.range(), checker.locator()),
            expr.range(),
        )));
    }
    checker.diagnostics.push(diagnostic);
}
