use ruff_macros::{ViolationMetadata, derive_message_formats};
use ruff_python_ast::token::TokenKind;
use ruff_text_size::{Ranged, TextRange};

use crate::AlwaysFixableViolation;
use crate::Edit;
use crate::Fix;
use crate::checkers::ast::LintContext;

use super::{LogicalLine, Whitespace};

/// ## What it does
/// Checks for the use of extraneous whitespace after "(", "[" or "{".
///
/// ## Why is this bad?
/// [PEP 8] recommends the omission of whitespace in the following cases:
/// - "Immediately inside parentheses, brackets or braces."
/// - "Immediately before a comma, semicolon, or colon."
///
/// ## Example
/// ```python
/// spam( ham[1], {eggs: 2})
/// spam(ham[ 1], {eggs: 2})
/// spam(ham[1], { eggs: 2})
/// ```
///
/// Use instead:
/// ```python
/// spam(ham[1], {eggs: 2})
/// ```
///
/// [PEP 8]: https://peps.python.org/pep-0008/#pet-peeves
#[derive(ViolationMetadata)]
#[violation_metadata(preview_since = "v0.0.269")]
pub(crate) struct WhitespaceAfterOpenBracket {
    symbol: char,
}

impl AlwaysFixableViolation for WhitespaceAfterOpenBracket {
    #[derive_message_formats]
    fn message(&self) -> String {
        let WhitespaceAfterOpenBracket { symbol } = self;
        format!("Whitespace after '{symbol}'")
    }

    fn fix_title(&self) -> String {
        let WhitespaceAfterOpenBracket { symbol } = self;
        format!("Remove whitespace before '{symbol}'")
    }
}

/// ## What it does
/// Checks for the use of extraneous whitespace before ")", "]" or "}".
///
/// ## Why is this bad?
/// [PEP 8] recommends the omission of whitespace in the following cases:
/// - "Immediately inside parentheses, brackets or braces."
/// - "Immediately before a comma, semicolon, or colon."
///
/// ## Example
/// ```python
/// spam(ham[1], {eggs: 2} )
/// spam(ham[1 ], {eggs: 2})
/// spam(ham[1], {eggs: 2 })
/// ```
///
/// Use instead:
/// ```python
/// spam(ham[1], {eggs: 2})
/// ```
///
/// [PEP 8]: https://peps.python.org/pep-0008/#pet-peeves
#[derive(ViolationMetadata)]
#[violation_metadata(preview_since = "v0.0.269")]
pub(crate) struct WhitespaceBeforeCloseBracket {
    symbol: char,
}

impl AlwaysFixableViolation for WhitespaceBeforeCloseBracket {
    #[derive_message_formats]
    fn message(&self) -> String {
        let WhitespaceBeforeCloseBracket { symbol } = self;
        format!("Whitespace before '{symbol}'")
    }

    fn fix_title(&self) -> String {
        let WhitespaceBeforeCloseBracket { symbol } = self;
        format!("Remove whitespace before '{symbol}'")
    }
}

/// ## What it does
/// Checks for the use of extraneous whitespace before ",", ";" or ":".
///
/// ## Why is this bad?
/// [PEP 8] recommends the omission of whitespace in the following cases:
/// - "Immediately inside parentheses, brackets or braces."
/// - "Immediately before a comma, semicolon, or colon."
///
/// ## Example
/// ```python
/// if x == 4: print(x, y); x, y = y , x
/// ```
///
/// Use instead:
/// ```python
/// if x == 4: print(x, y); x, y = y, x
/// ```
///
/// [PEP 8]: https://peps.python.org/pep-0008/#pet-peeves
#[derive(ViolationMetadata)]
#[violation_metadata(preview_since = "v0.0.269")]
pub(crate) struct WhitespaceBeforePunctuation {
    symbol: char,
}

impl AlwaysFixableViolation for WhitespaceBeforePunctuation {
    #[derive_message_formats]
    fn message(&self) -> String {
        let WhitespaceBeforePunctuation { symbol } = self;
        format!("Whitespace before '{symbol}'")
    }

    fn fix_title(&self) -> String {
        let WhitespaceBeforePunctuation { symbol } = self;
        format!("Remove whitespace before '{symbol}'")
    }
}

/// E201, E202, E203
pub(crate) fn extraneous_whitespace(line: &LogicalLine, context: &LintContext) {
    let mut interpolated_strings = 0u32;
    let mut brackets = vec![];
    let mut prev_token = None;
    let mut iter = line.tokens().iter().peekable();

    while let Some(token) = iter.next() {
        let kind = token.kind();
        match kind {
            TokenKind::FStringStart | TokenKind::TStringStart => interpolated_strings += 1,
            TokenKind::FStringEnd | TokenKind::TStringEnd => {
                interpolated_strings = interpolated_strings.saturating_sub(1);
            }
            TokenKind::Lsqb => {
                brackets.push(kind);
            }
            TokenKind::Rsqb => {
                brackets.pop();
            }
            TokenKind::Lbrace => {
                brackets.push(kind);
            }
            TokenKind::Rbrace => {
                brackets.pop();
            }
            _ => {}
        }
        if let Some(symbol) = BracketOrPunctuation::from_kind(kind) {
            // Whitespace before "{" or after "}" might be required in f-strings.
            // For example,
            //
            // ```python
            // f"{ {'a': 1} }"
            // ```
            //
            // Here, `{{` / `}} would be interpreted as a single raw `{` / `}`
            // character.
            match symbol {
                BracketOrPunctuation::OpenBracket(symbol)
                    if symbol != '{' || interpolated_strings == 0 =>
                {
                    let (trailing, trailing_len) = line.trailing_whitespace(token);
                    if !matches!(trailing, Whitespace::None) {
                        let range = TextRange::at(token.end(), trailing_len);
                        if let Some(mut diagnostic) = context.report_diagnostic_if_enabled(
                            WhitespaceAfterOpenBracket { symbol },
                            range,
                        ) {
                            diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                        }
                    }
                }
                BracketOrPunctuation::CloseBracket(symbol)
                    if (symbol != '}' || interpolated_strings == 0)
                        && !matches!(prev_token, Some(TokenKind::Comma)) =>
                {
                    if let (Whitespace::Single | Whitespace::Many | Whitespace::Tab, offset) =
                        line.leading_whitespace(token)
                    {
                        let range = TextRange::at(token.start() - offset, offset);
                        if let Some(mut diagnostic) = context.report_diagnostic_if_enabled(
                            WhitespaceBeforeCloseBracket { symbol },
                            range,
                        ) {
                            diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                        }
                    }
                }
                BracketOrPunctuation::Punctuation(symbol)
                    if !matches!(prev_token, Some(TokenKind::Comma)) =>
                {
                    let whitespace = line.leading_whitespace(token);
                    if let (Whitespace::Single | Whitespace::Many | Whitespace::Tab, offset) =
                        whitespace
                    {
                        // If we're in a slice, and the token is a colon, and it has
                        // equivalent spacing on both sides, allow it.
                        if symbol == ':'
                            && brackets
                                .last()
                                .is_some_and(|kind| matches!(kind, TokenKind::Lsqb))
                        {
                            // If we're in the second half of a double colon, disallow
                            // any whitespace (e.g., `foo[1: :2]` or `foo[1 : : 2]`).
                            if matches!(prev_token, Some(TokenKind::Colon)) {
                                let range = TextRange::at(token.start() - offset, offset);
                                if let Some(mut diagnostic) = context.report_diagnostic_if_enabled(
                                    WhitespaceBeforePunctuation { symbol },
                                    range,
                                ) {
                                    diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                                }
                            } else if iter.peek().is_some_and(|token| {
                                matches!(token.kind(), TokenKind::Rsqb | TokenKind::Comma)
                            }) {
                                // Allow `foo[1 :]`, but not `foo[1  :]`.
                                // Or `foo[index :, 2]`, but not `foo[index  :, 2]`.
                                if let (Whitespace::Many | Whitespace::Tab, offset) = whitespace {
                                    let range = TextRange::at(token.start() - offset, offset);
                                    if let Some(mut diagnostic) = context
                                        .report_diagnostic_if_enabled(
                                            WhitespaceBeforePunctuation { symbol },
                                            range,
                                        )
                                    {
                                        diagnostic
                                            .set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                                    }
                                }
                            } else if iter.peek().is_some_and(|token| {
                                matches!(
                                    token.kind(),
                                    TokenKind::NonLogicalNewline | TokenKind::Comment
                                )
                            }) {
                                // Allow [
                                //      long_expression_calculating_the_index() :
                                // ]
                                // But not [
                                //      long_expression_calculating_the_index()  :
                                // ]
                                // distinct from the above case, because ruff format produces a
                                // whitespace before the colon and so should the fix
                                if let (Whitespace::Many | Whitespace::Tab, offset) = whitespace {
                                    let range = TextRange::at(token.start() - offset, offset);
                                    if let Some(mut diagnostic) = context
                                        .report_diagnostic_if_enabled(
                                            WhitespaceBeforePunctuation { symbol },
                                            range,
                                        )
                                    {
                                        diagnostic.set_fix(Fix::safe_edits(
                                            Edit::range_deletion(range),
                                            [Edit::insertion(" ".into(), token.start() - offset)],
                                        ));
                                    }
                                }
                            } else {
                                // Allow, e.g., `foo[1:2]` or `foo[1 : 2]` or `foo[1 :: 2]`.
                                let token = iter
                                    .peek()
                                    .filter(|next| matches!(next.kind(), TokenKind::Colon))
                                    .unwrap_or(&token);
                                if line.trailing_whitespace(token) != whitespace {
                                    let range = TextRange::at(token.start() - offset, offset);
                                    if let Some(mut diagnostic) = context
                                        .report_diagnostic_if_enabled(
                                            WhitespaceBeforePunctuation { symbol },
                                            range,
                                        )
                                    {
                                        diagnostic
                                            .set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                                    }
                                }
                            }
                        } else {
                            if interpolated_strings > 0
                                && symbol == ':'
                                && matches!(prev_token, Some(TokenKind::Equal))
                            {
                                // Avoid removing any whitespace for f-string debug expressions.
                                continue;
                            }
                            let range = TextRange::at(token.start() - offset, offset);
                            if let Some(mut diagnostic) = context.report_diagnostic_if_enabled(
                                WhitespaceBeforePunctuation { symbol },
                                range,
                            ) {
                                diagnostic.set_fix(Fix::safe_edit(Edit::range_deletion(range)));
                            }
                        }
                    }
                }
                _ => {}
            }
        }

        prev_token = Some(kind);
    }
}

#[derive(Debug)]
enum BracketOrPunctuation {
    OpenBracket(char),
    CloseBracket(char),
    Punctuation(char),
}

impl BracketOrPunctuation {
    fn from_kind(kind: TokenKind) -> Option<BracketOrPunctuation> {
        match kind {
            TokenKind::Lbrace => Some(BracketOrPunctuation::OpenBracket('{')),
            TokenKind::Lpar => Some(BracketOrPunctuation::OpenBracket('(')),
            TokenKind::Lsqb => Some(BracketOrPunctuation::OpenBracket('[')),
            TokenKind::Rbrace => Some(BracketOrPunctuation::CloseBracket('}')),
            TokenKind::Rpar => Some(BracketOrPunctuation::CloseBracket(')')),
            TokenKind::Rsqb => Some(BracketOrPunctuation::CloseBracket(']')),
            TokenKind::Comma => Some(BracketOrPunctuation::Punctuation(',')),
            TokenKind::Colon => Some(BracketOrPunctuation::Punctuation(':')),
            TokenKind::Semi => Some(BracketOrPunctuation::Punctuation(';')),
            _ => None,
        }
    }
}
