use crate::storage::Error;
use change_case::snake_case;
use futures_util::TryStreamExt;
use sqlx::{Pool, Postgres, Row};
use std::collections::HashMap;

#[derive(Clone, Debug)]
pub struct Field {
    pub column_name: String,
    pub data_type: String,
    pub is_nullable: bool,
}

#[derive(Clone, Debug)]
pub struct Table {
    pub name: String,
    pub fields: Vec<Field>,
    pub fields_by_name: HashMap<String, Field>,
}

pub async fn load_schema(pool: &Pool<Postgres>) -> Result<HashMap<String, Table>, Error> {
    let query = "
        SELECT table_name, column_name, data_type, is_nullable
        FROM information_schema.columns
        WHERE table_schema = 'public'
    ";
    let mut rows = sqlx::query(query).fetch(pool);
    let mut tables_by_name: HashMap<String, Table> = HashMap::new();
    while let Some(row) = rows.try_next().await.map_err(|e| Error::internal(e.to_string()))? {
        let table_name: String = row.get(0);
        let column_name: String = row.get(1);
        let data_type: String = row.get(2);
        let is_nullable: bool = row.get::<&str, _>(3) == "YES";
        let column_name_snake_case = snake_case(&column_name);
        let field = Field {
            column_name,
            data_type,
            is_nullable,
        };
        if let Some(table) = tables_by_name.get_mut(&table_name) {
            table.fields.push(field.clone());
            table.fields_by_name.insert(column_name_snake_case, field);
        } else {
            let mut fields_by_name = HashMap::new();
            fields_by_name.insert(column_name_snake_case, field.clone());
            let table = Table {
                name: table_name.clone(),
                fields: vec![field],
                fields_by_name,
            };
            tables_by_name.insert(table_name, table);
        }
    }
    Ok(tables_by_name)
}
