db-cores 0.1.0

use sqlparser::{
    ast::{
        Expr,  LimitClause, ObjectName, ObjectType, 
        Statement,  ValueWithSpan,OffsetRows, Value as AstValue,Offset, Query,  SetExpr
    },
    dialect::{MsSqlDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect,GenericDialect},
    tokenizer::{Location as TokenizeLocation, Span},
    parser::{Parser, ParserError}
};
// use sqlparser::ast::helper::Span;
use crate::common::DatabaseKind;
#[derive(Debug, Clone, Default)]
pub struct ParseSQL {
    pub returns_result: bool,       // 判断是不是有返回值
    pub forbidden_in_txn: bool,     // 判断能不能在事务中执行
    pub table_name: Option<String>, // 获取表名
    pub kind: String,               // sql类型
    pub offset: Option<u64>,
    pub limit: Option<u64>,
    pub update_sql: Option<String>, // 对没有limit的语句,增加 limit限制后的 select 语句
}

impl ParseSQL {
    pub fn other_new(
        returns_result: bool,
        forbidden_in_txn: bool,
        table_name: Option<&str>,
        kind: &str,
    ) -> Self {
        Self {
            returns_result,
            forbidden_in_txn,
            table_name: table_name.filter(|s| !s.is_empty()).map(|s| s.to_string()),
            kind: kind.to_string(),
            offset: None,
            limit: None,
            update_sql: None,
        }
    }

    pub fn select_new(
        update_sql: Option<String>,
        table_name: Option<String>,
        limit: Option<u64>,
        offset: Option<u64>,
    ) -> Self {
        Self {
            returns_result: true,
            forbidden_in_txn: false,
            table_name,
            kind: "query".to_string(),
            offset,
            limit,
            update_sql,
        }
    }

    pub fn to_count_sql(sql: &str, kind: DatabaseKind) -> anyhow::Result<Option<String>> {
        let mut stmt = Self::to_statement(sql, &kind)?;
        if let Statement::Query(q) = &mut stmt {
            if q.limit_clause.is_some() {
                q.limit_clause = None;
                let count_sql = format!("SELECT COUNT(*) FROM ({}) t", stmt);
                return Ok(Some(count_sql));
            }
        }
        Ok(None) // 没有 LIMIT/OFFSET
    }

    pub fn to_statement(sql: &str, kind: &DatabaseKind) -> Result<Statement, ParserError> {
        let mut parser = match kind {
            DatabaseKind::Postgres => Parser::new(&PostgreSqlDialect {}),
            DatabaseKind::Sqlite => Parser::new(&SQLiteDialect {}),
            DatabaseKind::MySql => Parser::new(&MySqlDialect {}),
            DatabaseKind::Access => Parser::new(&MsSqlDialect {}),
            DatabaseKind::Api => Parser::new(&GenericDialect {}),
            _ => Parser::new(&GenericDialect {}),
        };
        parser = parser.try_with_sql(sql)?;
        parser.parse_statement()
    }

    pub fn full_table_name(object_name: &ObjectName) -> String {
        object_name
            .0
            .iter()
            .filter_map(|ident| ident.as_ident())
            .map(|ident| match ident.quote_style {
                Some(q) => format!("{q}{}{q}", ident.value),
                None => ident.value.clone(),
            })
            .collect::<Vec<_>>()
            .join(".")
    }

    /// 修改 Query 的 LIMIT，如果不存在则添加默认 LIMIT 返回 (limit, offset)
    pub fn modify_limit(query: &mut Box<Query>) -> (u64, u64) {
        // 默认值
        let default_limit = 1000;
        let default_offset = 0;

        match &mut query.limit_clause {
            // PostgreSQL 风格：LIMIT ... OFFSET ...
            Some(LimitClause::LimitOffset { limit, offset, .. }) => {
                // 补默认 limit
                if limit.is_none() {
                    *limit = Some(Expr::Value(ValueWithSpan {
                        value: AstValue::Number(default_limit.to_string(), false),
                        span: Span {
                            start: TokenizeLocation { line: 0, column: 0 },
                            end: TokenizeLocation { line: 0, column: 0 },
                        },
                    }));
                }
                // 获取 offset，如果没有则默认 0
                let offset_val = offset
                    .as_ref()
                    .and_then(|o| {
                        if let Expr::Value(ValueWithSpan {
                            value: AstValue::Number(n, _),
                            ..
                        }) = &o.value
                        {
                            n.parse::<u64>().ok()
                        } else {
                            None
                        }
                    })
                    .unwrap_or(default_offset);

                let limit_val = limit
                    .as_ref()
                    .and_then(|l| {
                        if let Expr::Value(ValueWithSpan {
                            value: AstValue::Number(n, _),
                            ..
                        }) = l
                        {
                            n.parse::<u64>().ok()
                        } else {
                            None
                        }
                    })
                    .unwrap_or(default_limit);

                (limit_val, offset_val)
            }

            // MySQL 风格：LIMIT offset, limit
            Some(LimitClause::OffsetCommaLimit { offset, limit }) => {
                let offset_val = if let Expr::Value(ValueWithSpan {
                    value: AstValue::Number(n, _),
                    ..
                }) = offset
                {
                    n.parse::<u64>().unwrap_or(default_offset)
                } else {
                    default_offset
                };

                let limit_val = if let Expr::Value(ValueWithSpan {
                    value: AstValue::Number(n, _),
                    ..
                }) = limit
                {
                    n.parse::<u64>().unwrap_or(default_limit)
                } else {
                    default_limit
                };

                (limit_val, offset_val)
            }

            // 没有 LIMIT，直接加默认 LIMIT
            None => {
                query.limit_clause = Some(LimitClause::LimitOffset {
                    limit: Some(Expr::Value(ValueWithSpan {
                        value: AstValue::Number(default_limit.to_string(), false),
                        span: Span {
                            start: TokenizeLocation { line: 0, column: 0 },
                            end: TokenizeLocation { line: 0, column: 0 },
                        },
                    })),
                    offset: Some(Offset {
                        value: Expr::Value(ValueWithSpan {
                            value: AstValue::Number(default_offset.to_string(), false),
                            span: Span {
                                start: TokenizeLocation { line: 0, column: 0 },
                                end: TokenizeLocation { line: 0, column: 0 },
                            },
                        }),
                        rows: OffsetRows::None,
                    }),
                    limit_by: vec![],
                });

                (default_limit, default_offset)
            }
        }
    }

    pub fn parse(sql: &str, kind: DatabaseKind) -> anyhow::Result<Self> {
        let mut stmt = Self::to_statement(sql, &kind)?;
        match &mut stmt {
            Statement::Query(query) => {
                let mut table_name = None;
                match &mut *query.body {
                    // 判断最外层是否可以加 LIMIT/OFFSET
                    SetExpr::Select(select) => {
                        for table_with_joins in &select.from {
                            if let sqlparser::ast::TableFactor::Table { name, .. } =
                                &table_with_joins.relation
                            {
                                table_name = Some(Self::full_table_name(name));
                                break;
                            }
                        }

                        let (limit, offset) = Self::modify_limit(query); // 只修改最外层
                        return Ok(Self::select_new(
                            Some(stmt.to_string()),
                            table_name,
                            Some(limit),
                            Some(offset),
                        ));
                    }
                    SetExpr::SetOperation { .. } => {
                        // SetOperation 就是多个 SELECT 合并的情况：
                        let (limit, offset) = Self::modify_limit(query); // 只修改最外层
                        return Ok(Self::select_new(
                            Some(stmt.to_string()),
                            table_name,
                            Some(limit),
                            Some(offset),
                        ));
                    }
                    _ => {}
                }
                Ok(Self::select_new(None, table_name, None, None))
            }
            Statement::Delete(d) if d.returning.is_some() => {
                Ok(Self::other_new(true, false, None, "delete"))
            }
            Statement::Insert(i) if i.returning.is_some() => {
                Ok(Self::other_new(true, false, None, "insert"))
            }
            Statement::Update {
                returning: Some(_), ..
            } => Ok(Self::other_new(true, false, None, "update")),
            Statement::ShowDatabases { .. }
            | Statement::ShowSchemas { .. }
            | Statement::ShowObjects { .. }
            | Statement::ShowVariable { .. }
            | Statement::ShowVariables { .. }
            | Statement::ShowStatus { .. }
            | Statement::ShowTables { .. }
            | Statement::ShowColumns { .. }
            | Statement::ShowViews { .. }
            | Statement::ShowCollation { .. }
            | Statement::ShowCreate { .. }
            | Statement::Explain { .. }
            | Statement::LISTEN { .. }
            | Statement::List { .. }
            | Statement::Fetch { .. }
            | Statement::Execute { .. }
            | Statement::Call { .. }
            | Statement::Merge {
                output: Some(_), ..
            } => Ok(Self::other_new(true, false, None, "other")),
            Statement::StartTransaction { .. }
            | Statement::Commit { .. }
            | Statement::Rollback { .. }
            | Statement::CreateDatabase { .. }
            | Statement::CreateSchema { .. }
            | Statement::AlterType { .. }
            | Statement::CreateIndex { .. }
            | Statement::Flush { .. } => Ok(Self::other_new(false, true, None, "ddl")),
            Statement::Drop { object_type, .. } => {
                let forbidden_in_txn =
                    matches!(object_type, ObjectType::Database | ObjectType::Table);
                Ok(Self::other_new(false, forbidden_in_txn, None, "other"))
            }
            Statement::Pragma {
                value: _,
                name,
                is_eq,
            } => {
                if *is_eq {
                    Ok(Self::other_new(false, false, None, "other"))
                } else {
                    let is_res = matches!(
                        name.to_string().to_lowercase().as_str(),
                        "table_info"
                            | "index_list"
                            | "index_info"
                            | "page_size"
                            | "encoding"
                            | "journal_mode"
                            | "foreign_key_list"
                            | "database_list"
                            | "collation_list"
                            | "pragma_list"
                            | "table_xinfo"
                    );
                    Ok(Self::other_new(is_res, false, None, "other"))
                }
            }
            _ => Ok(Self::other_new(false, false, None, "")),
        }
    }

    pub fn split_sql(sql: &str, kind: &DatabaseKind) -> anyhow::Result<Vec<String>> {
        use sqlparser::tokenizer::{Token, Tokenizer};
        // 方言（决定标识符/引号等规则）
        // 用给定方言和输入字符串创建 tokenizer
        let mut tokenizer = match kind {
            DatabaseKind::Postgres => Tokenizer::new(&PostgreSqlDialect {}, sql),
            DatabaseKind::MySql => Tokenizer::new(&MySqlDialect {}, sql),
            DatabaseKind::SqlServer => Tokenizer::new(&MsSqlDialect {}, sql),
            DatabaseKind::Sqlite => Tokenizer::new(&SQLiteDialect {}, sql),
            _ => Tokenizer::new(&GenericDialect {}, sql),
        };

        // 把输入分成 token 列表（可能返回错误，所以用 ? 传播）
        let tokens = tokenizer.tokenize()?;
        // 用来收集最终的 SQL 语句
        let mut stmts = Vec::new();
        // 正在构建的当前语句文本
        let mut current = String::new();

        for token in tokens {
            match token {
                Token::SemiColon => {
                    if !current.trim().is_empty() {
                        stmts.push(current.trim().to_string());
                        current.clear();
                    }
                }
                _ => {
                    current.push_str(&token.to_string());
                }
            }
        }

        // 循环结束后若还有残余（没有以分号结尾），也加入结果
        if !current.trim().is_empty() {
            stmts.push(current.trim().to_string());
        }

        Ok(stmts)
    }

    
}

#[test]
fn test_parse_sql() {
    let _ = "SELECT * FROM (SELECT * FROM table2 WHERE id = 1 AND name = 'John') WHERE id = 1 AND name = 'John'";

    pub fn test_sql_statements() -> Vec<&'static str> {
        vec![
        // 1. 简单 SELECT
        "SELECT id, name FROM users",

        // 2. SELECT 带 WHERE、ORDER BY
        "SELECT id, name FROM users WHERE age > 18 ORDER BY name ASC",

        // 3. SELECT 带 LIMIT/OFFSET
        "SELECT * FROM products LIMIT 50 OFFSET 10",

        // 4. SELECT 带子查询
        "SELECT u.id, u.name FROM users u WHERE u.id IN (SELECT user_id FROM orders WHERE total > 100)",

        // 5. 多个 SELECT 合并 (UNION)
        "SELECT id FROM users WHERE active = true UNION SELECT id FROM admins",

        // 6. INSERT 带 RETURNING
        "INSERT INTO users (name, age) VALUES ('Alice', 30) RETURNING id",

        // 7. UPDATE 带 RETURNING
        "UPDATE products SET price = price * 1.1 WHERE category = 'books' RETURNING id, price",

        // 8. DELETE 带 RETURNING
        "DELETE FROM sessions WHERE expires < NOW() RETURNING id",

        // 9. DDL: CREATE TABLE
        "CREATE TABLE logs (id SERIAL PRIMARY KEY, message TEXT, created_at TIMESTAMP)",

        // 10. SHOW / PRAGMA /其他
        "SHOW TABLES",
    ]
    }

    let test_sqls = test_sql_statements();
    for sql in test_sqls {
        let res = ParseSQL::parse(sql, DatabaseKind::Postgres).unwrap();
        println!("{}:{:#?}", sql, res);
    }
    // println!("1:{:?}", p);
    // let p = parse_sql("DELETE FROM tableq WHERE id = 1 AND name = 'John' RETURNING *;");
    // println!("2:{:?}", p);
    // let p = is_query_sql("VALUES ('2023-01-01', 100), ('2023-01-02', 200);");
    // println!("VALUES ('2023-01-01', 100), ('2023-01-02', 200);:{:?}", p);
}