// src/tables_data/handlers/get_table_data.rs

use tonic::Status;
use sqlx::{PgPool, Row};
use std::collections::HashMap;
use common::proto::komp_ac::tables_data::{GetTableDataRequest, GetTableDataResponse};
use crate::shared::schema_qualifier::qualify_table_name_for_data;

pub async fn get_table_data(
    db_pool: &PgPool,
    request: GetTableDataRequest,
) -> Result<GetTableDataResponse, Status> {
    let profile_name = request.profile_name;
    let table_name = request.table_name;
    let record_id = request.id;

    // Lookup profile
    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)))?;

    let schema_id = schema.ok_or_else(|| Status::not_found("Profile not found"))?.id;

    // Lookup table_definition
    let table_def = sqlx::query!(
        r#"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)))?;

    let table_def = table_def.ok_or_else(|| Status::not_found("Table not found"))?;

    // Parse user-defined columns from JSON
    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 user_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();
        user_columns.push(name);
    }

    // --- START OF FIX ---

    // 1. Get all foreign key columns for this table
    let fk_columns_query = 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)))?;

    // 2. Build the list of foreign key column names using full table names
    let mut foreign_key_columns = Vec::new();
    for fk in fk_columns_query {
        // Use the full table name, not a stripped version
        foreign_key_columns.push(format!("{}_id", fk.table_name));
    }

    // 3. Prepare a complete list of all columns to select
    let mut all_column_names = vec!["id".to_string(), "deleted".to_string()];
    all_column_names.extend(user_columns);
    all_column_names.extend(foreign_key_columns);

    // 4. Build the SELECT clause with all columns
    let columns_clause = all_column_names
        .iter()
        .map(|name| format!("COALESCE(\"{0}\"::TEXT, '') AS \"{0}\"", name))
        .collect::<Vec<_>>()
        .join(", ");

    // --- END OF FIX ---

    // Qualify table name with schema
    let qualified_table = qualify_table_name_for_data(
        db_pool,
        &profile_name,
        &table_name,
    )
    .await?;

    let sql = format!(
        "SELECT {} FROM {} WHERE id = $1 AND deleted = false",
        columns_clause, qualified_table
    );

    // Execute query with enhanced error handling
    let row_result = sqlx::query(&sql)
        .bind(record_id)
        .fetch_one(db_pool)
        .await;

    let row = match row_result {
        Ok(row) => row,
        Err(sqlx::Error::RowNotFound) => return Err(Status::not_found("Record not found")),
        Err(e) => {
            if let Some(db_err) = e.as_database_error() {
                if db_err.code() == Some(std::borrow::Cow::Borrowed("42P01")) {
                    return Err(Status::internal(format!(
                        "Table '{}' is defined but does not physically exist in the database as {}",
                        table_name, qualified_table
                    )));
                }
            }
            return Err(Status::internal(format!("Database error: {}", e)));
        }
    };

    // Build response data from the complete list of columns
    let mut data = HashMap::new();
    for column_name in &all_column_names {
        let value: String = row
            .try_get(column_name.as_str())
            .map_err(|e| Status::internal(format!("Failed to get column {}: {}", column_name, e)))?;
        data.insert(column_name.clone(), value);
    }

    Ok(GetTableDataResponse { data })
}
