// src/table_structure/handlers/table_structure.rs
use common::proto::komp_ac::table_structure::{
    GetTableStructureRequest, TableColumn, TableStructureResponse,
};
use sqlx::PgPool;
use tonic::Status;

// Helper struct to map query results
#[derive(sqlx::FromRow, Debug)]
struct DbColumnInfo {
    column_name: String,
    formatted_data_type: String,
    is_nullable: bool,
    is_primary_key: bool,
}

pub async fn get_table_structure(
    db_pool: &PgPool,
    request: GetTableStructureRequest,
) -> Result<TableStructureResponse, Status> {
    let profile_name = request.profile_name;
    let table_name = request.table_name;
    let table_schema = &profile_name;

    // 1. Validate 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!(
            "Failed to query profile '{}': {}",
            profile_name, e
        ))
    })?;

    let schema_id = match schema {
        Some(s) => s.id,
        None => {
            return Err(Status::not_found(format!(
                "Profile '{}' not found",
                profile_name
            )));
        }
    };

    // 2. Validate Table within Profile
    sqlx::query!(
        "SELECT id 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!("Failed to query table_definitions: {}", e)))?
    .ok_or_else(|| Status::not_found(format!(
        "Table '{}' not found in profile '{}'",
        table_name,
        profile_name
    )))?;

    // 3. Query information_schema for column details
    let query_str = r#"
        SELECT
            c.column_name,
            CASE
                WHEN c.udt_name = 'varchar' AND c.character_maximum_length IS NOT NULL THEN
                    'VARCHAR(' || c.character_maximum_length || ')'
                WHEN c.udt_name = 'bpchar' AND c.character_maximum_length IS NOT NULL THEN
                    'CHAR(' || c.character_maximum_length || ')'
                WHEN c.udt_name = 'numeric' AND c.numeric_precision IS NOT NULL AND c.numeric_scale IS NOT NULL THEN
                    'NUMERIC(' || c.numeric_precision || ',' || c.numeric_scale || ')'
                WHEN c.udt_name = 'numeric' AND c.numeric_precision IS NOT NULL THEN
                    'NUMERIC(' || c.numeric_precision || ')'
                WHEN STARTS_WITH(c.udt_name, '_') THEN
                    UPPER(SUBSTRING(c.udt_name FROM 2)) || '[]'
                ELSE
                    UPPER(c.udt_name)
            END AS formatted_data_type,
            c.is_nullable = 'YES' AS is_nullable,
            EXISTS (
                SELECT 1
                FROM information_schema.key_column_usage kcu
                JOIN information_schema.table_constraints tc
                ON kcu.constraint_name = tc.constraint_name
                AND kcu.table_schema = tc.table_schema
                AND kcu.table_name = tc.table_name
                WHERE tc.table_schema = c.table_schema
                  AND tc.table_name = c.table_name
                  AND tc.constraint_type = 'PRIMARY KEY'
                  AND kcu.column_name = c.column_name
            ) AS is_primary_key
        FROM
            information_schema.columns c
        WHERE
            c.table_schema = $1
            AND c.table_name = $2
        ORDER BY
            c.ordinal_position;
    "#;

    let db_columns = sqlx::query_as::<_, DbColumnInfo>(query_str)
        .bind(table_schema)
        .bind(&table_name) // Use the validated table_name
        .fetch_all(db_pool)
        .await
        .map_err(|e| {
            Status::internal(format!(
                "Failed to query column information for table '{}': {}",
                table_name, e
            ))
        })?;

    if db_columns.is_empty() {
        // This could mean the table exists in table_definitions but not in information_schema,
        // or it has no columns. The latter is unlikely for a created table.
        // Depending on desired behavior, you could return an error or an empty list.
        // For now, returning an empty list if the table was validated.
    }

    let columns = db_columns
        .into_iter()
        .map(|db_col| TableColumn {
            name: db_col.column_name,
            data_type: db_col.formatted_data_type,
            is_nullable: db_col.is_nullable,
            is_primary_key: db_col.is_primary_key,
        })
        .collect();

    Ok(TableStructureResponse { columns })
}
