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 target_schema = if let Some(path) = stem_path.filter(|p| !p.is_empty() && *p != "/") {
      self.resolve_stem(schema, path)?
    } else {
      schema
    };

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

  fn resolve_stem<'a>(
    &'a self,
    mut schema: &'a crate::database::schema::Schema,
    path: &str,
  ) -> Result<&'a crate::database::schema::Schema, String> {
    let parts: Vec<&str> = path.trim_start_matches('/').split('/').collect();
    for part in parts {
      let mut current = schema;
      let mut depth = 0;
      while let Some(r) = &current.obj.r#ref {
        if let Some(s) = self.db.schemas.get(r) {
          current = s;
        } else {
          break;
        }
        depth += 1;
        if depth > 20 {
          break;
        }
      }

      if current.obj.properties.is_none() && current.obj.items.is_some() {
        if let Some(items) = &current.obj.items {
          current = items;
          let mut depth2 = 0;
          while let Some(r) = &current.obj.r#ref {
            if let Some(s) = self.db.schemas.get(r) {
              current = s;
            } else {
              break;
            }
            depth2 += 1;
            if depth2 > 20 {
              break;
            }
          }
        }
      }

      if let Some(props) = &current.obj.properties {
        if let Some(next_schema) = props.get(part) {
          schema = next_schema;
        } else {
          return Err(format!("Stem part '{}' not found in schema", part));
        }
      } else {
        return Err(format!(
          "Cannot resolve stem part '{}': not an object",
          part
        ));
      }
    }

    let mut current = schema;
    let mut depth = 0;
    while let Some(r) = &current.obj.r#ref {
      if let Some(s) = self.db.schemas.get(r) {
        current = s;
      } else {
        break;
      }
      depth += 1;
      if depth > 20 {
        break;
      }
    }
    Ok(current)
  }

  /// 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],
  ) -> 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 {
          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,
              );
            }
          }
          let (item_sql, _) =
            self.walk_schema(items, parent_alias, prop_name_context, filter_keys)?;
          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(),
        ))
      }
      _ => {
        // Handle Objects & Direct Refs
        if let Some(ref_id) = &schema.obj.r#ref {
          // If it's a $ref, check if it points to an Entity Type
          if let Some(type_def) = self.db.types.get(ref_id) {
            return self.compile_entity_node(
              schema,
              type_def,
              parent_alias,
              prop_name_context,
              false,
              filter_keys,
            );
          }
          // 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);
          }
          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);
        }

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

  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],
  ) -> Result<(String, String), String> {
    // We are compiling a query block for an Entity.
    let mut select_args = Vec::new();

    // Mapping table hierarchy to aliases, e.g., ["person", "user", "organization", "entity"]
    let local_ctx = format!("{}_{}", parent_alias, prop_name.unwrap_or("obj"));
    // e.g., parent_t1_contact -> we'll use t1 for the first of this block, t2 for the second, etc.
    // Actually, local_ctx can just be exactly that prop's unique path.
    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 {
        // Join to previous
        let prev_alias = format!("{}_t{}", local_ctx, i);
        from_clauses.push(format!(
          "JOIN agreego.{} {} ON {}.id = {}.id",
          table_name, alias, alias, prev_alias
        ));
      }
    }

    // Now, let's map properties from the schema to the correct table alias using grouped_fields
    // grouped_fields is { "person": ["first_name", ...], "user": ["password"], ... }
    let grouped_fields = type_def.grouped_fields.as_ref().and_then(|v| v.as_object());

    if let Some(props) = &schema.obj.properties {
      for (prop_key, prop_schema) in props {
        // Find which table owns this property
        // Find which table owns this property
        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;
              }
            }
          }
        }

        // Now we know `owner_alias`, e.g., `parent_t1` or `parent_t3`.
        // Walk the property to get its SQL value
        let (val_sql, _) =
          self.walk_schema(prop_schema, &owner_alias, Some(prop_key), filter_keys)?;
        select_args.push(format!("'{}', {}", prop_key, val_sql));
      }
    }

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

    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));
    // Filter Mapping - Only append filters if this is the ROOT table query (i.e. parent_alias is "t1")
    // Because cue.filters operates strictly on top-level root properties right now.
    if parent_alias == "t1" && prop_name.is_none() {
      for (i, filter_key) in filter_keys.iter().enumerate() {
        // Find which table owns this filter key
        let mut filter_alias = base_alias.clone(); // default to root table (e.g. t3 entity)

        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(filter_key)) {
                filter_alias = table_aliases
                  .get(t_name)
                  .cloned()
                  .unwrap_or_else(|| base_alias.clone());
                break;
              }
            }
          }
        }

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

        // Check schema for filter_key to determine datatype operation
        if let Some(props) = &schema.obj.properties {
          if let Some(ps) = props.get(filter_key) {
            let is_enum = ps.obj.enum_.is_some();
            if let Some(crate::database::schema::SchemaTypeOrArray::Single(t)) = &ps.obj.type_ {
              if t == "string" {
                if ps.obj.format.as_deref() == Some("uuid") {
                  cast = "::uuid";
                } else if ps.obj.format.as_deref() == Some("date-time") {
                  cast = "::timestamptz";
                } else if !is_enum {
                  is_ilike = true;
                }
              } else if t == "boolean" {
                cast = "::boolean";
              } else if t == "integer" || t == "number" {
                cast = "::numeric";
              }
            }
          }
        }

        // Add to WHERE clause using 1-indexed args pointer: $1, $2
        if is_ilike {
          let param = format!("${}#>>'{{}}'", i + 1);
          where_clauses.push(format!("{}.{} ILIKE {}", filter_alias, filter_key, param));
        } else {
          let param = format!("(${}#>>'{{}}'){}", i + 1, cast);
          where_clauses.push(format!("{}.{} = {}", filter_alias, filter_key, param));
        }
      }
    }

    // Resolve FK relationship constraint if this is a nested subquery
    if let Some(_prop) = prop_name {
      // MOCK relation resolution (will integrate with `get_entity_relation` properly)
      // By default assume FK is parent_id on child
      where_clauses.push(format!("{}.parent_id = {}.id", base_alias, parent_alias));
    }

    // Wrap the object in the final array or object SELECT
    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 compile_inline_object(
    &self,
    props: &std::collections::BTreeMap<String, std::sync::Arc<crate::database::schema::Schema>>,
    parent_alias: &str,
    filter_keys: &[String],
  ) -> Result<(String, String), String> {
    let mut build_args = Vec::new();
    for (k, v) in props {
      let (child_sql, _) = self.walk_schema(v, parent_alias, Some(k), filter_keys)?;
      build_args.push(format!("'{}', {}", k, child_sql));
    }
    let combined = format!("jsonb_build_object({})", build_args.join(", "));
    Ok((combined, "object".to_string()))
  }
}