// src/steel/server/functions.rs

use steel::rvals::SteelVal;
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;
use sqlx::Row;

#[derive(Debug, Error)]
pub enum FunctionError {
    #[error("Column not found: {0}")]
    ColumnNotFound(String),
    #[error("Foreign key not found: {0}")]
    ForeignKeyNotFound(String),
    #[error("Table not found: {0}")]
    TableNotFound(String),
    #[error("Database error: {0}")]
    DatabaseError(String),
    #[error("Prohibited data type access: {0}")]
    ProhibitedTypeAccess(String),
}

/// Data types that Steel scripts are prohibited from accessing for security reasons
const PROHIBITED_TYPES: &[&str] = &["BIGINT", "DATE", "TIMESTAMPTZ"];

/// Execution context for Steel scripts with database access capabilities.
#[derive(Clone)]
pub struct SteelContext {
    pub current_table: String,
    pub schema_id: i64,
    pub schema_name: String,
    pub row_data: HashMap<String, String>,
    pub db_pool: Arc<PgPool>,
}

impl SteelContext {
    /// Resolves a base table name to its full qualified name in the current schema.
    /// Used for foreign key relationship traversal in Steel scripts.
    pub async fn get_related_table_name(&self, base_name: &str) -> Result<String, FunctionError> {
        let table_def = sqlx::query!(
            r#"SELECT table_name FROM table_definitions
               WHERE schema_id = $1 AND table_name LIKE $2"#,
            self.schema_id,
            format!("%_{}", base_name)
        )
        .fetch_optional(&*self.db_pool)
        .await
        .map_err(|e| FunctionError::DatabaseError(e.to_string()))?
        .ok_or_else(|| FunctionError::TableNotFound(base_name.to_string()))?;

        Ok(table_def.table_name)
    }

    /// Retrieves the SQL data type for a specific column in a table.
    /// Parses the JSON column definitions to find type information.
    async fn get_column_type(&self, table_name: &str, column_name: &str) -> Result<String, FunctionError> {
        let table_def = sqlx::query!(
            r#"SELECT columns FROM table_definitions
               WHERE schema_id = $1 AND table_name = $2"#,
            self.schema_id,
            table_name
        )
        .fetch_optional(&*self.db_pool)
        .await
        .map_err(|e| FunctionError::DatabaseError(e.to_string()))?
        .ok_or_else(|| FunctionError::TableNotFound(table_name.to_string()))?;

        let columns: Vec<String> = serde_json::from_value(table_def.columns)
            .map_err(|e| FunctionError::DatabaseError(format!("Invalid column data: {}", e)))?;

        // Parse column definitions to find the requested column type
        for column_def in columns {
            let mut parts = column_def.split_whitespace();
            if let (Some(name), Some(data_type)) = (parts.next(), parts.next()) {
                let column_name_clean = name.trim_matches('"');
                if column_name_clean == column_name {
                    return Ok(data_type.to_string());
                }
            }
        }

        Err(FunctionError::ColumnNotFound(format!(
            "Column '{}' not found in table '{}'",
            column_name,
            table_name
        )))
    }

    /// Converts database values to Steel script format based on column type.
    /// Currently handles boolean conversion to Steel's #true/#false syntax.
    fn convert_value_to_steel_format(&self, value: &str, column_type: &str) -> String {
        let normalized_type = normalize_data_type(column_type);

        match normalized_type.as_str() {
            "BOOLEAN" | "BOOL" => {
                // Convert database boolean representations to Steel boolean syntax
                match value.to_lowercase().as_str() {
                    "true" | "t" | "1" | "yes" | "on" => "#true".to_string(),
                    "false" | "f" | "0" | "no" | "off" => "#false".to_string(),
                    _ => value.to_string(), // Return as-is if not a recognized boolean
                }
            }
            "INTEGER" => value.to_string(),
            _ => value.to_string(), // Return as-is for other types
        }
    }

    /// Validates that a column type is allowed for Steel script access.
    /// Returns the column type if validation passes, error if prohibited.
    async fn validate_column_type_and_get_type(&self, table_name: &str, column_name: &str) -> Result<String, FunctionError> {
        let column_type = self.get_column_type(table_name, column_name).await?;

        if is_prohibited_type(&column_type) {
            return Err(FunctionError::ProhibitedTypeAccess(format!(
                "Cannot access column '{}' in table '{}' because it has prohibited type '{}'. Steel scripts cannot access columns of type: {}",
                column_name,
                table_name,
                column_type,
                PROHIBITED_TYPES.join(", ")
            )));
        }

        Ok(column_type)
    }

    /// Retrieves column value from current table or related tables via foreign keys.
    /// 
    /// # Behavior
    /// - Current table: Returns value directly from row_data with type conversion
    /// - Related table: Follows foreign key relationship and queries database
    /// - All accesses are subject to prohibited type validation
    pub fn steel_get_column(&self, table: &str, column: &str) -> Result<SteelVal, SteelVal> {
        if table == self.current_table {
            // Access current table data with type validation
            let column_type = tokio::task::block_in_place(|| {
                let handle = tokio::runtime::Handle::current();
                handle.block_on(async {
                    self.validate_column_type_and_get_type(table, column).await
                })
            });

            let column_type = match column_type {
                Ok(ct) => ct,
                Err(e) => return Err(SteelVal::StringV(e.to_string().into())),
            };

            return self.row_data.get(column)
                .map(|v| {
                    let converted_value = self.convert_value_to_steel_format(v, &column_type);
                    SteelVal::StringV(converted_value.into())
                })
                .ok_or_else(|| SteelVal::StringV(format!("Column {} not found", column).into()));
        }

        // Access related table via foreign key relationship
        let base_name = table.split_once('_')
            .map(|(_, rest)| rest)
            .unwrap_or(table);

        let fk_column = format!("{}_id", base_name);
        let fk_value = self.row_data.get(&fk_column)
            .ok_or_else(|| SteelVal::StringV(format!("Foreign key {} not found", fk_column).into()))?;

        let result = tokio::task::block_in_place(|| {
            let handle = tokio::runtime::Handle::current();
            handle.block_on(async {
                let actual_table = self.get_related_table_name(base_name).await
                    .map_err(|e| SteelVal::StringV(e.to_string().into()))?;

                // Validate column type and get type information
                let column_type = self.validate_column_type_and_get_type(&actual_table, column).await
                    .map_err(|e| SteelVal::StringV(e.to_string().into()))?;

                // Query the related table for the column value
                let raw_value = sqlx::query_scalar::<_, String>(
                    &format!("SELECT {} FROM \"{}\".\"{}\" WHERE id = $1", column, self.schema_name, actual_table)
                )
                .bind(fk_value.parse::<i64>().map_err(|_|
                    SteelVal::StringV("Invalid foreign key format".into()))?)
                .fetch_one(&*self.db_pool)
                .await
                .map_err(|e| SteelVal::StringV(e.to_string().into()))?;

                // Convert to appropriate Steel format
                let converted_value = self.convert_value_to_steel_format(&raw_value, &column_type);
                Ok(converted_value)
            })
        });

        result.map(|v| SteelVal::StringV(v.into()))
    }

    /// Retrieves a specific indexed element from a comma-separated column value.
    /// Useful for accessing elements from array-like string representations.
    pub fn steel_get_column_with_index(
        &self,
        table: &str,
        index: i64,
        column: &str
    ) -> Result<SteelVal, SteelVal> {
        // Get the full value with proper type conversion
        let value = self.steel_get_column(table, column)?;

        if let SteelVal::StringV(s) = value {
            let parts: Vec<_> = s.split(',').collect();

            if let Some(part) = parts.get(index as usize) {
                let trimmed_part = part.trim();

                // Apply type conversion to the indexed part based on original column type
                let column_type = tokio::task::block_in_place(|| {
                    let handle = tokio::runtime::Handle::current();
                    handle.block_on(async {
                        self.get_column_type(table, column).await
                    })
                });

                match column_type {
                    Ok(ct) => {
                        let converted_part = self.convert_value_to_steel_format(trimmed_part, &ct);
                        Ok(SteelVal::StringV(converted_part.into()))
                    }
                    Err(_) => {
                        // If type cannot be determined, return value as-is
                        Ok(SteelVal::StringV(trimmed_part.into()))
                    }
                }
            } else {
                Err(SteelVal::StringV("Index out of bounds".into()))
            }
        } else {
            Err(SteelVal::StringV("Expected comma-separated string".into()))
        }
    }

    /// Executes read-only SQL queries from Steel scripts with safety restrictions.
    /// 
    /// # Security Features
    /// - Only SELECT, SHOW, and EXPLAIN queries allowed
    /// - Prohibited column type access validation
    /// - Returns first column of all rows as comma-separated string
    pub fn steel_query_sql(&self, query: &str) -> Result<SteelVal, SteelVal> {
        if !is_read_only_query(query) {
            return Err(SteelVal::StringV(
                "Only SELECT queries are allowed".into()
            ));
        }

        if contains_prohibited_column_access(query) {
            return Err(SteelVal::StringV(format!(
                "SQL query may access prohibited column types. Steel scripts cannot access columns of type: {}",
                PROHIBITED_TYPES.join(", ")
            ).into()));
        }

        let pool = self.db_pool.clone();

        let result = tokio::task::block_in_place(|| {
            let handle = tokio::runtime::Handle::current();
            handle.block_on(async {
                let rows = sqlx::query(query)
                    .fetch_all(&*pool)
                    .await
                    .map_err(|e| SteelVal::StringV(e.to_string().into()))?;

                let mut results = Vec::new();
                for row in rows {
                    let val: String = row.try_get(0)
                        .map_err(|e| SteelVal::StringV(e.to_string().into()))?;
                    results.push(val);
                }

                Ok(results.join(","))
            })
        });

        result.map(|s| SteelVal::StringV(s.into()))
    }
}

/// Checks if a data type is prohibited for Steel script access.
fn is_prohibited_type(data_type: &str) -> bool {
    let normalized_type = normalize_data_type(data_type);
    PROHIBITED_TYPES.iter().any(|&prohibited| normalized_type.starts_with(prohibited))
}

/// Normalizes data type strings for consistent comparison.
/// Handles variations like NUMERIC(10,2) by extracting base type.
fn normalize_data_type(data_type: &str) -> String {
    data_type.to_uppercase()
        .split('(')  // Remove precision/scale from NUMERIC(x,y)
        .next()
        .unwrap_or(data_type)
        .trim()
        .to_string()
}

/// Performs basic heuristic check for prohibited column type access in SQL queries.
/// Looks for common patterns that might indicate access to restricted types.
fn contains_prohibited_column_access(query: &str) -> bool {
    let query_upper = query.to_uppercase();

    let patterns = [
        "EXTRACT(", // Common with DATE/TIMESTAMPTZ
        "DATE_PART(", // Common with DATE/TIMESTAMPTZ
        "::DATE",
        "::TIMESTAMPTZ",
        "::BIGINT",
    ];

    patterns.iter().any(|pattern| query_upper.contains(pattern))
}

/// Validates that a query is read-only and safe for Steel script execution.
fn is_read_only_query(query: &str) -> bool {
    let query = query.trim_start().to_uppercase();
    query.starts_with("SELECT") ||
    query.starts_with("SHOW") ||
    query.starts_with("EXPLAIN")
}

/// Converts row data boolean values to Steel script format during context initialization.
pub async fn convert_row_data_for_steel(
    db_pool: &PgPool,
    schema_id: i64,
    table_name: &str,
    row_data: &mut HashMap<String, String>,
) -> Result<(), sqlx::Error> {
    let table_def = sqlx::query!(
        r#"SELECT columns FROM table_definitions
           WHERE schema_id = $1 AND table_name = $2"#,
        schema_id,
        table_name
    )
    .fetch_optional(db_pool)
    .await?
    .ok_or_else(|| sqlx::Error::RowNotFound)?;

    // Parse column definitions to identify boolean columns for conversion
    if let Ok(columns) = serde_json::from_value::<Vec<String>>(table_def.columns) {
        for column_def in columns {
            let mut parts = column_def.split_whitespace();
            if let (Some(name), Some(data_type)) = (parts.next(), parts.next()) {
                let column_name = name.trim_matches('"');
                let normalized_type = normalize_data_type(data_type);

                if normalized_type == "BOOLEAN" || normalized_type == "BOOL" {
                    if let Some(value) = row_data.get_mut(column_name) {
                        // Convert boolean value to Steel format
                        *value = match value.to_lowercase().as_str() {
                            "true" | "t" | "1" | "yes" | "on" => "#true".to_string(),
                            "false" | "f" | "0" | "no" | "off" => "#false".to_string(),
                            _ => value.clone(), // Keep original if not recognized
                        };
                    }
                }
            }
        }
    }

    Ok(())
}
