// src/tables_data/handlers/put_table_data.rs

use tonic::Status;
use sqlx::{PgPool, Arguments, Row};
use sqlx::postgres::PgArguments;
use chrono::{DateTime, Utc};
use common::proto::komp_ac::tables_data::{PutTableDataRequest, PutTableDataResponse};

use std::sync::Arc;
use prost_types::value::Kind;
use rust_decimal::Decimal;
use std::str::FromStr;
use std::collections::HashMap;

use crate::steel::server::execution::{self, Value};
use crate::indexer::{IndexCommand, IndexCommandData};
use tokio::sync::mpsc;
use tracing::error;

/// Updates table data with validation through associated scripts.
/// 
/// This function handles updating records in dynamically defined tables while:
/// - Validating permissions and column existence
/// - Executing validation scripts for computed columns
/// - Performing type-safe database updates
/// - Maintaining search index consistency
pub async fn put_table_data(
    db_pool: &PgPool,
    request: PutTableDataRequest,
    indexer_tx: &mpsc::Sender<IndexCommand>,
) -> Result<PutTableDataResponse, Status> {
    let profile_name = request.profile_name;
    let table_name = request.table_name;
    let record_id = request.id;

    if request.data.is_empty() {
        return Ok(PutTableDataResponse {
            success: true,
            message: "No fields to update.".into(),
            updated_id: record_id,
        });
    }

    // Fetch schema and table metadata
    let schema = sqlx::query!("SELECT id FROM schemas WHERE name = $1", profile_name)
        .fetch_optional(db_pool).await
        .map_err(|e| Status::internal(format!("Profile lookup error: {}", e)))?
        .ok_or_else(|| Status::not_found("Profile not found"))?;
    let schema_id = schema.id;

    let table_def = sqlx::query!(
        "SELECT id, columns FROM table_definitions WHERE schema_id = $1 AND table_name = $2",
        schema_id, table_name
    )
    .fetch_optional(db_pool).await
    .map_err(|e| Status::internal(format!("Table lookup error: {}", e)))?
    .ok_or_else(|| Status::not_found("Table not found"))?;

    // Parse column definitions from JSON format
    let columns_json: Vec<String> = serde_json::from_value(table_def.columns.clone())
        .map_err(|e| Status::internal(format!("Column parsing error: {}", e)))?;

    let mut columns = Vec::new();
    for col_def in columns_json {
        let parts: Vec<&str> = col_def.splitn(2, ' ').collect();
        if parts.len() != 2 {
            return Err(Status::internal("Invalid column format"));
        }
        let name = parts[0].trim_matches('"').to_string();
        let sql_type = parts[1].to_string();
        columns.push((name, sql_type));
    }

    // Build list of valid system columns (foreign keys and special columns)
    let fk_columns = sqlx::query!(
        r#"SELECT ltd.table_name
        FROM table_definition_links tdl
        JOIN table_definitions ltd ON tdl.linked_table_id = ltd.id
        WHERE tdl.source_table_id = $1"#,
        table_def.id
    )
    .fetch_all(db_pool)
    .await
    .map_err(|e| Status::internal(format!("Foreign key lookup error: {}", e)))?;

    let mut system_columns = vec!["deleted".to_string()];
    for fk in &fk_columns {
        system_columns.push(format!("{}_id", fk.table_name));
    }

    // Validate all requested columns exist and are editable
    let system_columns_set: std::collections::HashSet<_> = system_columns.iter().map(|s| s.as_str()).collect();
    let user_columns: Vec<&String> = columns.iter().map(|(name, _)| name).collect();

    for key in request.data.keys() {
        if !system_columns_set.contains(key.as_str()) &&
           !user_columns.contains(&&key.to_string()) {
            return Err(Status::invalid_argument(format!("Invalid column: {}", key)));
        }
    }

    // Fetch validation scripts for this table
    let scripts = sqlx::query!("SELECT id, target_column, script FROM table_scripts WHERE table_definitions_id = $1", table_def.id)
        .fetch_all(db_pool).await
        .map_err(|e| Status::internal(format!("Failed to fetch scripts: {}", e)))?;

    // Determine minimal set of columns needed for script execution context
    let mut required_columns = std::collections::HashSet::new();

    // Always need: id, target columns of scripts, and columns being updated
    required_columns.insert("id".to_string());
    for script_record in &scripts {
        required_columns.insert(script_record.target_column.clone());
    }
    for key in request.data.keys() {
        required_columns.insert(key.clone());
    }

    // Add columns referenced by scripts based on pre-computed dependencies
    if !scripts.is_empty() {
        let script_ids: Vec<i64> = scripts.iter().map(|s| s.id).collect();

        let dependencies = sqlx::query!(
            r#"SELECT sd.target_table_id, sd.dependency_type, sd.context_info, td.table_name as target_table
               FROM script_dependencies sd
               JOIN table_definitions td ON sd.target_table_id = td.id
               WHERE sd.script_id = ANY($1)"#,
            &script_ids
        )
        .fetch_all(db_pool)
        .await
        .map_err(|e| Status::internal(format!("Failed to fetch script dependencies: {}", e)))?;

        for dep in dependencies {
            if dep.target_table == table_name {
                // Script references this table - add specific columns it uses
                match dep.dependency_type.as_str() {
                    "column_access" | "indexed_access" => {
                        if let Some(context) = dep.context_info {
                            if let Some(column) = context.get("column").and_then(|v| v.as_str()) {
                                required_columns.insert(column.to_string());
                            }
                        }
                    }
                    _ => {} // SQL queries handled differently
                }
            } else {
                // Script references linked table - add foreign key column
                let fk_column = format!("{}_id", dep.target_table.split('_').last().unwrap_or(&dep.target_table));
                required_columns.insert(fk_column);
            }
        }
    }

    // Fetch current record state with only required columns for efficiency
    let qualified_table = crate::shared::schema_qualifier::qualify_table_name_for_data(db_pool, &profile_name, &table_name).await?;

    let columns_clause = required_columns
        .iter()
        .map(|name| format!("COALESCE(\"{0}\"::TEXT, '') AS \"{0}\"", name))
        .collect::<Vec<_>>()
        .join(", ");

    let select_sql = format!("SELECT {} FROM {} WHERE id = $1", columns_clause, qualified_table);

    let current_row = sqlx::query(&select_sql).bind(record_id).fetch_optional(db_pool).await
        .map_err(|e| Status::internal(format!("Failed to fetch current row state: {}", e)))?
        .ok_or_else(|| Status::not_found("Record not found"))?;

    let mut current_row_data = HashMap::new();
    for col_name in &required_columns {
        let value: String = current_row.try_get(col_name.as_str()).unwrap_or_default();
        current_row_data.insert(col_name.clone(), value);
    }

    // Extract and normalize update data from request
    let mut update_data = HashMap::new();
    for (key, proto_value) in &request.data {
        let str_val = match &proto_value.kind {
            Some(Kind::StringValue(s)) => s.trim().to_string(),
            Some(Kind::NumberValue(n)) => n.to_string(),
            Some(Kind::BoolValue(b)) => b.to_string(),
            Some(Kind::NullValue(_)) | None => String::new(),
            _ => return Err(Status::invalid_argument(format!("Unsupported type for column '{}'", key))),
        };
        update_data.insert(key.clone(), str_val);
    }

    // Create final context by merging current data with updates
    let mut final_context_data = current_row_data.clone();
    final_context_data.extend(update_data.clone());

    // Execute and validate scripts for computed columns
    for script_record in scripts {
        let target_column = script_record.target_column;

        // Determine SQL type for type-aware validation
        let target_sql_type = if let Some((_, stype)) = columns.iter().find(|(name, _)| name == &target_column) {
            stype.as_str()
        } else {
            "TEXT" // Default fallback for system columns
        };

        // Execute script with current context
        let script_result = execution::execute_script(
            script_record.script.clone(),
            "STRINGS",
            Arc::new(db_pool.clone()),
            schema_id,
            profile_name.clone(),
            table_name.clone(),
            final_context_data.clone(),
        )
        .await
        .map_err(|e| {
            error!("Script execution failed for '{}': {:?}", target_column, e);
            Status::invalid_argument(format!("Script execution failed for '{}': {}", target_column, e))
        })?;

        let Value::Strings(mut script_output_vec) = script_result;

        let script_output = script_output_vec.pop().ok_or_else(|| Status::internal("Script returned no values"))?;

        if update_data.contains_key(&target_column) {
            // Case A: Column is being updated - validate user input against script output
            let user_value = update_data.get(&target_column).unwrap();

            // Perform type-aware comparison based on SQL column type
            let values_match = match target_sql_type {
                s if s.starts_with("NUMERIC") => {
                    let user_decimal = Decimal::from_str(user_value).map_err(|_| Status::invalid_argument(format!("Invalid decimal format for column '{}'", target_column)))?;
                    let script_decimal = Decimal::from_str(&script_output).map_err(|_| Status::internal(format!("Script for '{}' produced invalid decimal", target_column)))?;
                    user_decimal == script_decimal
                },
                "INTEGER" | "BIGINT" => {
                    let user_int: i64 = user_value.parse().map_err(|_| Status::invalid_argument(format!("Invalid integer format for column '{}'", target_column)))?;
                    let script_int: i64 = script_output.parse().map_err(|_| Status::internal(format!("Script for '{}' produced invalid integer", target_column)))?;
                    user_int == script_int
                },
                "BOOLEAN" => {
                    let user_bool: bool = user_value.parse().map_err(|_| Status::invalid_argument(format!("Invalid boolean format for column '{}'", target_column)))?;
                    let script_bool: bool = script_output.parse().map_err(|_| Status::internal(format!("Script for '{}' produced invalid boolean", target_column)))?;
                    user_bool == script_bool
                },
                _ => user_value == &script_output
            };

            if !values_match {
                return Err(Status::invalid_argument(format!("Validation failed for column '{}': Script calculated '{}', but user provided '{}'", target_column, script_output, user_value)));
            }
        } else {
            // Case B: Column is NOT being updated - prevent unauthorized script-triggered changes
            let current_value = current_row_data.get(&target_column).cloned().unwrap_or_default();

            let values_match = match target_sql_type {
                s if s.starts_with("NUMERIC") => {
                    let current_decimal = Decimal::from_str(&current_value).unwrap_or_default();
                    let script_decimal = Decimal::from_str(&script_output).unwrap_or_default();
                    current_decimal == script_decimal
                },
                "INTEGER" | "BIGINT" => {
                    let current_int: i64 = current_value.parse().unwrap_or_default();
                    let script_int: i64 = script_output.parse().unwrap_or_default();
                    current_int == script_int
                },
                "BOOLEAN" => {
                    let current_bool: bool = current_value.parse().unwrap_or(false);
                    let script_bool: bool = script_output.parse().unwrap_or(false);
                    current_bool == script_bool
                },
                _ => current_value == script_output
            };

            if !values_match {
                return Err(Status::failed_precondition(format!("Script for column '{}' was triggered and would change its value from '{}' to '{}'. To apply this change, please include '{}' in your update request.", target_column, current_value, script_output, target_column)));
            }
        }
    }

    // Build parameterized UPDATE query with type-safe parameter binding
    let mut params = PgArguments::default();
    let mut set_clauses = Vec::new();
    let mut param_idx = 1;

    for (col, proto_value) in request.data {
        // Determine SQL type for proper parameter binding
        let sql_type = if system_columns_set.contains(col.as_str()) {
            match col.as_str() {
                "deleted" => "BOOLEAN",
                _ if col.ends_with("_id") => "BIGINT",
                _ => return Err(Status::invalid_argument("Invalid system column")),
            }
        } else {
            columns.iter()
                .find(|(name, _)| name == &col)
                .map(|(_, sql_type)| sql_type.as_str())
                .ok_or_else(|| Status::invalid_argument(format!("Column not found: {}", col)))?
        };

        // Handle NULL values first
        let kind = match &proto_value.kind {
            None | Some(Kind::NullValue(_)) => {
                match sql_type {
                    "BOOLEAN" => params.add(None::<bool>),
                    "TEXT" => params.add(None::<String>),
                    "TIMESTAMPTZ" => params.add(None::<DateTime<Utc>>),
                    "BIGINT" => params.add(None::<i64>),
                    "INTEGER" => params.add(None::<i32>),
                    s if s.starts_with("NUMERIC") => params.add(None::<Decimal>),
                    _ => return Err(Status::invalid_argument(format!("Unsupported type for null value: {}", sql_type))),
                }.map_err(|e| Status::internal(format!("Failed to add null parameter for {}: {}", col, e)))?;

                set_clauses.push(format!("\"{}\" = ${}", col, param_idx));
                param_idx += 1;
                continue;
            }
            Some(k) => k,
        };

        // Handle typed values with validation
        match sql_type {
            "TEXT" => {
                if let Kind::StringValue(value) = kind {
                    let trimmed_value = value.trim();
                    if trimmed_value.is_empty() {
                        params.add(None::<String>).map_err(|e| Status::internal(format!("Failed to add null parameter for {}: {}", col, e)))?;
                    } else {
                        // Apply business rule validation (e.g., phone number length)
                        if col == "telefon" && trimmed_value.len() > 15 {
                            return Err(Status::internal(format!("Value too long for {}", col)));
                        }
                        params.add(trimmed_value).map_err(|e| Status::invalid_argument(format!("Failed to add text parameter for {}: {}", col, e)))?;
                    }
                } else {
                    return Err(Status::invalid_argument(format!("Expected string for column '{}'", col)));
                }
            },
            "BOOLEAN" => {
                if let Kind::BoolValue(val) = kind {
                    params.add(val).map_err(|e| Status::invalid_argument(format!("Failed to add boolean parameter for {}: {}", col, e)))?;
                } else {
                    return Err(Status::invalid_argument(format!("Expected boolean for column '{}'", col)));
                }
            },
            "TIMESTAMPTZ" => {
                if let Kind::StringValue(value) = kind {
                    let dt = DateTime::parse_from_rfc3339(value).map_err(|_| Status::invalid_argument(format!("Invalid timestamp for {}", col)))?;
                    params.add(dt.with_timezone(&Utc)).map_err(|e| Status::invalid_argument(format!("Failed to add timestamp parameter for {}: {}", col, e)))?;
                } else {
                    return Err(Status::invalid_argument(format!("Expected ISO 8601 string for column '{}'", col)));
                }
            },
            "BIGINT" => {
                if let Kind::NumberValue(val) = kind {
                    if val.fract() != 0.0 {
                        return Err(Status::invalid_argument(format!("Expected integer for column '{}', but got a float", col)));
                    }
                    let as_i64 = *val as i64;
                    if (as_i64 as f64) != *val {
                        return Err(Status::invalid_argument(format!("Integer value out of range for BIGINT column '{}'", col)));
                    }
                    params.add(as_i64).map_err(|e| Status::invalid_argument(format!("Failed to add bigint parameter for {}: {}", col, e)))?;
                } else {
                    return Err(Status::invalid_argument(format!("Expected number for column '{}'", col)));
                }
            },
            "INTEGER" => {
                if let Kind::NumberValue(val) = kind {
                    if val.fract() != 0.0 {
                        return Err(Status::invalid_argument(format!("Expected integer for column '{}', but got a float", col)));
                    }
                    let as_i32 = *val as i32;
                    if (as_i32 as f64) != *val {
                        return Err(Status::invalid_argument(format!("Integer value out of range for INTEGER column '{}'", col)));
                    }
                    params.add(as_i32).map_err(|e| Status::invalid_argument(format!("Failed to add integer parameter for {}: {}", col, e)))?;
                } else {
                    return Err(Status::invalid_argument(format!("Expected number for column '{}'", col)));
                }
            },
            s if s.starts_with("NUMERIC") => {
                let decimal_val = match kind {
                    Kind::StringValue(s) => {
                        let trimmed = s.trim();
                        if trimmed.is_empty() {
                            None
                        } else {
                            Some(Decimal::from_str(trimmed).map_err(|_| {
                                Status::invalid_argument(format!(
                                    "Invalid decimal string format for column '{}': {}",
                                    col, s
                                ))
                            })?)
                        }
                    }
                    _ => {
                        return Err(Status::invalid_argument(format!(
                            "Expected a string representation for decimal column '{}', but received a different type.",
                            col
                        )));
                    }
                };
                params.add(decimal_val).map_err(|e| {
                    Status::invalid_argument(format!(
                        "Failed to add decimal parameter for {}: {}",
                        col, e
                    ))
                })?;
            },
            _ => return Err(Status::invalid_argument(format!("Unsupported type {}", sql_type))),
        }

        set_clauses.push(format!("\"{}\" = ${}", col, param_idx));
        param_idx += 1;
    }

    if set_clauses.is_empty() {
        return Ok(PutTableDataResponse {
            success: true,
            message: "No valid fields to update after processing.".into(),
            updated_id: record_id,
        });
    }

    // Execute the UPDATE query
    let set_clause = set_clauses.join(", ");
    let sql = format!(
        "UPDATE {} SET {} WHERE id = ${} RETURNING id",
        qualified_table,
        set_clause,
        param_idx
    );

    params.add(record_id).map_err(|e| Status::internal(format!("Failed to add record_id parameter: {}", e)))?;

    let result = sqlx::query_scalar_with::<_, i64, _>(&sql, params)
        .fetch_optional(db_pool)
        .await;

    let updated_id = match result {
        Ok(Some(id)) => id,
        Ok(None) => return Err(Status::not_found("Record not found")),
        Err(e) => {
            // Handle specific database errors with user-friendly messages
            if let Some(db_err) = e.as_database_error() {
                if db_err.code() == Some(std::borrow::Cow::Borrowed("22P02")) ||
                   db_err.code() == Some(std::borrow::Cow::Borrowed("22003")) {
                     return Err(Status::invalid_argument(format!(
                        "Numeric field overflow or invalid format. Check precision and scale. Details: {}", db_err.message()
                    )));
                }
            }
            return Err(Status::internal(format!("Update failed: {}", e)));
        }
    };

    // Queue record for search index update
    let command = IndexCommand::AddOrUpdate(IndexCommandData {
        table_name: table_name.clone(),
        row_id: updated_id,
    });

    if let Err(e) = indexer_tx.send(command).await {
        error!(
            "CRITICAL: DB update for table '{}' (id: {}) succeeded but failed to queue for indexing: {}. Search index is now inconsistent.",
            table_name, updated_id, e
        );
    }

    Ok(PutTableDataResponse {
        success: true,
        message: "Data updated successfully".into(),
        updated_id,
    })
}
