// - `Unsupported`: a recognizable SQL command that this classifier does not yet classify
//   (e.g. GRANT, COPY, DO, MERGE). Risk cannot be assessed, so execution requires user
//   acknowledgment (see policy::write::sql_risk).
// - `Other`: no statement keyword was found (e.g. empty input, comment-only, SELECT INTO).
//   Non-empty input is gated like `Unsupported`; empty / comment-only input executes
//   immediately (nothing to run). Invalid SQL is rejected by the database.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StatementKind {
    Select,
    Insert,
    Update { has_where: bool },
    Delete { has_where: bool },
    Create,
    Alter,
    Drop,
    Truncate,
    Transaction,
    Unsupported,
    Other,
}

pub fn classify(sql: &str) -> StatementKind {
    let lower = sql.trim().to_lowercase();
    if lower.is_empty() {
        return StatementKind::Other;
    }
    let chars: Vec<(usize, char)> = lower.char_indices().collect();
    classify_inner(&lower, &chars)
}

pub fn drop_subtype(sql: &str) -> Option<String> {
    let trimmed = sql.trim();
    let chars: Vec<(usize, char)> = trimmed.char_indices().collect();
    let tokens = collect_top_level_tokens(trimmed, &chars);
    let lowers: Vec<String> = tokens.iter().map(|(_, t)| t.to_lowercase()).collect();
    let drop_idx = lowers.iter().position(|t| t == "drop")?;
    lowers.get(drop_idx + 1).cloned()
}

pub fn first_keyword(sql: &str) -> Option<String> {
    let trimmed = sql.trim();
    let chars: Vec<(usize, char)> = trimmed.char_indices().collect();
    collect_top_level_tokens(trimmed, &chars)
        .into_iter()
        .next()
        .map(|(_, token)| token.to_uppercase())
}

pub fn extract_target_name(sql: &str, kind: &StatementKind) -> Option<String> {
    let original_trimmed = sql.trim();
    // Avoids byte-length mismatch when Unicode identifiers change size under case folding.
    let chars: Vec<(usize, char)> = original_trimmed.char_indices().collect();

    match kind {
        StatementKind::Drop => extract_drop_object_name(original_trimmed, &chars),
        StatementKind::Truncate => extract_truncate_table_name(original_trimmed, &chars),
        StatementKind::Delete { .. } => extract_delete_table_name(original_trimmed, &chars),
        StatementKind::Update { .. } => extract_update_table_name(original_trimmed, &chars),
        _ => None,
    }
}

pub(crate) fn skip_line_comment(chars: &[(usize, char)], i: usize, ch: char) -> Option<usize> {
    if ch != '-' || !next_char_is(chars, i, '-') {
        return None;
    }
    let mut cursor = i;
    while cursor < chars.len() && chars[cursor].1 != '\n' {
        cursor += 1;
    }
    Some(cursor)
}

pub(crate) fn skip_block_comment(chars: &[(usize, char)], i: usize, ch: char) -> Option<usize> {
    if ch != '/' || !next_char_is(chars, i, '*') {
        return None;
    }
    let mut cursor = i + 2;
    while cursor + 1 < chars.len() && !(chars[cursor].1 == '*' && chars[cursor + 1].1 == '/') {
        cursor += 1;
    }
    Some(cursor + 2)
}

pub(crate) fn advance_single_quote(
    chars: &[(usize, char)],
    i: usize,
    ch: char,
    in_string: &mut bool,
) -> Option<usize> {
    if ch != '\'' {
        return None;
    }
    if *in_string {
        if next_char_is(chars, i, '\'') {
            return Some(i + 2);
        }
        *in_string = false;
    } else {
        *in_string = true;
    }
    Some(i + 1)
}

pub(crate) fn skip_double_quoted_identifier(
    chars: &[(usize, char)],
    i: usize,
    ch: char,
) -> Option<usize> {
    if ch != '"' {
        return None;
    }
    let mut cursor = i + 1;
    while cursor < chars.len() {
        if chars[cursor].1 == '"' {
            if next_char_is(chars, cursor, '"') {
                cursor += 2;
            } else {
                cursor += 1;
                break;
            }
        } else {
            cursor += 1;
        }
    }
    Some(cursor)
}

pub(crate) fn skip_dollar_quoted_string(
    lower: &str,
    chars: &[(usize, char)],
    i: usize,
    byte_pos: usize,
    ch: char,
) -> Option<usize> {
    if ch != '$' {
        return None;
    }
    let mut cursor = i + 1;
    while cursor < chars.len() && (chars[cursor].1.is_alphanumeric() || chars[cursor].1 == '_') {
        cursor += 1;
    }
    if cursor >= chars.len() || chars[cursor].1 != '$' {
        return None;
    }
    let tag = &lower[byte_pos..=chars[cursor].0];
    cursor += 1;
    while cursor + tag.len() <= chars.len() {
        let candidate_start = chars[cursor].0;
        if chars[cursor].1 == '$' {
            let candidate_end = candidate_start + tag.len();
            if candidate_end <= lower.len() && &lower[candidate_start..candidate_end] == tag {
                let mut next = cursor;
                while next < chars.len() && chars[next].0 < candidate_end {
                    next += 1;
                }
                return Some(next);
            }
        }
        cursor += 1;
    }
    Some(cursor)
}

pub(crate) fn update_parentheses_depth(ch: char, depth: &mut i32) {
    if ch == '(' {
        *depth += 1;
    } else if ch == ')' {
        *depth -= 1;
    }
}

pub(crate) fn next_char_is(chars: &[(usize, char)], i: usize, expected: char) -> bool {
    i + 1 < chars.len() && chars[i + 1].1 == expected
}

fn is_word_start(chars: &[(usize, char)], i: usize) -> bool {
    if i == 0 {
        return true;
    }
    let prev = chars[i - 1].1;
    !prev.is_alphanumeric() && prev != '_'
}

fn is_keyword(s: &str, keyword: &str) -> bool {
    if !s.starts_with(keyword) {
        return false;
    }
    s[keyword.len()..]
        .chars()
        .next()
        .is_none_or(|c| !c.is_alphanumeric() && c != '_')
}

fn classify_inner(lower: &str, chars: &[(usize, char)]) -> StatementKind {
    let mut i = 0;
    let mut depth: i32 = 0;
    let mut in_string = false;
    let mut in_cte = false;
    let mut is_explain = false;
    let mut has_analyze = false;
    let mut kind: Option<StatementKind> = None;

    while i < chars.len() {
        let (byte_pos, ch) = chars[i];

        if let Some(next_i) = skip_line_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = skip_block_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = advance_single_quote(chars, i, ch, &mut in_string) {
            i = next_i;
            continue;
        }
        if in_string {
            i += 1;
            continue;
        }
        if let Some(next_i) = skip_double_quoted_identifier(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = skip_dollar_quoted_string(lower, chars, i, byte_pos, ch) {
            i = next_i;
            continue;
        }

        update_parentheses_depth(ch, &mut depth);

        // Detect ANALYZE inside EXPLAIN(...) option list, e.g. EXPLAIN (ANALYZE, VERBOSE)
        if is_explain
            && kind.is_none()
            && !has_analyze
            && depth == 1
            && (ch.is_alphabetic() || ch == '_')
            && is_word_start(chars, i)
            && is_keyword(&lower[byte_pos..], "analyze")
        {
            has_analyze = true;
        }

        // Multi-statement is SAB-102 scope; reject here for safety
        if depth == 0 && ch == ';' {
            if has_non_whitespace_after(lower, byte_pos) {
                return StatementKind::Other;
            }
            break;
        }

        if depth == 0 && (ch.is_alphabetic() || ch == '_') && is_word_start(chars, i) {
            let rest = &lower[byte_pos..];

            if kind.is_none() && is_keyword(rest, "explain") {
                is_explain = true;
                i += 1;
                continue;
            }

            if is_explain && kind.is_none() {
                if is_keyword(rest, "analyze") {
                    has_analyze = true;
                    i += 1;
                    continue;
                }
                if match_keyword(rest).is_none() {
                    // EXPLAIN option keyword (VERBOSE, COSTS, BUFFERS, …)
                    i += 1;
                    continue;
                }
                if !has_analyze {
                    // Plain EXPLAIN never executes the inner statement; always read-only.
                    return StatementKind::Select;
                }
                // EXPLAIN ANALYZE executes the inner statement; classify it normally.
            }

            // SELECT INTO creates a table; treat as Other to avoid silent execution.
            if matches!(kind, Some(StatementKind::Select)) && is_keyword(rest, "into") {
                return StatementKind::Other;
            }

            if kind.is_none() && is_keyword(rest, "with") {
                in_cte = true;
                i += 1;
                continue;
            }

            // CTE: keep overriding with each top-level keyword (last one wins).
            // Non-CTE: first keyword determines the kind.
            if in_cte || kind.is_none() {
                if let Some(k) = match_keyword(rest) {
                    kind = Some(k);
                    if matches!(
                        kind,
                        Some(StatementKind::Update { .. } | StatementKind::Delete { .. })
                    ) {
                        let has_where = scan_for_where(lower, chars, i);
                        kind = Some(match kind.unwrap() {
                            StatementKind::Update { .. } => StatementKind::Update { has_where },
                            StatementKind::Delete { .. } => StatementKind::Delete { has_where },
                            other => other,
                        });
                    }
                } else if !in_cte {
                    kind = Some(StatementKind::Unsupported);
                }
            }
        }

        i += 1;
    }

    kind.unwrap_or(StatementKind::Other)
}

fn match_keyword(rest: &str) -> Option<StatementKind> {
    if is_keyword(rest, "select") {
        Some(StatementKind::Select)
    } else if is_keyword(rest, "insert") {
        Some(StatementKind::Insert)
    } else if is_keyword(rest, "update") {
        Some(StatementKind::Update { has_where: false })
    } else if is_keyword(rest, "delete") {
        Some(StatementKind::Delete { has_where: false })
    } else if is_keyword(rest, "create") {
        Some(StatementKind::Create)
    } else if is_keyword(rest, "alter") {
        Some(StatementKind::Alter)
    } else if is_keyword(rest, "drop") {
        Some(StatementKind::Drop)
    } else if is_keyword(rest, "truncate") {
        Some(StatementKind::Truncate)
    } else if is_keyword(rest, "begin")
        || is_keyword(rest, "commit")
        || is_keyword(rest, "rollback")
        || is_keyword(rest, "savepoint")
        || is_keyword(rest, "release")
    {
        Some(StatementKind::Transaction)
    } else if is_keyword(rest, "start") {
        if rest.len() > 6 {
            let after_start = rest[6..].trim_start();
            if is_keyword(after_start, "transaction") {
                return Some(StatementKind::Transaction);
            }
        }
        None
    } else if is_keyword(rest, "show") {
        Some(StatementKind::Select)
    } else {
        None
    }
}

fn scan_for_where(lower: &str, chars: &[(usize, char)], start: usize) -> bool {
    let mut i = start;
    let mut depth: i32 = 0;
    let mut in_string = false;

    while i < chars.len() {
        let (byte_pos, ch) = chars[i];

        if let Some(next_i) = skip_line_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = skip_block_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = advance_single_quote(chars, i, ch, &mut in_string) {
            i = next_i;
            continue;
        }
        if in_string {
            i += 1;
            continue;
        }
        if let Some(next_i) = skip_double_quoted_identifier(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = skip_dollar_quoted_string(lower, chars, i, byte_pos, ch) {
            i = next_i;
            continue;
        }

        update_parentheses_depth(ch, &mut depth);

        if depth == 0 && is_word_start(chars, i) {
            let rest = &lower[byte_pos..];
            if is_keyword(rest, "where") {
                return true;
            }
        }

        i += 1;
    }

    false
}

fn has_non_whitespace_after(lower: &str, byte_pos: usize) -> bool {
    lower
        .get(byte_pos + 1..)
        .is_some_and(|tail| !tail.trim().is_empty())
}

fn collect_top_level_tokens(original: &str, chars: &[(usize, char)]) -> Vec<(usize, String)> {
    let mut tokens = Vec::new();
    let mut i = 0;
    let mut depth: i32 = 0;
    let mut in_string = false;

    while i < chars.len() {
        let (byte_pos, ch) = chars[i];

        if let Some(next_i) = skip_line_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = skip_block_comment(chars, i, ch) {
            i = next_i;
            continue;
        }
        if let Some(next_i) = advance_single_quote(chars, i, ch, &mut in_string) {
            i = next_i;
            continue;
        }
        if in_string {
            i += 1;
            continue;
        }

        if ch == '"' {
            let start_i = i;
            if let Some(next_i) = skip_double_quoted_identifier(chars, i, ch) {
                if depth == 0 {
                    let start_byte = chars[start_i].0;
                    let end_byte = if next_i < chars.len() {
                        chars[next_i].0
                    } else {
                        original.len()
                    };
                    tokens.push((start_byte, original[start_byte..end_byte].to_string()));
                }
                i = next_i;
                continue;
            }
        }

        if let Some(next_i) = skip_dollar_quoted_string(original, chars, i, byte_pos, ch) {
            i = next_i;
            continue;
        }

        update_parentheses_depth(ch, &mut depth);

        if depth == 0 && ch == ',' {
            tokens.push((byte_pos, ",".to_string()));
            i += 1;
            continue;
        }

        if depth == 0 && (ch.is_alphanumeric() || ch == '_') && is_word_start(chars, i) {
            let start_byte = byte_pos;
            let mut end_i = i;
            while end_i + 1 < chars.len()
                && (chars[end_i + 1].1.is_alphanumeric() || chars[end_i + 1].1 == '_')
            {
                end_i += 1;
            }
            let end_byte = if end_i + 1 < chars.len() {
                chars[end_i + 1].0
            } else {
                original.len()
            };
            tokens.push((start_byte, original[start_byte..end_byte].to_string()));
            i = end_i + 1;
            continue;
        }

        if depth == 0 && ch == '.' && !tokens.is_empty() {
            let prev = tokens.last_mut().unwrap();
            prev.1.push('.');
            i += 1;
            if i < chars.len() {
                let (next_byte, next_ch) = chars[i];
                if next_ch == '"' {
                    if let Some(next_i) = skip_double_quoted_identifier(chars, i, next_ch) {
                        let end_byte = if next_i < chars.len() {
                            chars[next_i].0
                        } else {
                            original.len()
                        };
                        prev.1.push_str(&original[next_byte..end_byte]);
                        i = next_i;
                        continue;
                    }
                } else if next_ch.is_alphanumeric() || next_ch == '_' {
                    let mut end_i = i;
                    while end_i + 1 < chars.len()
                        && (chars[end_i + 1].1.is_alphanumeric() || chars[end_i + 1].1 == '_')
                    {
                        end_i += 1;
                    }
                    let end_byte = if end_i + 1 < chars.len() {
                        chars[end_i + 1].0
                    } else {
                        original.len()
                    };
                    prev.1.push_str(&original[next_byte..end_byte]);
                    i = end_i + 1;
                    continue;
                }
            }
            continue;
        }

        i += 1;
    }

    tokens
}

fn unquote_simple(name: &str) -> String {
    if let Some(dot_pos) = find_unquoted_dot(name) {
        let schema = &name[..dot_pos];
        let table = &name[dot_pos + 1..];
        let schema_unquoted = unquote_single_ident(schema);
        let table_unquoted = unquote_single_ident(table);
        return format!("{schema_unquoted}.{table_unquoted}");
    }
    unquote_single_ident(name)
}

fn find_unquoted_dot(name: &str) -> Option<usize> {
    let mut in_quote = false;
    for (i, ch) in name.char_indices() {
        if ch == '"' {
            in_quote = !in_quote;
        } else if ch == '.' && !in_quote {
            return Some(i);
        }
    }
    None
}

fn unquote_single_ident(s: &str) -> String {
    if s.starts_with('"') && s.ends_with('"') && s.len() >= 2 {
        s[1..s.len() - 1].replace("\"\"", "\"")
    } else {
        s.to_string()
    }
}

const KNOWN_DROP_TYPES: &[&str] = &[
    "table",
    "index",
    "policy",
    "schema",
    "view",
    "type",
    "sequence",
    "trigger",
    "extension",
    "rule",
    "domain",
    "function",
    "procedure",
    "database",
    "role",
    "tablespace",
    "collation",
    "conversion",
    "server",
    "publication",
    "subscription",
    "statistics",
    "aggregate",
];

fn extract_drop_object_name(original: &str, chars: &[(usize, char)]) -> Option<String> {
    let tokens = collect_top_level_tokens(original, chars);
    let lowers: Vec<String> = tokens.iter().map(|(_, t)| t.to_lowercase()).collect();

    let drop_idx = lowers.iter().position(|t| t == "drop")?;
    let type_kw = lowers.get(drop_idx + 1)?;

    let second_kw = lowers.get(drop_idx + 2).map(String::as_str);
    let is_two_word_type = (type_kw == "materialized" && second_kw == Some("view"))
        || (type_kw == "foreign" && second_kw == Some("table"));
    let mut name_idx = if is_two_word_type {
        drop_idx + 3
    } else if KNOWN_DROP_TYPES.contains(&type_kw.as_str()) {
        drop_idx + 2
    } else {
        return None;
    };

    // Skip CONCURRENTLY (DROP INDEX CONCURRENTLY)
    if lowers.get(name_idx).map(String::as_str) == Some("concurrently") {
        name_idx += 1;
    }

    // Skip IF EXISTS
    if lowers.get(name_idx).map(String::as_str) == Some("if")
        && lowers.get(name_idx + 1).map(String::as_str) == Some("exists")
    {
        name_idx += 2;
    }

    let raw = tokens.get(name_idx).map(|(_, t)| t.as_str())?;
    if tokens.get(name_idx + 1).map(|(_, t)| t.as_str()) == Some(",") {
        return None;
    }
    Some(unquote_simple(raw))
}

fn extract_truncate_table_name(original: &str, chars: &[(usize, char)]) -> Option<String> {
    let tokens = collect_top_level_tokens(original, chars);
    let lowers: Vec<String> = tokens.iter().map(|(_, t)| t.to_lowercase()).collect();

    let trunc_idx = lowers.iter().position(|t| t == "truncate")?;
    let mut name_idx = trunc_idx + 1;

    if lowers.get(name_idx).map(String::as_str) == Some("table") {
        name_idx += 1;
    }

    if lowers.get(name_idx).map(String::as_str) == Some("only") {
        name_idx += 1;
    }

    let raw = tokens.get(name_idx).map(|(_, t)| t.as_str())?;
    if tokens.get(name_idx + 1).map(|(_, t)| t.as_str()) == Some(",") {
        return None;
    }
    Some(unquote_simple(raw))
}

fn extract_delete_table_name(original: &str, chars: &[(usize, char)]) -> Option<String> {
    let tokens = collect_top_level_tokens(original, chars);
    let lowers: Vec<String> = tokens.iter().map(|(_, t)| t.to_lowercase()).collect();

    let delete_idx = lowers.iter().position(|t| t == "delete")?;

    if lowers.get(delete_idx + 1).map(String::as_str) != Some("from") {
        return None;
    }

    let mut name_idx = delete_idx + 2;

    if lowers.get(name_idx).map(String::as_str) == Some("only") {
        name_idx += 1;
    }

    let raw = tokens.get(name_idx).map(|(_, t)| t.as_str())?;
    Some(unquote_simple(raw))
}

fn extract_update_table_name(original: &str, chars: &[(usize, char)]) -> Option<String> {
    let tokens = collect_top_level_tokens(original, chars);
    let lowers: Vec<String> = tokens.iter().map(|(_, t)| t.to_lowercase()).collect();

    let update_idx = lowers.iter().position(|t| t == "update")?;

    let mut name_idx = update_idx + 1;

    if lowers.get(name_idx).map(String::as_str) == Some("only") {
        name_idx += 1;
    }

    let raw = tokens.get(name_idx).map(|(_, t)| t.as_str())?;
    Some(unquote_simple(raw))
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;

    mod classify_tests {
        use super::*;

        #[rstest]
        #[case::plain_select("SELECT * FROM users", StatementKind::Select)]
        #[case::lowercase_select("select id from users", StatementKind::Select)]
        #[case::cte_select("WITH cte AS (SELECT 1) SELECT * FROM cte", StatementKind::Select)]
        #[case::recursive_cte_select(
            "WITH RECURSIVE tree AS (SELECT 1) SELECT * FROM tree",
            StatementKind::Select
        )]
        #[case::multiple_ctes_select(
            "WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b",
            StatementKind::Select
        )]
        #[case::explain_select("EXPLAIN SELECT * FROM users", StatementKind::Select)]
        #[case::explain_analyze_select(
            "EXPLAIN ANALYZE SELECT * FROM users",
            StatementKind::Select
        )]
        #[case::explain_verbose_select(
            "EXPLAIN VERBOSE SELECT * FROM users",
            StatementKind::Select
        )]
        #[case::explain_costs_off("EXPLAIN COSTS OFF SELECT * FROM users", StatementKind::Select)]
        #[case::explain_update("EXPLAIN UPDATE users SET x = 1", StatementKind::Select)]
        #[case::explain_delete("EXPLAIN DELETE FROM users", StatementKind::Select)]
        #[case::explain_merge(
            "EXPLAIN MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE SET x = 1",
            StatementKind::Select
        )]
        #[case::explain_paren_analyze_select(
            "EXPLAIN (ANALYZE) SELECT * FROM users",
            StatementKind::Select
        )]
        #[case::show("SHOW search_path", StatementKind::Select)]
        #[case::show_all("SHOW ALL", StatementKind::Select)]
        fn select_variants(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::explain_analyze_update(
            "EXPLAIN ANALYZE UPDATE users SET name = 'x'",
            StatementKind::Update { has_where: false }
        )]
        #[case::explain_analyze_update_where(
            "EXPLAIN ANALYZE UPDATE users SET name = 'x' WHERE id = 1",
            StatementKind::Update { has_where: true }
        )]
        #[case::explain_analyze_delete(
            "EXPLAIN ANALYZE DELETE FROM users",
            StatementKind::Delete { has_where: false }
        )]
        #[case::explain_analyze_insert(
            "EXPLAIN ANALYZE INSERT INTO users VALUES (1)",
            StatementKind::Insert
        )]
        #[case::explain_paren_analyze_update(
            "EXPLAIN (ANALYZE) UPDATE users SET x = 1",
            StatementKind::Update { has_where: false }
        )]
        #[case::explain_paren_analyze_verbose_delete(
            "EXPLAIN (ANALYZE, VERBOSE) DELETE FROM users",
            StatementKind::Delete { has_where: false }
        )]
        fn explain_analyze_dml(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::cte_update(
            "WITH cte AS (SELECT 1) UPDATE users SET name = 'x'",
            StatementKind::Update { has_where: false }
        )]
        #[case::cte_delete(
            "WITH cte AS (SELECT 1) DELETE FROM users",
            StatementKind::Delete { has_where: false }
        )]
        #[case::cte_delete_where(
            "WITH cte AS (SELECT 1) DELETE FROM users WHERE id = 1",
            StatementKind::Delete { has_where: true }
        )]
        fn cte_dml(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::plain_insert("INSERT INTO users VALUES (1)", StatementKind::Insert)]
        #[case::update_no_where("UPDATE users SET name = 'x'", StatementKind::Update { has_where: false })]
        #[case::update_with_where("UPDATE users SET name = 'x' WHERE id = 1", StatementKind::Update { has_where: true })]
        #[case::delete_no_where("DELETE FROM users", StatementKind::Delete { has_where: false })]
        #[case::delete_with_where("DELETE FROM users WHERE id = 1", StatementKind::Delete { has_where: true })]
        fn dml(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::create_table("CREATE TABLE foo (id INT)", StatementKind::Create)]
        #[case::alter_table("ALTER TABLE users ADD COLUMN foo INT", StatementKind::Alter)]
        #[case::drop_table("DROP TABLE users", StatementKind::Drop)]
        #[case::truncate("TRUNCATE users", StatementKind::Truncate)]
        fn ddl(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::begin("BEGIN", StatementKind::Transaction)]
        #[case::commit("COMMIT", StatementKind::Transaction)]
        #[case::rollback("ROLLBACK", StatementKind::Transaction)]
        #[case::savepoint("SAVEPOINT sp1", StatementKind::Transaction)]
        #[case::start_transaction("START TRANSACTION", StatementKind::Transaction)]
        #[case::rollback_to_savepoint("ROLLBACK TO SAVEPOINT sp1", StatementKind::Transaction)]
        #[case::release_savepoint("RELEASE SAVEPOINT sp1", StatementKind::Transaction)]
        fn transaction(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::grant("GRANT SELECT ON users TO role1", StatementKind::Unsupported)]
        #[case::revoke("REVOKE SELECT ON users FROM role1", StatementKind::Unsupported)]
        #[case::copy("COPY users FROM '/tmp/data.csv'", StatementKind::Unsupported)]
        #[case::do_block("DO $$ BEGIN RAISE NOTICE 'hi'; END $$", StatementKind::Unsupported)]
        #[case::call("CALL my_procedure()", StatementKind::Unsupported)]
        #[case::merge(
            "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE SET x = 1",
            StatementKind::Unsupported
        )]
        fn unsupported(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::empty("", StatementKind::Other)]
        #[case::whitespace_only("   ", StatementKind::Other)]
        #[case::comment_only("-- just a comment", StatementKind::Other)]
        #[case::block_comment_only("/* nothing */", StatementKind::Other)]
        fn other(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }

        #[rstest]
        #[case::string_literal_keyword(
            "SELECT * FROM t WHERE action = 'delete'",
            StatementKind::Select
        )]
        #[case::double_quoted_keyword("SELECT \"update\" FROM t", StatementKind::Select)]
        #[case::dollar_quoted_keyword("SELECT $$delete$$ AS label", StatementKind::Select)]
        #[case::mixed_case("SeLeCt * FROM users", StatementKind::Select)]
        #[case::leading_comment("-- comment\nSELECT * FROM t", StatementKind::Select)]
        #[case::leading_block_comment("/* comment */ SELECT * FROM t", StatementKind::Select)]
        #[case::multiple_statements("SELECT 1; DELETE FROM users", StatementKind::Other)]
        #[case::trailing_semicolon("SELECT 1;", StatementKind::Select)]
        #[case::trailing_semicolon_whitespace("SELECT 1;   ", StatementKind::Select)]
        #[case::where_in_subquery(
            "DELETE FROM users WHERE id IN (SELECT id FROM old_users)",
            StatementKind::Delete { has_where: true }
        )]
        #[case::where_in_string_not_real(
            "UPDATE users SET note = 'where is my cat'",
            StatementKind::Update { has_where: false }
        )]
        #[case::where_in_dollar_quote(
            "UPDATE users SET note = $$where$$ WHERE id = 1",
            StatementKind::Update { has_where: true }
        )]
        #[case::tagged_dollar_quote_keyword(
            "SELECT $tag$delete from here$tag$ AS s",
            StatementKind::Select
        )]
        #[case::semicolon_inside_string("SELECT * FROM t WHERE x = ';'", StatementKind::Select)]
        #[case::semicolon_in_dollar_quote("SELECT $$semi;colon$$ AS label", StatementKind::Select)]
        #[case::non_ascii_identifier("SELECT * FROM \"ユーザー\"", StatementKind::Select)]
        #[case::non_ascii_literal(
            "SELECT name FROM users WHERE name = '日本語'",
            StatementKind::Select
        )]
        #[case::unterminated_dollar_quote("SELECT $$unclosed", StatementKind::Select)]
        #[case::nested_block_comment(
            "SELECT /* outer /* inner */ still comment */ 1",
            StatementKind::Select
        )]
        #[case::identifier_contains_keyword("SELECT delete_flag FROM t", StatementKind::Select)]
        #[case::table_name_contains_keyword("SELECT * FROM users_to_delete", StatementKind::Select)]
        #[case::double_quoted_escaped("SELECT \"up\"\"date\" FROM t", StatementKind::Select)]
        #[case::cte_with_update_in_subquery(
            "WITH x AS (UPDATE users SET name='a' RETURNING *) SELECT * FROM x",
            StatementKind::Select
        )]
        #[case::select_with_parenthesized_expr(
            "WITH cte AS (SELECT 1) SELECT (1+2)",
            StatementKind::Select
        )]
        #[case::delete_then_select("DELETE FROM users; SELECT 1", StatementKind::Other)]
        #[case::select_then_select("SELECT 1; SELECT 2", StatementKind::Other)]
        #[case::select_into("SELECT * INTO backup FROM users", StatementKind::Other)]
        #[case::select_into_columns("SELECT id, name INTO backup FROM users", StatementKind::Other)]
        fn edge_cases(#[case] sql: &str, #[case] expected: StatementKind) {
            assert_eq!(classify(sql), expected);
        }
    }

    mod first_keyword_tests {
        use super::*;

        #[rstest]
        #[case::grant("GRANT SELECT ON users TO role1", Some("GRANT"))]
        #[case::lowercase("do $$ BEGIN RAISE NOTICE 'hi'; END $$", Some("DO"))]
        #[case::leading_comment("-- note\nCOPY users FROM '/tmp/x'", Some("COPY"))]
        #[case::leading_whitespace("  call my_procedure()", Some("CALL"))]
        #[case::symbols_only("???", None)]
        #[case::empty("", None)]
        fn returns_first_token_uppercased(#[case] sql: &str, #[case] expected: Option<&str>) {
            assert_eq!(first_keyword(sql), expected.map(ToString::to_string));
        }
    }

    mod extract_target_name_tests {
        use super::*;

        #[rstest]
        #[case::simple("DROP TABLE users", StatementKind::Drop, Some("users"))]
        #[case::if_exists("DROP TABLE IF EXISTS users", StatementKind::Drop, Some("users"))]
        #[case::schema_qualified(
            "DROP TABLE IF EXISTS public.users",
            StatementKind::Drop,
            Some("public.users")
        )]
        #[case::quoted("DROP TABLE \"Users\"", StatementKind::Drop, Some("Users"))]
        #[case::quoted_schema(
            "DROP TABLE \"public\".\"Users\"",
            StatementKind::Drop,
            Some("public.Users")
        )]
        fn drop_table(
            #[case] sql: &str,
            #[case] kind: StatementKind,
            #[case] expected: Option<&str>,
        ) {
            assert_eq!(
                extract_target_name(sql, &kind),
                expected.map(ToString::to_string)
            );
        }

        #[rstest]
        #[case::simple("TRUNCATE users", StatementKind::Truncate, Some("users"))]
        #[case::with_table_keyword("TRUNCATE TABLE users", StatementKind::Truncate, Some("users"))]
        #[case::multiple("TRUNCATE a, b", StatementKind::Truncate, None)]
        fn truncate(
            #[case] sql: &str,
            #[case] kind: StatementKind,
            #[case] expected: Option<&str>,
        ) {
            assert_eq!(
                extract_target_name(sql, &kind),
                expected.map(ToString::to_string)
            );
        }

        #[rstest]
        #[case::simple(
            "DELETE FROM users",
            StatementKind::Delete { has_where: false },
            Some("users")
        )]
        #[case::with_only(
            "DELETE FROM ONLY users",
            StatementKind::Delete { has_where: false },
            Some("users")
        )]
        #[case::schema_qualified(
            "DELETE FROM public.users",
            StatementKind::Delete { has_where: false },
            Some("public.users")
        )]
        fn delete(#[case] sql: &str, #[case] kind: StatementKind, #[case] expected: Option<&str>) {
            assert_eq!(
                extract_target_name(sql, &kind),
                expected.map(ToString::to_string)
            );
        }

        #[rstest]
        #[case::simple(
            "UPDATE users SET name = 'x'",
            StatementKind::Update { has_where: false },
            Some("users")
        )]
        #[case::with_only(
            "UPDATE ONLY users SET name = 'x'",
            StatementKind::Update { has_where: false },
            Some("users")
        )]
        #[case::schema_qualified(
            "UPDATE public.users SET name = 'x'",
            StatementKind::Update { has_where: false },
            Some("public.users")
        )]
        fn update(#[case] sql: &str, #[case] kind: StatementKind, #[case] expected: Option<&str>) {
            assert_eq!(
                extract_target_name(sql, &kind),
                expected.map(ToString::to_string)
            );
        }

        #[rstest]
        #[case::select("SELECT * FROM users", StatementKind::Select)]
        #[case::insert("INSERT INTO users VALUES (1)", StatementKind::Insert)]
        #[case::create("CREATE TABLE users (id INT)", StatementKind::Create)]
        #[case::alter("ALTER TABLE users ADD COLUMN x INT", StatementKind::Alter)]
        #[case::transaction("BEGIN", StatementKind::Transaction)]
        #[case::other("GRANT SELECT ON users TO role1", StatementKind::Other)]
        fn not_applicable(#[case] sql: &str, #[case] kind: StatementKind) {
            assert_eq!(extract_target_name(sql, &kind), None);
        }

        #[rstest]
        #[case::truncate_schema(
            "TRUNCATE TABLE public.events",
            StatementKind::Truncate,
            Some("public.events")
        )]
        #[case::truncate_only(
            "TRUNCATE TABLE ONLY partitioned_table",
            StatementKind::Truncate,
            Some("partitioned_table")
        )]
        #[case::delete_quoted(
            "DELETE FROM \"MyTable\"",
            StatementKind::Delete { has_where: false },
            Some("MyTable")
        )]
        #[case::update_quoted_schema(
            "UPDATE \"public\".\"MyTable\" SET x = 1",
            StatementKind::Update { has_where: false },
            Some("public.MyTable")
        )]
        #[case::drop_index("DROP INDEX my_index", StatementKind::Drop, Some("my_index"))]
        #[case::drop_index_if_exists(
            "DROP INDEX IF EXISTS my_index",
            StatementKind::Drop,
            Some("my_index")
        )]
        #[case::drop_index_concurrently(
            "DROP INDEX CONCURRENTLY my_index",
            StatementKind::Drop,
            Some("my_index")
        )]
        #[case::drop_index_concurrently_if_exists(
            "DROP INDEX CONCURRENTLY IF EXISTS my_index",
            StatementKind::Drop,
            Some("my_index")
        )]
        #[case::drop_policy(
            "DROP POLICY tenant_isolation ON projects",
            StatementKind::Drop,
            Some("tenant_isolation")
        )]
        #[case::drop_policy_if_exists(
            "DROP POLICY IF EXISTS tenant_isolation ON projects",
            StatementKind::Drop,
            Some("tenant_isolation")
        )]
        #[case::drop_schema("DROP SCHEMA my_schema", StatementKind::Drop, Some("my_schema"))]
        #[case::drop_view("DROP VIEW my_view", StatementKind::Drop, Some("my_view"))]
        #[case::drop_materialized_view(
            "DROP MATERIALIZED VIEW my_mv",
            StatementKind::Drop,
            Some("my_mv")
        )]
        #[case::drop_type("DROP TYPE my_type", StatementKind::Drop, Some("my_type"))]
        #[case::drop_function(
            "DROP FUNCTION my_func(int, text)",
            StatementKind::Drop,
            Some("my_func")
        )]
        #[case::drop_sequence("DROP SEQUENCE my_seq", StatementKind::Drop, Some("my_seq"))]
        #[case::drop_trigger(
            "DROP TRIGGER my_trigger ON my_table",
            StatementKind::Drop,
            Some("my_trigger")
        )]
        #[case::drop_extension(
            "DROP EXTENSION IF EXISTS pgcrypto",
            StatementKind::Drop,
            Some("pgcrypto")
        )]
        #[case::drop_rule("DROP RULE my_rule ON my_table", StatementKind::Drop, Some("my_rule"))]
        #[case::drop_domain("DROP DOMAIN my_domain", StatementKind::Drop, Some("my_domain"))]
        #[case::drop_schema_qualified(
            "DROP INDEX public.my_index",
            StatementKind::Drop,
            Some("public.my_index")
        )]
        #[case::drop_database("DROP DATABASE production", StatementKind::Drop, Some("production"))]
        #[case::drop_role("DROP ROLE app_user", StatementKind::Drop, Some("app_user"))]
        #[case::drop_tablespace("DROP TABLESPACE fastdisk", StatementKind::Drop, Some("fastdisk"))]
        #[case::drop_publication("DROP PUBLICATION my_pub", StatementKind::Drop, Some("my_pub"))]
        #[case::drop_foreign_table(
            "DROP FOREIGN TABLE remote_users",
            StatementKind::Drop,
            Some("remote_users")
        )]
        #[case::drop_foreign_table_if_exists(
            "DROP FOREIGN TABLE IF EXISTS remote_users",
            StatementKind::Drop,
            Some("remote_users")
        )]
        #[case::drop_aggregate("DROP AGGREGATE my_agg(int)", StatementKind::Drop, Some("my_agg"))]
        #[case::drop_procedure("DROP PROCEDURE my_proc(int)", StatementKind::Drop, Some("my_proc"))]
        #[case::drop_collation(
            "DROP COLLATION my_collation",
            StatementKind::Drop,
            Some("my_collation")
        )]
        #[case::drop_conversion("DROP CONVERSION my_conv", StatementKind::Drop, Some("my_conv"))]
        #[case::drop_server("DROP SERVER my_server", StatementKind::Drop, Some("my_server"))]
        #[case::drop_subscription("DROP SUBSCRIPTION my_sub", StatementKind::Drop, Some("my_sub"))]
        #[case::drop_statistics("DROP STATISTICS my_stats", StatementKind::Drop, Some("my_stats"))]
        #[case::drop_function_multiple("DROP FUNCTION f(int), g(text)", StatementKind::Drop, None)]
        #[case::drop_procedure_multiple("DROP PROCEDURE p(int), q(int)", StatementKind::Drop, None)]
        #[case::drop_materialized_invalid("DROP MATERIALIZED foo", StatementKind::Drop, None)]
        #[case::drop_foreign_invalid("DROP FOREIGN foo", StatementKind::Drop, None)]
        #[case::drop_owned_by("DROP OWNED BY my_role", StatementKind::Drop, None)]
        #[case::drop_cast("DROP CAST (int AS text)", StatementKind::Drop, None)]
        #[case::drop_multiple("DROP TABLE a, b", StatementKind::Drop, None)]
        #[case::drop_multiple_index("DROP INDEX a, b", StatementKind::Drop, None)]
        #[case::truncate_multiple("TRUNCATE a, b, c", StatementKind::Truncate, None)]
        fn additional_extraction(
            #[case] sql: &str,
            #[case] kind: StatementKind,
            #[case] expected: Option<&str>,
        ) {
            assert_eq!(
                extract_target_name(sql, &kind),
                expected.map(ToString::to_string)
            );
        }
    }
}
