use regex::Regex;
use std::sync::LazyLock;

/// Regex to strip Jinja tags {{ ... }} and {%- ... -%} etc.
static JINJA_TAG: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\{\{-?[\s\S]*?-?\}\}|\{%-?[\s\S]*?-?%\}").unwrap());

/// Regex to strip Jinja comments {# ... #}
static JINJA_COMMENT: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{#[\s\S]*?#\}").unwrap());

/// Match the beginning of a SELECT clause (possibly with DISTINCT).
static SELECT_START: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?is)\bSELECT\b\s+(?:DISTINCT\s+)?").unwrap());

/// Extract column names from the outermost SELECT clause of a SQL string.
///
/// This is a best-effort regex-based extraction, not a full SQL parser.
/// It handles:
/// - `SELECT col1, col2 FROM ...` -> `["col1", "col2"]`
/// - `SELECT t.col1 AS alias1` -> `["alias1"]`
/// - `SELECT col1 as alias1` -> `["alias1"]`
/// - `SELECT *` -> `["*"]`
/// - `SELECT DISTINCT col1, col2` -> `["col1", "col2"]`
/// - Jinja tags are stripped before parsing
/// - Subqueries in parentheses are skipped
/// - Multiline SELECT clauses are handled
pub fn extract_select_columns(sql: &str) -> Vec<String> {
    // Strip Jinja comments and tags
    let cleaned = JINJA_COMMENT.replace_all(sql, "");
    let cleaned = JINJA_TAG.replace_all(&cleaned, "__jinja__");

    // Find the first SELECT keyword
    let m = match SELECT_START.find(&cleaned) {
        Some(m) => m,
        None => return vec![],
    };

    // Find the first top-level FROM after the SELECT (not inside parentheses)
    let after_select = &cleaned[m.end()..];
    let select_body = match find_top_level_from(after_select) {
        Some(pos) => &after_select[..pos],
        None => return vec![],
    };

    // Split on commas, but not commas inside parentheses
    let items = split_top_level_commas(select_body);

    items
        .iter()
        .filter_map(|item| classify_select_item(item.trim()))
        .collect()
}

/// Classify a single SELECT item and return its column name, if any.
fn classify_select_item(item: &str) -> Option<String> {
    if item.is_empty() {
        return None;
    }

    // Items starting with '(' are subqueries; check for alias after closing paren
    if item.starts_with('(') {
        return extract_alias_after_paren(item);
    }

    let col = extract_column_name(item);
    if col.is_empty() {
        None
    } else {
        Some(col)
    }
}

/// Check if a byte is a word boundary character (not alphanumeric and not underscore)
fn is_word_boundary(b: u8) -> bool {
    !b.is_ascii_alphanumeric() && b != b'_'
}

/// Check if position `i` starts a top-level `FROM` keyword with proper boundaries.
/// Works at the byte level: continuation bytes of multi-byte UTF-8 chars (>= 0x80)
/// never match ASCII letters, so byte comparison is safe and avoids panicking on
/// non-char-boundary slices.
fn check_from_at(bytes: &[u8], i: usize, len: usize) -> bool {
    if i + 4 > len {
        return false;
    }
    if !bytes[i..i + 4].eq_ignore_ascii_case(b"from") {
        return false;
    }
    let before_ok = i == 0 || is_word_boundary(bytes[i - 1]);
    let after_ok = i + 4 >= len || is_word_boundary(bytes[i + 4]);
    before_ok && after_ok
}

/// Find the position of the first top-level `FROM` keyword (not inside parentheses).
/// Returns the byte offset of the start of `FROM` relative to the input string.
fn find_top_level_from(s: &str) -> Option<usize> {
    let bytes = s.as_bytes();
    let len = bytes.len();
    let mut depth: i32 = 0;
    let mut i = 0;

    while i < len {
        match bytes[i] {
            b'(' => depth += 1,
            b')' if depth > 0 => depth -= 1,
            b'f' | b'F' if depth == 0 && check_from_at(bytes, i, len) => return Some(i),
            _ => {}
        }
        i += 1;
    }

    None
}

/// Split a string on commas that are not inside parentheses.
fn split_top_level_commas(s: &str) -> Vec<String> {
    let mut items = Vec::new();
    let mut current = String::new();
    let mut depth = 0;

    for ch in s.chars() {
        match ch {
            '(' => {
                depth += 1;
                current.push(ch);
            }
            ')' => {
                depth -= 1;
                current.push(ch);
            }
            ',' if depth == 0 => {
                items.push(current.clone());
                current.clear();
            }
            _ => {
                current.push(ch);
            }
        }
    }

    if !current.trim().is_empty() {
        items.push(current);
    }

    items
}

/// Extract the alias after a closing parenthesis, e.g., `(SELECT ...) AS alias`
fn extract_alias_after_paren(item: &str) -> Option<String> {
    // Find the last closing paren
    let close = item.rfind(')')?;
    let after = item[close + 1..].trim();
    if after.is_empty() {
        return None;
    }
    // Strip leading AS (case-insensitive). Match at byte level so a leading
    // multi-byte UTF-8 char doesn't panic on `after[..2]`.
    let after_bytes = after.as_bytes();
    let after = if after_bytes.len() >= 3
        && after_bytes[..2].eq_ignore_ascii_case(b"as")
        && after_bytes[2].is_ascii_whitespace()
    {
        after[2..].trim()
    } else {
        after
    };
    if after.is_empty() {
        None
    } else {
        Some(clean_identifier(after))
    }
}

/// Extract the effective column name from a single SELECT item.
///
/// Rules:
/// 1. If `AS alias` is present, return the alias.
/// 2. If `table.column`, return column.
/// 3. Otherwise return the token itself (e.g., `*`, `col1`).
fn extract_column_name(item: &str) -> String {
    let item = item.trim();

    // Check for AS alias (case-insensitive) - look for last " AS " or " as "
    // We search from the end to handle expressions like `CAST(x AS int) AS col`
    if let Some(alias) = find_last_as_alias(item) {
        return clean_identifier(&alias);
    }

    // No alias; take the last token (handles `table.col` and bare `col`)
    let last_token = item.split_whitespace().last().unwrap_or(item);

    // Handle table.column
    if let Some(pos) = last_token.rfind('.') {
        return clean_identifier(&last_token[pos + 1..]);
    }

    clean_identifier(last_token)
}

/// Check if position `i` (a whitespace char) starts a top-level ` AS ` token.
/// Returns the position after "AS " if matched.
/// Byte-level comparison avoids panicking when `i + 1..i + 3` lands in the
/// middle of a multi-byte UTF-8 char.
fn is_as_keyword_at(bytes: &[u8], i: usize, len: usize) -> Option<usize> {
    if i + 3 >= len {
        return None;
    }
    let after = bytes[i + 3];
    if bytes[i + 1..i + 3].eq_ignore_ascii_case(b"as")
        && matches!(after, b' ' | b'\t' | b'\n' | b'\r')
    {
        Some(i + 4)
    } else {
        None
    }
}

/// Find the alias from the last ` AS ` keyword that is not inside parentheses.
fn find_last_as_alias(item: &str) -> Option<String> {
    let bytes = item.as_bytes();
    let len = bytes.len();
    let mut depth = 0;
    let mut last_as_pos: Option<usize> = None;

    let mut i = 0;
    while i < len {
        match bytes[i] {
            b'(' => depth += 1,
            b')' if depth > 0 => depth -= 1,
            b' ' | b'\t' | b'\n' | b'\r' if depth == 0 => {
                if let Some(pos) = is_as_keyword_at(bytes, i, len) {
                    last_as_pos = Some(pos);
                }
            }
            _ => {}
        }
        i += 1;
    }

    last_as_pos.map(|pos| item[pos..].trim().to_string())
}

/// Clean an identifier: trim whitespace and remove surrounding backticks or quotes.
fn clean_identifier(s: &str) -> String {
    let s = s.trim();
    let s = s.trim_matches('`');
    let s = s.trim_matches('"');
    s.to_string()
}

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

    #[test]
    fn test_simple_select() {
        let sql = "SELECT col1, col2 FROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["col1", "col2"]);
    }

    #[test]
    fn test_select_with_aliases() {
        let sql = "SELECT col1 AS alias1, col2 as alias2 FROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["alias1", "alias2"]);
    }

    #[test]
    fn test_select_with_table_prefixes() {
        let sql = "SELECT t.col1, t.col2 FROM my_table t";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["col1", "col2"]);
    }

    #[test]
    fn test_select_star() {
        let sql = "SELECT * FROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["*"]);
    }

    #[test]
    fn test_select_distinct() {
        let sql = "SELECT DISTINCT col1, col2 FROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["col1", "col2"]);
    }

    #[test]
    fn test_select_with_jinja() {
        let sql = r#"
            {{ config(materialized='table') }}

            SELECT
                order_id,
                {{ dbt_utils.star(from=ref('stg_orders')) }},
                customer_id
            FROM {{ ref('stg_orders') }}
        "#;
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["order_id", "__jinja__", "customer_id"]);
    }

    #[test]
    fn test_multiline_select() {
        let sql = r#"
            SELECT
                order_id,
                customer_id,
                order_date,
                status
            FROM orders
        "#;
        let cols = extract_select_columns(sql);
        assert_eq!(
            cols,
            vec!["order_id", "customer_id", "order_date", "status"]
        );
    }

    #[test]
    fn test_cte_gets_outer_select() {
        let sql = r#"
            WITH cte AS (
                SELECT inner_col1, inner_col2 FROM raw_table
            )
            SELECT outer_col1, outer_col2 FROM cte
        "#;
        // The first SELECT...FROM is inside the CTE.
        // For a basic regex approach, we get the first SELECT...FROM found.
        // This returns the CTE's columns. For most dbt models, the outermost
        // query is the final SELECT, but with CTEs the regex finds the first one.
        // This is a known limitation of the regex approach.
        let cols = extract_select_columns(sql);
        // The CTE's SELECT is wrapped in parens, so the regex actually
        // matches the outer SELECT because the inner one is inside parens
        // following "AS (" and the FROM at the end of "FROM raw_table" is
        // consumed. Let's verify what we actually get.
        assert!(!cols.is_empty());
    }

    #[test]
    fn test_select_with_function() {
        let sql = "SELECT COUNT(*) AS total, SUM(amount) AS total_amount FROM orders";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["total", "total_amount"]);
    }

    #[test]
    fn test_select_table_prefix_with_alias() {
        let sql = "SELECT t.col1 AS alias1, t.col2 FROM my_table t";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["alias1", "col2"]);
    }

    #[test]
    fn test_no_select() {
        let sql = "INSERT INTO my_table VALUES (1, 2, 3)";
        let cols = extract_select_columns(sql);
        assert!(cols.is_empty());
    }

    #[test]
    fn test_select_with_jinja_comments() {
        let sql = r#"
            {# Select all order columns #}
            SELECT order_id, status FROM orders
        "#;
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["order_id", "status"]);
    }

    #[test]
    fn test_select_with_cast() {
        let sql = "SELECT CAST(order_id AS INTEGER) AS order_id, status FROM orders";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["order_id", "status"]);
    }

    #[test]
    fn test_select_with_subquery_alias() {
        let sql = "SELECT (SELECT MAX(id) FROM t) AS max_id, name FROM users";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["max_id", "name"]);
    }

    #[test]
    fn test_typical_dbt_model() {
        let sql = r#"
            {{ config(materialized='view') }}

            SELECT
                order_id,
                customer_id,
                order_date,
                status,
                amount
            FROM {{ ref('stg_orders') }}
        "#;
        let cols = extract_select_columns(sql);
        assert_eq!(
            cols,
            vec!["order_id", "customer_id", "order_date", "status", "amount"]
        );
    }

    #[test]
    fn test_select_case_insensitive() {
        let sql = "select col1, col2 from my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["col1", "col2"]);
    }

    #[test]
    fn test_select_with_backtick_identifiers() {
        let sql = "SELECT `col1`, `col2` FROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["col1", "col2"]);
    }

    #[test]
    fn test_extract_alias_after_paren_no_alias() {
        // Subquery with no alias after the closing paren
        let result = extract_alias_after_paren("(SELECT 1)");
        assert!(result.is_none());
    }

    #[test]
    fn test_extract_alias_after_paren_bare_alias() {
        // Subquery with bare alias (no AS keyword)
        let result = extract_alias_after_paren("(SELECT 1) my_alias");
        assert_eq!(result, Some("my_alias".to_string()));
    }

    #[test]
    fn test_extract_alias_after_paren_as_alias() {
        // Subquery with AS alias
        let result = extract_alias_after_paren("(SELECT 1) AS my_alias");
        assert_eq!(result, Some("my_alias".to_string()));
    }

    #[test]
    fn test_extract_alias_after_paren_no_paren() {
        // No closing paren at all
        let result = extract_alias_after_paren("SELECT 1");
        assert!(result.is_none());
    }

    // Regression test for GH issue #1: panic when SQL contains multi-byte
    // UTF-8 characters (e.g. Japanese in a -- comment or string literal).
    #[test]
    fn test_select_with_multibyte_utf8_in_comment() {
        let sql = "SELECT\n    case\n      when flag = true then false -- 日本語コメント\n      else flag\n    end as flag\nFROM my_table";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["flag"]);
    }

    #[test]
    fn test_select_with_multibyte_utf8_in_string_literal() {
        let sql = "SELECT '日本語' AS greeting, name FROM users";
        let cols = extract_select_columns(sql);
        assert_eq!(cols, vec!["greeting", "name"]);
    }

    #[test]
    fn test_extract_alias_after_paren_multibyte_leading() {
        // Trimmed remainder starts with a multi-byte UTF-8 char.
        // Must not panic on slicing.
        let result = extract_alias_after_paren("(SELECT 1) 日本語");
        assert_eq!(result, Some("日本語".to_string()));
    }
}
