use pg_query::protobuf::Token;
use text_size::{TextRange, TextSize};

/// Find the qualified function name and its name-token range after `EXECUTE FUNCTION`.
///
/// Scans `stmt_text` (a single statement slice, not the full file).
/// Returns `(symbol_str, range)` where `symbol_str` is `"schema/name"` and `range`
/// is the absolute byte range of the function name identifier token.
/// Returns `None` on scan failure, missing token sequence, or unexpected ident count.
pub(super) fn function_name_after_execute(
    stmt_text: &str,
    stmt_start: TextSize,
) -> Option<(String, TextRange)> {
    let tokens = match pg_query::scan(stmt_text) {
        Ok(result) => result.tokens,
        Err(e) => {
            tracing::debug!("lex scan failed: {e}");
            return None;
        }
    };
    collect_func_name_after_execute(&tokens, stmt_text, stmt_start)
}

fn collect_func_name_after_execute(
    tokens: &[pg_query::protobuf::ScanToken],
    stmt_text: &str,
    stmt_start: TextSize,
) -> Option<(String, TextRange)> {
    let func_token = Token::Function as i32;
    let open_paren = Token::Ascii40 as i32;
    let dot = Token::Ascii46 as i32;
    let ident = Token::Ident as i32;

    let anchor_pos = tokens.iter().position(|t| t.token == func_token)?;

    let mut parts: Vec<(String, TextRange)> = Vec::new();
    for tok in &tokens[anchor_pos + 1..] {
        tracing::trace!(
            token = tok.token,
            start = tok.start,
            end = tok.end,
            "examining func name token"
        );
        if tok.token == open_paren {
            break;
        }
        if tok.token == dot {
            continue;
        }
        if tok.token == ident {
            let range = scan_token_to_range(tok, stmt_start);
            let raw_start = usize::try_from(tok.start.max(0)).unwrap_or(0);
            let raw_end = usize::try_from(tok.end.max(0)).unwrap_or(0);
            let name = stmt_text.get(raw_start..raw_end).unwrap_or("").to_owned();
            parts.push((name, range));
        }
    }

    match parts.as_slice() {
        [(schema, _), (name, name_range)] => Some((format!("{schema}/{name}"), *name_range)),
        [(name, name_range)] => Some((format!("public/{name}"), *name_range)),
        _ => None,
    }
}

/// Find all Ident token ranges in the first balanced `(…)` group after `anchor`.
///
/// Scans `stmt_text` (a single statement slice, not the full file).
/// Offsets in returned ranges are absolute: `stmt_start` is added to each
/// token position. Returns empty on scan failure, missing anchor, or missing
/// parens.
pub(super) fn idents_in_parens_after(
    stmt_text: &str,
    stmt_start: TextSize,
    anchor: Token,
) -> Vec<TextRange> {
    let tokens = match pg_query::scan(stmt_text) {
        Ok(result) => result.tokens,
        Err(e) => {
            tracing::debug!("lex scan failed: {e}");
            return vec![];
        }
    };
    collect_idents_after_anchor(&tokens, stmt_start, anchor)
}

/// Find all Ident token ranges in the first balanced `(…)` group after `anchor`,
/// using a pre-scanned token slice.
///
/// Use this when the caller already holds the token slice from a prior scan
/// to avoid re-scanning the same statement text. Semantics are identical to
/// [`idents_in_parens_after`].
pub(super) fn idents_in_parens_after_tokens(
    tokens: &[pg_query::protobuf::ScanToken],
    stmt_start: TextSize,
    anchor: Token,
) -> Vec<TextRange> {
    collect_idents_after_anchor(tokens, stmt_start, anchor)
}

fn collect_idents_after_anchor(
    tokens: &[pg_query::protobuf::ScanToken],
    stmt_start: TextSize,
    anchor: Token,
) -> Vec<TextRange> {
    let anchor_i32 = anchor as i32;
    let open_paren = Token::Ascii40 as i32;
    let close_paren = Token::Ascii41 as i32;
    let ident = Token::Ident as i32;

    let anchor_pos = tokens.iter().position(|t| t.token == anchor_i32);
    let Some(anchor_pos) = anchor_pos else {
        return vec![];
    };

    let paren_pos = tokens[anchor_pos + 1..]
        .iter()
        .position(|t| t.token == open_paren);
    let Some(paren_pos) = paren_pos else {
        return vec![];
    };
    let start = anchor_pos + 1 + paren_pos + 1; // first token inside '('

    let mut depth: u32 = 1;
    let mut result = Vec::new();
    for tok in &tokens[start..] {
        tracing::trace!(
            token = tok.token,
            start = tok.start,
            end = tok.end,
            "examining token"
        );
        if tok.token == open_paren {
            depth += 1;
        } else if tok.token == close_paren {
            depth -= 1;
            if depth == 0 {
                break;
            }
        } else if tok.token == ident {
            let range = scan_token_to_range(tok, stmt_start);
            result.push(range);
        }
    }
    result
}

fn scan_token_to_range(tok: &pg_query::protobuf::ScanToken, stmt_start: TextSize) -> TextRange {
    let start = stmt_start + TextSize::from(u32::try_from(tok.start.max(0)).unwrap_or(0));
    let end = stmt_start + TextSize::from(u32::try_from(tok.end.max(0)).unwrap_or(0));
    TextRange::new(start, end)
}

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

    fn text_at(sql: &str, range: TextRange) -> &str {
        &sql[std::ops::Range::<usize>::from(range)]
    }

    #[test]
    fn idents_after_on_finds_index_columns() {
        let sql = "CREATE INDEX i ON t (col_a, col_b);";
        let ranges = idents_in_parens_after(sql, TextSize::from(0u32), Token::On);
        assert_eq!(ranges.len(), 2, "expected 2 column ranges, got: {ranges:?}");
        assert_eq!(text_at(sql, ranges[0]), "col_a");
        assert_eq!(text_at(sql, ranges[1]), "col_b");
    }

    #[test]
    fn idents_after_key_finds_fk_column() {
        let sql = "ALTER TABLE t ADD CONSTRAINT c FOREIGN KEY (user_id) REFERENCES t2(id);";
        let ranges = idents_in_parens_after(sql, TextSize::from(0u32), Token::Key);
        assert_eq!(ranges.len(), 1, "expected 1 fk_attr range, got: {ranges:?}");
        assert_eq!(text_at(sql, ranges[0]), "user_id");
    }

    #[test]
    fn idents_after_references_finds_pk_column() {
        let sql = "ALTER TABLE t ADD CONSTRAINT c FOREIGN KEY (user_id) REFERENCES t2(id);";
        let ranges = idents_in_parens_after(sql, TextSize::from(0u32), Token::References);
        assert_eq!(ranges.len(), 1, "expected 1 pk_attr range, got: {ranges:?}");
        assert_eq!(text_at(sql, ranges[0]), "id");
    }

    #[test]
    fn returns_empty_on_missing_anchor() {
        let sql = "CREATE INDEX i ON t (col_a);";
        let ranges = idents_in_parens_after(sql, TextSize::from(0u32), Token::Truncate);
        assert!(ranges.is_empty(), "expected empty, got: {ranges:?}");
    }

    #[test]
    fn stmt_start_offset_applied() {
        let prefix = "SELECT 1; ";
        let sql = "CREATE INDEX i ON t (col_a);";
        let full = format!("{prefix}{sql}");
        let stmt_start = TextSize::from(prefix.len() as u32);
        let ranges = idents_in_parens_after(sql, stmt_start, Token::On);
        assert_eq!(ranges.len(), 1, "expected 1 range, got: {ranges:?}");
        // The range should point into `full`, not into `sql`.
        let col_text = &full[std::ops::Range::<usize>::from(ranges[0])];
        assert_eq!(col_text, "col_a");
    }

    #[test]
    fn function_name_schema_qualified() {
        let sql = "CREATE TRIGGER documents_updated_at BEFORE UPDATE ON public.documents \
                   FOR EACH ROW EXECUTE FUNCTION public.update_updated_at_column();";
        let result = function_name_after_execute(sql, TextSize::from(0u32));
        let (sym, range) = result.expect("must find function name");
        assert_eq!(sym, "public/update_updated_at_column");
        assert_eq!(text_at(sql, range), "update_updated_at_column");
    }

    #[test]
    fn function_name_unqualified() {
        let sql = "CREATE TRIGGER t AFTER INSERT ON t FOR EACH ROW EXECUTE FUNCTION my_func();";
        let result = function_name_after_execute(sql, TextSize::from(0u32));
        let (sym, _range) = result.expect("must find function name");
        assert_eq!(sym, "public/my_func");
    }

    #[test]
    fn function_name_stmt_start_offset() {
        let prefix = "SELECT 1; ";
        let sql = "CREATE TRIGGER t AFTER INSERT ON t FOR EACH ROW EXECUTE FUNCTION public.fn();";
        let stmt_start = TextSize::from(prefix.len() as u32);
        let result = function_name_after_execute(sql, stmt_start);
        let (_sym, range) = result.expect("must find function name");
        // The range should point into the full document (prefix + sql), not just sql.
        let full = format!("{prefix}{sql}");
        let name_text = &full[std::ops::Range::<usize>::from(range)];
        assert_eq!(name_text, "fn");
    }

    #[test]
    fn function_name_no_execute_returns_none() {
        let sql = "SELECT 1;";
        let result = function_name_after_execute(sql, TextSize::from(0u32));
        assert!(result.is_none(), "no EXECUTE FUNCTION must return None");
    }

    #[test]
    fn token_name_in_index_name_not_matched() {
        // "idx_scope_token" contains "token" but the anchor is ON, so we only
        // look inside the paren list that follows ON — not inside the index name.
        let sql = "CREATE INDEX idx_scope_token ON t (token);";
        let ranges = idents_in_parens_after(sql, TextSize::from(0u32), Token::On);
        assert_eq!(ranges.len(), 1, "expected exactly 1 range, got: {ranges:?}");
        assert_eq!(text_at(sql, ranges[0]), "token");
        // The range must start after the '(' character.
        let paren_pos = sql.find('(').expect("must have paren") as u32;
        assert!(
            u32::from(ranges[0].start()) > paren_pos,
            "range must be inside parens, not in the index name"
        );
    }

    #[test]
    fn idents_in_parens_after_tokens_fk_and_pk_from_same_scan() {
        let sql = "ALTER TABLE t ADD CONSTRAINT c FOREIGN KEY (user_id) REFERENCES t2(id);";
        let tokens = pg_query::scan(sql).expect("scan must succeed").tokens;
        let stmt_start = TextSize::from(0u32);

        let fk_ranges = idents_in_parens_after_tokens(&tokens, stmt_start, Token::Key);
        assert_eq!(
            fk_ranges.len(),
            1,
            "expected 1 fk range, got: {fk_ranges:?}"
        );
        assert_eq!(text_at(sql, fk_ranges[0]), "user_id");

        let pk_ranges = idents_in_parens_after_tokens(&tokens, stmt_start, Token::References);
        assert_eq!(
            pk_ranges.len(),
            1,
            "expected 1 pk range, got: {pk_ranges:?}"
        );
        assert_eq!(text_at(sql, pk_ranges[0]), "id");

        // Results must agree with the scanning entry point.
        let fk_via_scan = idents_in_parens_after(sql, stmt_start, Token::Key);
        let pk_via_scan = idents_in_parens_after(sql, stmt_start, Token::References);
        assert_eq!(
            fk_ranges, fk_via_scan,
            "pre-scanned fk ranges must match scanned"
        );
        assert_eq!(
            pk_ranges, pk_via_scan,
            "pre-scanned pk ranges must match scanned"
        );
    }
}
