use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::keywords::Keyword;
use sqlparser::tokenizer::{Token, Tokenizer};

pub fn requires_primary(sql: &str) -> bool {
    let Ok(tokens) = Tokenizer::new(&PostgreSqlDialect {}, sql).tokenize() else {
        return true;
    };
    let mut tokens = tokens
        .into_iter()
        .filter(|token| !matches!(token, Token::Whitespace(_)))
        .peekable();
    let mut previous = [Keyword::NoKeyword; 2];
    let mut statement_start = true;

    while let Some(token) = tokens.next() {
        match token {
            Token::SemiColon => {
                previous = [Keyword::NoKeyword; 2];
                statement_start = true;
                continue;
            }
            Token::Word(word) => {
                if word.quote_style.is_none() {
                    if matches!(
                        word.keyword,
                        Keyword::ABORT
                            | Keyword::ALTER
                            | Keyword::ANALYZE
                            | Keyword::BEGIN
                            | Keyword::CALL
                            | Keyword::CLUSTER
                            | Keyword::COMMENT
                            | Keyword::COMMIT
                            | Keyword::COPY
                            | Keyword::CREATE
                            | Keyword::DEALLOCATE
                            | Keyword::DELETE
                            | Keyword::DISCARD
                            | Keyword::DO
                            | Keyword::DROP
                            | Keyword::EXECUTE
                            | Keyword::GRANT
                            | Keyword::INSERT
                            | Keyword::INTO
                            | Keyword::LISTEN
                            | Keyword::LOAD
                            | Keyword::LOCK
                            | Keyword::MERGE
                            | Keyword::NOTIFY
                            | Keyword::PREPARE
                            | Keyword::REFRESH
                            | Keyword::REINDEX
                            | Keyword::RELEASE
                            | Keyword::RESET
                            | Keyword::REVOKE
                            | Keyword::ROLLBACK
                            | Keyword::SAVEPOINT
                            | Keyword::SET
                            | Keyword::START
                            | Keyword::TRUNCATE
                            | Keyword::UNLISTEN
                            | Keyword::UPDATE
                            | Keyword::VACUUM
                    ) {
                        return true;
                    }
                    if statement_start
                        && (word.keyword == Keyword::END
                            || [
                                "checkpoint",
                                "reassign",
                                "declare",
                                "close",
                                "fetch",
                                "move",
                            ]
                            .iter()
                            .any(|command| word.value.eq_ignore_ascii_case(command)))
                    {
                        return true;
                    }
                    if word.keyword == Keyword::SHARE
                        && (previous[1] == Keyword::FOR || previous == [Keyword::FOR, Keyword::KEY])
                    {
                        return true;
                    }
                }
                if matches!(tokens.peek(), Some(Token::LParen)) && primary_function(&word.value) {
                    return true;
                }
                previous = [previous[1], word.keyword];
            }
            _ => previous = [Keyword::NoKeyword; 2],
        }
        statement_start = false;
    }
    false
}

fn primary_function(name: &str) -> bool {
    let name = name.to_ascii_lowercase();
    matches!(
        name.as_str(),
        "nextval"
            | "setval"
            | "set_config"
            | "pg_notify"
            | "pg_cancel_backend"
            | "pg_terminate_backend"
    ) || name.starts_with("pg_advisory_")
        || name.starts_with("pg_try_advisory_")
}