jspg/src/queryer/compiler.rs

use crate::database::Database;
use std::sync::Arc;

pub struct SqlCompiler {
  pub db: Arc<Database>,
}

impl SqlCompiler {
  pub fn new(db: Arc<Database>) -> Self {
    Self { db }
  }

  /// Compiles a JSON schema into a nested PostgreSQL query returning JSONB
  pub fn compile(
    &self,
    schema_id: &str,
    stem_path: Option<&str>,
    filter_keys: &[String],
  ) -> Result<String, String> {
    let schema = self
      .db
      .schemas
      .get(schema_id)
      .ok_or_else(|| format!("Schema not found: {}", schema_id))?;

    let resolved_arc;
    let target_schema = if let Some(path) = stem_path.filter(|p| !p.is_empty() && *p != "/") {
      if let Some(stems_map) = self.db.stems.get(schema_id) {
        if let Some(stem) = stems_map.get(path) {
          resolved_arc = stem.schema.clone();
        } else {
          return Err(format!(
            "Stem entity type '{}' not found in schema '{}'",
            path, schema_id
          ));
        }
      } else {
        return Err(format!(
          "Stem entity type '{}' not found in schema '{}'",
          path, schema_id
        ));
      }
      resolved_arc.as_ref()
    } else {
      schema
    };

    // We expect the top level to typically be an Object or Array
    let is_stem_query = stem_path.is_some();
    let (sql, _) = self.walk_schema(target_schema, "t1", None, filter_keys, is_stem_query, 0, String::new())?;
    Ok(sql)
  }

  /// Recursively walks the schema AST emitting native PostgreSQL jsonb mapping
  /// Returns a tuple of (SQL_String, Field_Type)
  fn walk_schema(
    &self,
    schema: &crate::database::schema::Schema,
    parent_alias: &str,
    prop_name_context: Option<&str>,
    filter_keys: &[String],
    is_stem_query: bool,
    depth: usize,
    current_path: String,
  ) -> Result<(String, String), String> {
    // Determine the base schema type (could be an array, object, or literal)
    match &schema.obj.type_ {
      Some(crate::database::schema::SchemaTypeOrArray::Single(t)) if t == "array" => {
        // Handle Arrays:
        if let Some(items) = &schema.obj.items {
          let next_path = if current_path.is_empty() {
            String::from("#")
          } else {
            format!("{}.#", current_path)
          };

          if let Some(ref_id) = &items.obj.r#ref {
            if let Some(type_def) = self.db.types.get(ref_id) {
              return self.compile_entity_node(
                items,
                type_def,
                parent_alias,
                prop_name_context,
                true,
                filter_keys,
                is_stem_query,
                depth,
                next_path,
              );
            }
          }
          let (item_sql, _) = self.walk_schema(
            items,
            parent_alias,
            prop_name_context,
            filter_keys,
            is_stem_query,
            depth + 1,
            next_path,
          )?;
          return Ok((
            format!("(SELECT jsonb_agg({}) FROM TODO)", item_sql),
            "array".to_string(),
          ));
        }

        Ok((
          "SELECT jsonb_agg(TODO) FROM TODO".to_string(),
          "array".to_string(),
        ))
      }
      _ => {
        // Determine if this schema represents a Database Entity
        let mut resolved_type = None;

        // Target is generally a specific schema (e.g. 'base.person'), but it tells us what physical
        // database table hierarchy it maps to via the `schema.id` prefix/suffix convention.
        if let Some(lookup_key) = schema.obj.id.as_ref().or(schema.obj.r#ref.as_ref()) {
          let base_type_name = lookup_key.split('.').next_back().unwrap_or("").to_string();
          resolved_type = self.db.types.get(&base_type_name);
        }

        if let Some(type_def) = resolved_type {
          return self.compile_entity_node(
            schema,
            type_def,
            parent_alias,
            prop_name_context,
            false,
            filter_keys,
            is_stem_query,
            depth,
            current_path,
          );
        }

        // Handle Direct Refs
        if let Some(ref_id) = &schema.obj.r#ref {
          // If it's just an ad-hoc struct ref, we should resolve it
          if let Some(target_schema) = self.db.schemas.get(ref_id) {
            return self.walk_schema(
              target_schema,
              parent_alias,
              prop_name_context,
              filter_keys,
              is_stem_query,
              depth,
              current_path,
            );
          }
          return Err(format!("Unresolved $ref: {}", ref_id));
        }

        // Just an inline object definition?
        if let Some(props) = &schema.obj.properties {
          return self.compile_inline_object(
            props,
            parent_alias,
            filter_keys,
            is_stem_query,
            depth,
            current_path,
          );
        }

        // Literal fallback
        Ok((
          format!(
            "{}.{}",
            parent_alias,
            prop_name_context.unwrap_or("unknown_prop")
          ),
          "string".to_string(),
        ))
      }
    }
  }

  fn get_merged_properties(
    &self,
    schema: &crate::database::schema::Schema,
  ) -> std::collections::BTreeMap<String, Arc<crate::database::schema::Schema>> {
    let mut props = std::collections::BTreeMap::new();

    if let Some(ref_id) = &schema.obj.r#ref {
      if let Some(parent_schema) = self.db.schemas.get(ref_id) {
        props.extend(self.get_merged_properties(parent_schema));
      }
    }

    if let Some(local_props) = &schema.obj.properties {
      for (k, v) in local_props {
        props.insert(k.clone(), v.clone());
      }
    }

    props
  }

  fn compile_entity_node(
    &self,
    schema: &crate::database::schema::Schema,
    type_def: &crate::database::r#type::Type,
    parent_alias: &str,
    prop_name: Option<&str>,
    is_array: bool,
    filter_keys: &[String],
    is_stem_query: bool,
    depth: usize,
    current_path: String,
  ) -> Result<(String, String), String> {
    let local_ctx = format!("{}_{}", parent_alias, prop_name.unwrap_or("obj"));

    // 1. Build FROM clauses and table aliases
    let (table_aliases, from_clauses) = self.build_hierarchy_from_clauses(type_def, &local_ctx);

    // 2. Map properties and build jsonb_build_object args
    let select_args = self.map_properties_to_aliases(
      schema,
      type_def,
      &table_aliases,
      parent_alias,
      filter_keys,
      is_stem_query,
      depth,
      &current_path,
    )?;

    let jsonb_obj_sql = if select_args.is_empty() {
      "jsonb_build_object()".to_string()
    } else {
      format!("jsonb_build_object({})", select_args.join(", "))
    };

    // 3. Build WHERE clauses
    let where_clauses = self.build_filter_where_clauses(
      schema,
      type_def,
      &table_aliases,
      parent_alias,
      prop_name,
      filter_keys,
      &current_path,
    )?;

    let selection = if is_array {
      format!("COALESCE(jsonb_agg({}), '[]'::jsonb)", jsonb_obj_sql)
    } else {
      jsonb_obj_sql
    };

    let full_sql = format!(
      "(SELECT {} FROM {} WHERE {})",
      selection,
      from_clauses.join(" "),
      where_clauses.join(" AND ")
    );

    Ok((
      full_sql,
      if is_array {
        "array".to_string()
      } else {
        "object".to_string()
      },
    ))
  }

  fn build_hierarchy_from_clauses(
    &self,
    type_def: &crate::database::r#type::Type,
    local_ctx: &str,
  ) -> (std::collections::HashMap<String, String>, Vec<String>) {
    let mut table_aliases = std::collections::HashMap::new();
    let mut from_clauses = Vec::new();

    for (i, table_name) in type_def.hierarchy.iter().enumerate() {
      let alias = format!("{}_t{}", local_ctx, i + 1);
      table_aliases.insert(table_name.clone(), alias.clone());

      if i == 0 {
        from_clauses.push(format!("agreego.{} {}", table_name, alias));
      } else {
        let prev_alias = format!("{}_t{}", local_ctx, i);
        from_clauses.push(format!(
          "JOIN agreego.{} {} ON {}.id = {}.id",
          table_name, alias, alias, prev_alias
        ));
      }
    }
    (table_aliases, from_clauses)
  }

  fn map_properties_to_aliases(
    &self,
    schema: &crate::database::schema::Schema,
    type_def: &crate::database::r#type::Type,
    table_aliases: &std::collections::HashMap<String, String>,
    parent_alias: &str,
    filter_keys: &[String],
    is_stem_query: bool,
    depth: usize,
    current_path: &str,
  ) -> Result<Vec<String>, String> {
    let mut select_args = Vec::new();
    let grouped_fields = type_def.grouped_fields.as_ref().and_then(|v| v.as_object());
    let merged_props = self.get_merged_properties(schema);

    for (prop_key, prop_schema) in &merged_props {
      let mut owner_alias = table_aliases
        .get("entity")
        .cloned()
        .unwrap_or_else(|| format!("{}_t_err", parent_alias));

      if let Some(gf) = grouped_fields {
        for (t_name, fields_val) in gf {
          if let Some(fields_arr) = fields_val.as_array() {
            if fields_arr.iter().any(|v| v.as_str() == Some(prop_key)) {
              owner_alias = table_aliases
                .get(t_name)
                .cloned()
                .unwrap_or_else(|| parent_alias.to_string());
              break;
            }
          }
        }
      }

      let next_path = if current_path.is_empty() {
        prop_key.clone()
      } else {
        format!("{}.{}", current_path, prop_key)
      };

      let (val_sql, val_type) = self.walk_schema(
        prop_schema,
        &owner_alias,
        Some(prop_key),
        filter_keys,
        is_stem_query,
        depth + 1,
        next_path,
      )?;

      if val_type != "abort" {
        select_args.push(format!("'{}', {}", prop_key, val_sql));
      }
    }
    Ok(select_args)
  }

  fn build_filter_where_clauses(
    &self,
    schema: &crate::database::schema::Schema,
    type_def: &crate::database::r#type::Type,
    table_aliases: &std::collections::HashMap<String, String>,
    parent_alias: &str,
    prop_name: Option<&str>,
    filter_keys: &[String],
    current_path: &str,
  ) -> Result<Vec<String>, String> {
    let base_alias = table_aliases
      .get(&type_def.name)
      .cloned()
      .unwrap_or_else(|| "err".to_string());

    let mut where_clauses = Vec::new();
    where_clauses.push(format!("NOT {}.archived", base_alias));

    for (i, filter_key) in filter_keys.iter().enumerate() {
      let mut parts = filter_key.split(':');
      let full_field_path = parts.next().unwrap_or(filter_key);
      let op = parts.next().unwrap_or("$eq");

      let field_name = if current_path.is_empty() {
        if full_field_path.contains('.') || full_field_path.contains('#') {
          continue;
        }
        full_field_path
      } else {
        let prefix = format!("{}.", current_path);
        if full_field_path.starts_with(&prefix) {
          let remainder = &full_field_path[prefix.len()..];
          if remainder.contains('.') || remainder.contains('#') {
            continue;
          }
          remainder
        } else {
          continue;
        }
      };

      let mut filter_alias = base_alias.clone();

      if let Some(gf) = type_def.grouped_fields.as_ref().and_then(|v| v.as_object()) {
          for (t_name, fields_val) in gf {
            if let Some(fields_arr) = fields_val.as_array() {
              if fields_arr.iter().any(|v| v.as_str() == Some(field_name)) {
                filter_alias = table_aliases
                  .get(t_name)
                  .cloned()
                  .unwrap_or_else(|| base_alias.clone());
                break;
              }
            }
          }
        }

        let mut is_ilike = false;
        let mut cast = "";

        if let Some(field_types) = type_def.field_types.as_ref().and_then(|v| v.as_object()) {
          if let Some(pg_type_val) = field_types.get(field_name) {
            if let Some(pg_type) = pg_type_val.as_str() {
              if pg_type == "uuid" {
                cast = "::uuid";
              } else if pg_type == "boolean" || pg_type == "bool" {
                cast = "::boolean";
              } else if pg_type.contains("timestamp")
                || pg_type == "timestamptz"
                || pg_type == "date"
              {
                cast = "::timestamptz";
              } else if pg_type == "numeric"
                || pg_type.contains("int")
                || pg_type == "real"
                || pg_type == "double precision"
              {
                cast = "::numeric";
              } else if pg_type == "text" || pg_type.contains("char") {
                let mut is_enum = false;
                if let Some(props) = &schema.obj.properties {
                  if let Some(ps) = props.get(field_name) {
                    is_enum = ps.obj.enum_.is_some();
                  }
                }
                if !is_enum {
                  is_ilike = true;
                }
              }
            }
          }
        }

        let param_index = i + 1;
        let p_val = format!("${}#>>'{{}}'", param_index);

        if op == "$in" || op == "$nin" {
          let sql_op = if op == "$in" { "IN" } else { "NOT IN" };
          let subquery = format!(
            "(SELECT value{} FROM jsonb_array_elements_text(({})::jsonb))",
            cast, p_val
          );
          where_clauses.push(format!(
            "{}.{} {} {}",
            filter_alias, field_name, sql_op, subquery
          ));
        } else {
          let sql_op = match op {
            "$eq" => {
              if is_ilike {
                "ILIKE"
              } else {
                "="
              }
            }
            "$ne" => {
              if is_ilike {
                "NOT ILIKE"
              } else {
                "!="
              }
            }
            "$gt" => ">",
            "$gte" => ">=",
            "$lt" => "<",
            "$lte" => "<=",
            _ => {
              if is_ilike {
                "ILIKE"
              } else {
                "="
              }
            }
          };

          let param_sql = if is_ilike && (op == "$eq" || op == "$ne") {
            p_val
          } else {
            format!("({}){}", p_val, cast)
          };

          where_clauses.push(format!(
            "{}.{} {} {}",
            filter_alias, field_name, sql_op, param_sql
          ));
      }
    }

    if let Some(_prop) = prop_name {
      where_clauses.push(format!("{}.parent_id = {}.id", base_alias, parent_alias));
    }

    Ok(where_clauses)
  }

  fn compile_inline_object(
    &self,
    props: &std::collections::BTreeMap<String, std::sync::Arc<crate::database::schema::Schema>>,
    parent_alias: &str,
    filter_keys: &[String],
    is_stem_query: bool,
    depth: usize,
    current_path: String,
  ) -> Result<(String, String), String> {
    let mut build_args = Vec::new();
    for (k, v) in props {
      let next_path = if current_path.is_empty() {
        k.clone()
      } else {
        format!("{}.{}", current_path, k)
      };
      
      let (child_sql, val_type) = self.walk_schema(
        v,
        parent_alias,
        Some(k),
        filter_keys,
        is_stem_query,
        depth + 1,
        next_path,
      )?;
      if val_type == "abort" {
        continue;
      }
      build_args.push(format!("'{}', {}", k, child_sql));
    }
    let combined = format!("jsonb_build_object({})", build_args.join(", "));
    Ok((combined, "object".to_string()))
  }
}