#![allow(dead_code)]

use sqlx::PgPool;

pub const STORAGE_FILES_LEGACY_COMPAT_SQL: &str =
    include_str!("../../../sql/files/storage-files-legacy-compat.sql");
pub const STORAGE_FILES_PROVISION_SQL: &str = include_str!("../../../sql/provision_files.sql");

#[derive(Debug, Clone)]
enum SqlScannerState {
    Normal,
    SingleQuotedString,
    DoubleQuotedIdentifier,
    LineComment,
    BlockComment(usize),
    DollarQuoted(String),
}

pub async fn execute_storage_files_sql(pool: &PgPool) -> Result<usize, sqlx::Error> {
    execute_sql_file(pool, STORAGE_FILES_PROVISION_SQL).await
}

async fn execute_sql_file(pool: &PgPool, sql: &str) -> Result<usize, sqlx::Error> {
    let statements: Vec<&str> = split_sql_statements(sql);
    let total = statements.len();

    for statement in statements {
        sqlx::query(statement).execute(pool).await?;
    }

    Ok(total)
}

fn split_sql_statements(sql: &str) -> Vec<&str> {
    let bytes = sql.as_bytes();
    let mut statements: Vec<&str> = Vec::new();
    let mut state = SqlScannerState::Normal;
    let mut statement_start = 0usize;
    let mut i = 0usize;

    while i < bytes.len() {
        match &mut state {
            SqlScannerState::Normal => {
                if bytes[i] == b'-' && i + 1 < bytes.len() && bytes[i + 1] == b'-' {
                    state = SqlScannerState::LineComment;
                    i += 2;
                    continue;
                }
                if bytes[i] == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
                    state = SqlScannerState::BlockComment(1);
                    i += 2;
                    continue;
                }
                if bytes[i] == b'\'' {
                    state = SqlScannerState::SingleQuotedString;
                    i += 1;
                    continue;
                }
                if bytes[i] == b'"' {
                    state = SqlScannerState::DoubleQuotedIdentifier;
                    i += 1;
                    continue;
                }
                if bytes[i] == b'$'
                    && let Some((tag, len)) = parse_dollar_quote_tag(bytes, i)
                {
                    state = SqlScannerState::DollarQuoted(tag);
                    i += len;
                    continue;
                }
                if bytes[i] == b';' {
                    push_statement(&mut statements, sql, statement_start, i);
                    statement_start = i + 1;
                    i += 1;
                    continue;
                }
                i += 1;
            }
            SqlScannerState::SingleQuotedString => {
                if bytes[i] == b'\'' {
                    if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
                        i += 2;
                    } else {
                        state = SqlScannerState::Normal;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            SqlScannerState::DoubleQuotedIdentifier => {
                if bytes[i] == b'"' {
                    if i + 1 < bytes.len() && bytes[i + 1] == b'"' {
                        i += 2;
                    } else {
                        state = SqlScannerState::Normal;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            SqlScannerState::LineComment => {
                if bytes[i] == b'\n' {
                    state = SqlScannerState::Normal;
                }
                i += 1;
            }
            SqlScannerState::BlockComment(depth) => {
                if bytes[i] == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
                    *depth += 1;
                    i += 2;
                    continue;
                }
                if bytes[i] == b'*' && i + 1 < bytes.len() && bytes[i + 1] == b'/' {
                    *depth = depth.saturating_sub(1);
                    i += 2;
                    if *depth == 0 {
                        state = SqlScannerState::Normal;
                    }
                    continue;
                }
                i += 1;
            }
            SqlScannerState::DollarQuoted(tag) => {
                if matches_dollar_quote_end(bytes, i, tag) {
                    i += tag.len() + 2;
                    state = SqlScannerState::Normal;
                    continue;
                }
                i += 1;
            }
        }
    }

    push_statement(&mut statements, sql, statement_start, bytes.len());
    statements
}

fn push_statement<'a>(statements: &mut Vec<&'a str>, sql: &'a str, start: usize, end: usize) {
    let Some((trim_start, trim_end)) = trim_bounds(sql.as_bytes(), start, end) else {
        return;
    };
    let statement = trim_leading_comments(&sql[trim_start..trim_end]);
    if statement_is_comment_only(statement) {
        return;
    }
    statements.push(statement);
}

fn trim_bounds(bytes: &[u8], start: usize, end: usize) -> Option<(usize, usize)> {
    if start >= end || end > bytes.len() {
        return None;
    }
    let mut trim_start = start;
    while trim_start < end && bytes[trim_start].is_ascii_whitespace() {
        trim_start += 1;
    }
    if trim_start >= end {
        return None;
    }

    let mut trim_end = end;
    while trim_end > trim_start && bytes[trim_end - 1].is_ascii_whitespace() {
        trim_end -= 1;
    }
    if trim_end <= trim_start {
        return None;
    }

    Some((trim_start, trim_end))
}

fn parse_dollar_quote_tag(bytes: &[u8], start: usize) -> Option<(String, usize)> {
    if start >= bytes.len() || bytes[start] != b'$' {
        return None;
    }

    let mut idx = start + 1;
    while idx < bytes.len() && bytes[idx] != b'$' {
        let byte = bytes[idx];
        if !(byte.is_ascii_alphanumeric() || byte == b'_') {
            return None;
        }
        idx += 1;
    }

    if idx >= bytes.len() || bytes[idx] != b'$' {
        return None;
    }

    let tag = String::from_utf8(bytes[start + 1..idx].to_vec()).ok()?;
    Some((tag, idx - start + 1))
}

fn matches_dollar_quote_end(bytes: &[u8], start: usize, tag: &str) -> bool {
    let needed = tag.len() + 2;
    if start + needed > bytes.len() || bytes[start] != b'$' {
        return false;
    }

    let end = start + needed;
    if bytes[end - 1] != b'$' {
        return false;
    }

    bytes[start + 1..end - 1] == *tag.as_bytes()
}

fn statement_is_comment_only(statement: &str) -> bool {
    let bytes = statement.as_bytes();
    let mut state = SqlScannerState::Normal;
    let mut i = 0usize;

    while i < bytes.len() {
        match &mut state {
            SqlScannerState::Normal => {
                if bytes[i].is_ascii_whitespace() {
                    i += 1;
                    continue;
                }
                if bytes[i] == b'-' && i + 1 < bytes.len() && bytes[i + 1] == b'-' {
                    state = SqlScannerState::LineComment;
                    i += 2;
                    continue;
                }
                if bytes[i] == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
                    state = SqlScannerState::BlockComment(1);
                    i += 2;
                    continue;
                }
                return false;
            }
            SqlScannerState::LineComment => {
                if bytes[i] == b'\n' {
                    state = SqlScannerState::Normal;
                }
                i += 1;
            }
            SqlScannerState::BlockComment(depth) => {
                if bytes[i] == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
                    *depth += 1;
                    i += 2;
                    continue;
                }
                if bytes[i] == b'*' && i + 1 < bytes.len() && bytes[i + 1] == b'/' {
                    *depth = depth.saturating_sub(1);
                    i += 2;
                    if *depth == 0 {
                        state = SqlScannerState::Normal;
                    }
                    continue;
                }
                i += 1;
            }
            SqlScannerState::SingleQuotedString
            | SqlScannerState::DoubleQuotedIdentifier
            | SqlScannerState::DollarQuoted(_) => return false,
        }
    }

    true
}

fn trim_leading_comments(mut statement: &str) -> &str {
    loop {
        statement = statement.trim_start();
        if statement.is_empty() {
            return statement;
        }

        if let Some(rest) = statement.strip_prefix("--") {
            if let Some(newline_idx) = rest.find('\n') {
                statement = &rest[newline_idx + 1..];
                continue;
            }
            return "";
        }

        if let Some(rest) = statement.strip_prefix("/*") {
            if let Some(end_idx) = rest.find("*/") {
                statement = &rest[end_idx + 2..];
                continue;
            }
            return "";
        }

        return statement;
    }
}

#[cfg(test)]
mod tests {
    use super::{
        STORAGE_FILES_LEGACY_COMPAT_SQL, STORAGE_FILES_PROVISION_SQL, split_sql_statements,
    };

    #[test]
    fn split_sql_statements_ignores_leading_comments() {
        let sql = "-- comment\nCREATE TABLE test_one (id int);\n/* block */\nCREATE INDEX idx ON test_one(id);";
        let statements = split_sql_statements(sql);

        assert_eq!(statements.len(), 2);
        assert!(statements[0].starts_with("CREATE TABLE"));
        assert!(statements[1].starts_with("CREATE INDEX"));
    }

    #[test]
    fn bundled_sql_files_split_into_statements() {
        assert!(!split_sql_statements(STORAGE_FILES_PROVISION_SQL).is_empty());
        assert!(!split_sql_statements(STORAGE_FILES_LEGACY_COMPAT_SQL).is_empty());
    }

    #[test]
    fn split_sql_statements_preserves_semicolons_in_strings() {
        let statements =
            split_sql_statements("INSERT INTO logs(message) VALUES ('first;second');\nSELECT 1;");

        assert_eq!(statements.len(), 2);
        assert_eq!(
            statements[0],
            "INSERT INTO logs(message) VALUES ('first;second')"
        );
        assert_eq!(statements[1], "SELECT 1");
    }

    #[test]
    fn split_sql_statements_keeps_dollar_quoted_blocks_intact() {
        let sql = "CREATE TABLE demo (id int);\nDO $$\nBEGIN\n    PERFORM 1;\nEND $$;\nSELECT 1;";
        let statements = split_sql_statements(sql);

        assert_eq!(statements.len(), 3);
        assert_eq!(statements[0], "CREATE TABLE demo (id int)");
        assert_eq!(statements[1], "DO $$\nBEGIN\n    PERFORM 1;\nEND $$");
        assert_eq!(statements[2], "SELECT 1");
    }
}
