foxhunt/database/src/query.rs

use crate::error::{DatabaseError, DatabaseResult, ErrorContext};
use serde_json::Value as JsonValue;
use sqlx::postgres::{PgArguments, PgRow};
use sqlx::{Arguments, Executor, FromRow, Postgres, Row};
use std::collections::HashMap;
use std::fmt;

/// Type-safe query builder for PostgreSQL
#[derive(Debug, Clone)]
pub struct QueryBuilder {
    query: String,
    args: PgArguments,
    prepared: bool,
}

impl QueryBuilder {
    /// Create a new query builder
    pub fn new() -> Self {
        Self {
            query: String::new(),
            args: PgArguments::default(),
            prepared: false,
        }
    }

    /// Create a SELECT query
    pub fn select<T: AsRef<str>>(columns: &[T]) -> SelectBuilder {
        let columns_str = if columns.is_empty() {
            "*".to_string()
        } else {
            columns.iter()
                .map(|c| c.as_ref())
                .collect::<Vec<_>>()
                .join(", ")
        };
        
        SelectBuilder {
            query: QueryBuilder::new(),
            columns: columns_str,
            from_clause: None,
            where_conditions: Vec::new(),
            order_by: Vec::new(),
            group_by: Vec::new(),
            having_conditions: Vec::new(),
            limit_clause: None,
            offset_clause: None,
            joins: Vec::new(),
        }
    }

    /// Create an INSERT query
    pub fn insert(table: &str) -> InsertBuilder {
        InsertBuilder {
            query: QueryBuilder::new(),
            table: table.to_string(),
            columns: Vec::new(),
            values: Vec::new(),
            on_conflict: None,
            returning: None,
        }
    }

    /// Create an UPDATE query
    pub fn update(table: &str) -> UpdateBuilder {
        UpdateBuilder {
            query: QueryBuilder::new(),
            table: table.to_string(),
            set_clauses: Vec::new(),
            where_conditions: Vec::new(),
            returning: None,
        }
    }

    /// Create a DELETE query
    pub fn delete(table: &str) -> DeleteBuilder {
        DeleteBuilder {
            query: QueryBuilder::new(),
            table: table.to_string(),
            where_conditions: Vec::new(),
            returning: None,
        }
    }

    /// Add a parameter to the query
    pub fn bind<T>(&mut self, value: T) -> &mut Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        self.args.add(value);
        self
    }

    /// Get the query string
    pub fn sql(&self) -> &str {
        &self.query
    }

    /// Get the arguments
    pub fn arguments(&self) -> &PgArguments {
        &self.args
    }

    /// Execute the query and return affected rows
    pub async fn execute<'e, E>(&self, executor: E) -> DatabaseResult<u64>
    where
        E: Executor<'e, Database = Postgres>,
    {
        let result = sqlx::query_with(&self.query, self.args.clone())
            .execute(executor)
            .await
            .with_query_context(&self.query)?;
        
        Ok(result.rows_affected())
    }

    /// Execute the query and fetch all rows
    pub async fn fetch_all<'e, E, T>(&self, executor: E) -> DatabaseResult<Vec<T>>
    where
        E: Executor<'e, Database = Postgres>,
        T: for<'r> FromRow<'r, PgRow> + Send + Unpin,
    {
        let rows = sqlx::query_as_with(&self.query, self.args.clone())
            .fetch_all(executor)
            .await
            .with_query_context(&self.query)?;
        
        Ok(rows)
    }

    /// Execute the query and fetch one row
    pub async fn fetch_one<'e, E, T>(&self, executor: E) -> DatabaseResult<T>
    where
        E: Executor<'e, Database = Postgres>,
        T: for<'r> FromRow<'r, PgRow> + Send + Unpin,
    {
        let row = sqlx::query_as_with(&self.query, self.args.clone())
            .fetch_one(executor)
            .await
            .with_query_context(&self.query)?;
        
        Ok(row)
    }

    /// Execute the query and fetch optional row
    pub async fn fetch_optional<'e, E, T>(&self, executor: E) -> DatabaseResult<Option<T>>
    where
        E: Executor<'e, Database = Postgres>,
        T: for<'r> FromRow<'r, PgRow> + Send + Unpin,
    {
        let row = sqlx::query_as_with(&self.query, self.args.clone())
            .fetch_optional(executor)
            .await
            .with_query_context(&self.query)?;
        
        Ok(row)
    }
}

impl Default for QueryBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// SELECT query builder
#[derive(Debug, Clone)]
pub struct SelectBuilder {
    query: QueryBuilder,
    columns: String,
    from_clause: Option<String>,
    where_conditions: Vec<String>,
    order_by: Vec<String>,
    group_by: Vec<String>,
    having_conditions: Vec<String>,
    limit_clause: Option<u64>,
    offset_clause: Option<u64>,
    joins: Vec<String>,
}

impl SelectBuilder {
    /// Add FROM clause
    pub fn from(mut self, table: &str) -> Self {
        self.from_clause = Some(table.to_string());
        self
    }

    /// Add WHERE condition
    pub fn where_eq<T>(mut self, column: &str, value: T) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        let placeholder = format!("${}", self.query.args.len() + 1);
        self.where_conditions.push(format!("{} = {}", column, placeholder));
        self.query.bind(value);
        self
    }

    /// Add WHERE IN condition
    pub fn where_in<T>(mut self, column: &str, values: Vec<T>) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        if values.is_empty() {
            return self;
        }

        let placeholders: Vec<String> = values
            .into_iter()
            .map(|value| {
                let placeholder = format!("${}", self.query.args.len() + 1);
                self.query.bind(value);
                placeholder
            })
            .collect();

        self.where_conditions.push(format!("{} IN ({})", column, placeholders.join(", ")));
        self
    }

    /// Add custom WHERE condition
    pub fn where_raw(mut self, condition: &str) -> Self {
        self.where_conditions.push(condition.to_string());
        self
    }

    /// Add INNER JOIN
    pub fn inner_join(mut self, table: &str, on: &str) -> Self {
        self.joins.push(format!("INNER JOIN {} ON {}", table, on));
        self
    }

    /// Add LEFT JOIN
    pub fn left_join(mut self, table: &str, on: &str) -> Self {
        self.joins.push(format!("LEFT JOIN {} ON {}", table, on));
        self
    }

    /// Add ORDER BY clause
    pub fn order_by(mut self, column: &str, direction: OrderDirection) -> Self {
        self.order_by.push(format!("{} {}", column, direction));
        self
    }

    /// Add GROUP BY clause
    pub fn group_by(mut self, column: &str) -> Self {
        self.group_by.push(column.to_string());
        self
    }

    /// Add HAVING condition
    pub fn having(mut self, condition: &str) -> Self {
        self.having_conditions.push(condition.to_string());
        self
    }

    /// Add LIMIT clause
    pub fn limit(mut self, count: u64) -> Self {
        self.limit_clause = Some(count);
        self
    }

    /// Add OFFSET clause
    pub fn offset(mut self, count: u64) -> Self {
        self.offset_clause = Some(count);
        self
    }

    /// Build the final query
    pub fn build(mut self) -> DatabaseResult<QueryBuilder> {
        let mut query_parts = vec![format!("SELECT {}", self.columns)];

        if let Some(from) = &self.from_clause {
            query_parts.push(format!("FROM {}", from));
        } else {
            return Err(DatabaseError::Query {
                query: "incomplete".to_string(),
                message: "SELECT query must have FROM clause".to_string(),
            });
        }

        // Add JOINs
        query_parts.extend(self.joins);

        // Add WHERE conditions
        if !self.where_conditions.is_empty() {
            query_parts.push(format!("WHERE {}", self.where_conditions.join(" AND ")));
        }

        // Add GROUP BY
        if !self.group_by.is_empty() {
            query_parts.push(format!("GROUP BY {}", self.group_by.join(", ")));
        }

        // Add HAVING
        if !self.having_conditions.is_empty() {
            query_parts.push(format!("HAVING {}", self.having_conditions.join(" AND ")));
        }

        // Add ORDER BY
        if !self.order_by.is_empty() {
            query_parts.push(format!("ORDER BY {}", self.order_by.join(", ")));
        }

        // Add LIMIT
        if let Some(limit) = self.limit_clause {
            query_parts.push(format!("LIMIT {}", limit));
        }

        // Add OFFSET
        if let Some(offset) = self.offset_clause {
            query_parts.push(format!("OFFSET {}", offset));
        }

        self.query.query = query_parts.join(" ");
        Ok(self.query)
    }
}

/// INSERT query builder
#[derive(Debug, Clone)]
pub struct InsertBuilder {
    query: QueryBuilder,
    table: String,
    columns: Vec<String>,
    values: Vec<Vec<String>>,
    on_conflict: Option<String>,
    returning: Option<String>,
}

impl InsertBuilder {
    /// Add a single row of values
    pub fn values<T>(mut self, values: &[(&str, T)]) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send + Clone,
    {
        if self.columns.is_empty() {
            self.columns = values.iter().map(|(col, _)| col.to_string()).collect();
        }

        let placeholders: Vec<String> = values
            .iter()
            .map(|(_, value)| {
                let placeholder = format!("${}", self.query.args.len() + 1);
                self.query.bind(value.clone());
                placeholder
            })
            .collect();

        self.values.push(placeholders);
        self
    }

    /// Add ON CONFLICT clause
    pub fn on_conflict(mut self, conflict: &str) -> Self {
        self.on_conflict = Some(conflict.to_string());
        self
    }

    /// Add RETURNING clause
    pub fn returning(mut self, columns: &str) -> Self {
        self.returning = Some(columns.to_string());
        self
    }

    /// Build the final query
    pub fn build(mut self) -> DatabaseResult<QueryBuilder> {
        if self.columns.is_empty() || self.values.is_empty() {
            return Err(DatabaseError::Query {
                query: "incomplete".to_string(),
                message: "INSERT query must have columns and values".to_string(),
            });
        }

        let mut query_parts = vec![
            format!("INSERT INTO {} ({})", self.table, self.columns.join(", ")),
            format!("VALUES {}", 
                self.values
                    .iter()
                    .map(|row| format!("({})", row.join(", ")))
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
        ];

        if let Some(conflict) = &self.on_conflict {
            query_parts.push(format!("ON CONFLICT {}", conflict));
        }

        if let Some(returning) = &self.returning {
            query_parts.push(format!("RETURNING {}", returning));
        }

        self.query.query = query_parts.join(" ");
        Ok(self.query)
    }
}

/// UPDATE query builder
#[derive(Debug, Clone)]
pub struct UpdateBuilder {
    query: QueryBuilder,
    table: String,
    set_clauses: Vec<String>,
    where_conditions: Vec<String>,
    returning: Option<String>,
}

impl UpdateBuilder {
    /// Add SET clause
    pub fn set<T>(mut self, column: &str, value: T) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        let placeholder = format!("${}", self.query.args.len() + 1);
        self.set_clauses.push(format!("{} = {}", column, placeholder));
        self.query.bind(value);
        self
    }

    /// Add WHERE condition
    pub fn where_eq<T>(mut self, column: &str, value: T) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        let placeholder = format!("${}", self.query.args.len() + 1);
        self.where_conditions.push(format!("{} = {}", column, placeholder));
        self.query.bind(value);
        self
    }

    /// Add RETURNING clause
    pub fn returning(mut self, columns: &str) -> Self {
        self.returning = Some(columns.to_string());
        self
    }

    /// Build the final query
    pub fn build(mut self) -> DatabaseResult<QueryBuilder> {
        if self.set_clauses.is_empty() {
            return Err(DatabaseError::Query {
                query: "incomplete".to_string(),
                message: "UPDATE query must have SET clauses".to_string(),
            });
        }

        let mut query_parts = vec![
            format!("UPDATE {}", self.table),
            format!("SET {}", self.set_clauses.join(", ")),
        ];

        if !self.where_conditions.is_empty() {
            query_parts.push(format!("WHERE {}", self.where_conditions.join(" AND ")));
        }

        if let Some(returning) = &self.returning {
            query_parts.push(format!("RETURNING {}", returning));
        }

        self.query.query = query_parts.join(" ");
        Ok(self.query)
    }
}

/// DELETE query builder
#[derive(Debug, Clone)]
pub struct DeleteBuilder {
    query: QueryBuilder,
    table: String,
    where_conditions: Vec<String>,
    returning: Option<String>,
}

impl DeleteBuilder {
    /// Add WHERE condition
    pub fn where_eq<T>(mut self, column: &str, value: T) -> Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        let placeholder = format!("${}", self.query.args.len() + 1);
        self.where_conditions.push(format!("{} = {}", column, placeholder));
        self.query.bind(value);
        self
    }

    /// Add custom WHERE condition
    pub fn where_raw(mut self, condition: &str) -> Self {
        self.where_conditions.push(condition.to_string());
        self
    }

    /// Add RETURNING clause
    pub fn returning(mut self, columns: &str) -> Self {
        self.returning = Some(columns.to_string());
        self
    }

    /// Build the final query
    pub fn build(mut self) -> DatabaseResult<QueryBuilder> {
        let mut query_parts = vec![format!("DELETE FROM {}", self.table)];

        if !self.where_conditions.is_empty() {
            query_parts.push(format!("WHERE {}", self.where_conditions.join(" AND ")));
        }

        if let Some(returning) = &self.returning {
            query_parts.push(format!("RETURNING {}", returning));
        }

        self.query.query = query_parts.join(" ");
        Ok(self.query)
    }
}

/// Order direction for ORDER BY clauses
#[derive(Debug, Clone, Copy)]
pub enum OrderDirection {
    Asc,
    Desc,
}

impl fmt::Display for OrderDirection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            OrderDirection::Asc => write!(f, "ASC"),
            OrderDirection::Desc => write!(f, "DESC"),
        }
    }
}

/// Helper for building dynamic WHERE conditions
#[derive(Debug, Clone)]
pub struct WhereBuilder {
    conditions: Vec<String>,
    args: PgArguments,
}

impl WhereBuilder {
    /// Create a new WHERE builder
    pub fn new() -> Self {
        Self {
            conditions: Vec::new(),
            args: PgArguments::default(),
        }
    }

    /// Add an equality condition
    pub fn eq<T>(&mut self, column: &str, value: T) -> &mut Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        let placeholder = format!("${}", self.args.len() + 1);
        self.conditions.push(format!("{} = {}", column, placeholder));
        self.args.add(value);
        self
    }

    /// Add an IN condition
    pub fn in_values<T>(&mut self, column: &str, values: Vec<T>) -> &mut Self
    where
        T: for<'q> sqlx::Encode<'q, Postgres> + sqlx::Type<Postgres> + Send,
    {
        if values.is_empty() {
            return self;
        }

        let placeholders: Vec<String> = values
            .into_iter()
            .map(|value| {
                let placeholder = format!("${}", self.args.len() + 1);
                self.args.add(value);
                placeholder
            })
            .collect();

        self.conditions.push(format!("{} IN ({})", column, placeholders.join(", ")));
        self
    }

    /// Add a custom condition
    pub fn custom(&mut self, condition: &str) -> &mut Self {
        self.conditions.push(condition.to_string());
        self
    }

    /// Build the WHERE clause
    pub fn build(self) -> (String, PgArguments) {
        let clause = if self.conditions.is_empty() {
            String::new()
        } else {
            format!("WHERE {}", self.conditions.join(" AND "))
        };
        (clause, self.args)
    }
}

impl Default for WhereBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_select_builder() {
        let query = QueryBuilder::select(&["id", "name"])
            .from("users")
            .where_eq("active", true)
            .order_by("name", OrderDirection::Asc)
            .limit(10)
            .build()
            .unwrap();

        assert!(query.sql().contains("SELECT id, name"));
        assert!(query.sql().contains("FROM users"));
        assert!(query.sql().contains("WHERE active = $1"));
        assert!(query.sql().contains("ORDER BY name ASC"));
        assert!(query.sql().contains("LIMIT 10"));
    }

    #[test]
    fn test_insert_builder() {
        let query = QueryBuilder::insert("users")
            .values(&[("name", "John"), ("email", "john@example.com")])
            .returning("id")
            .build()
            .unwrap();

        assert!(query.sql().contains("INSERT INTO users"));
        assert!(query.sql().contains("(name, email)"));
        assert!(query.sql().contains("VALUES ($1, $2)"));
        assert!(query.sql().contains("RETURNING id"));
    }

    #[test]
    fn test_update_builder() {
        let query = QueryBuilder::update("users")
            .set("name", "Jane")
            .set("updated_at", "NOW()")
            .where_eq("id", 1)
            .build()
            .unwrap();

        assert!(query.sql().contains("UPDATE users"));
        assert!(query.sql().contains("SET name = $1"));
        assert!(query.sql().contains("WHERE id = $"));
    }

    #[test]
    fn test_delete_builder() {
        let query = QueryBuilder::delete("users")
            .where_eq("id", 1)
            .returning("name")
            .build()
            .unwrap();

        assert!(query.sql().contains("DELETE FROM users"));
        assert!(query.sql().contains("WHERE id = $1"));
        assert!(query.sql().contains("RETURNING name"));
    }

    #[test]
    fn test_where_builder() {
        let mut where_builder = WhereBuilder::new();
        where_builder.eq("active", true);
        where_builder.in_values("status", vec!["active", "pending"]);
        where_builder.custom("created_at > NOW() - INTERVAL '1 day'");

        let (clause, _) = where_builder.build();
        assert!(clause.contains("WHERE"));
        assert!(clause.contains("active = $1"));
        assert!(clause.contains("status IN"));
        assert!(clause.contains("created_at >"));
    }
}