use sqlparser::{
    ast::{Query, SetExpr, Statement},
    dialect::PostgreSqlDialect,
    parser::Parser,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SqlRoute {
    Read,
    Write,
}

pub fn route_sql(sql: &str) -> SqlRoute {
    let trimmed = sql.trim_start();

    if starts_with_hint(trimmed, "write") {
        return SqlRoute::Write;
    }

    if starts_with_hint(trimmed, "read") {
        return SqlRoute::Read;
    }

    if is_read_query_by_ast(sql) {
        SqlRoute::Read
    } else {
        SqlRoute::Write
    }
}

fn starts_with_hint(sql: &str, hint: &str) -> bool {
    let expected = format!("/*+ {hint} */");
    sql.starts_with(&expected)
}

fn is_read_query_by_ast(sql: &str) -> bool {
    let dialect = PostgreSqlDialect {};

    let Ok(statements) = Parser::parse_sql(&dialect, sql) else {
        return false;
    };

    if statements.is_empty() {
        return false;
    }

    statements.iter().all(is_read_statement)
}

fn is_read_statement(statement: &Statement) -> bool {
    match statement {
        Statement::Query(query) => is_read_query_ast(query),
        Statement::ShowVariable { .. }
        | Statement::ShowVariables { .. }
        | Statement::ShowCreate { .. }
        | Statement::ShowColumns { .. }
        | Statement::ShowTables { .. } => true,
        _ => false,
    }
}

fn is_read_query_ast(query: &Query) -> bool {
    if !query.locks.is_empty() {
        return false;
    }

    match query.body.as_ref() {
        SetExpr::Select(_) => true,
        SetExpr::Query(inner) => is_read_query_ast(inner),
        SetExpr::SetOperation { left, right, .. } => {
            is_read_set_expr(left) && is_read_set_expr(right)
        }
        SetExpr::Values(_) => true,
        _ => false,
    }
}

fn is_read_set_expr(expr: &SetExpr) -> bool {
    match expr {
        SetExpr::Select(_) => true,
        SetExpr::Query(query) => is_read_query_ast(query),
        SetExpr::SetOperation { left, right, .. } => {
            is_read_set_expr(left) && is_read_set_expr(right)
        }
        SetExpr::Values(_) => true,
        _ => false,
    }
}