komp_ac/web/src/pages/analytics/loader.rs

//! gRPC calls behind the analytics endpoints, plus the shaping of their
//! responses into the view models the templates render.

use axum::http::HeaderMap;

use crate::{
    AppState,
    analytics::{
        AnalyticsTable, ExecuteAnalyticsQueryRequest, GetAnalyticsCatalogRequest,
        GetAnalyticsCatalogResponse, analytics_value,
    },
    definitions::common::Empty,
    services::authenticated_request,
};

use super::state::{
    CatalogColumnView, CatalogLinkView, CatalogTableView, CatalogView, ColumnOutput, ProfileOption,
    QueryOutput,
};

pub(crate) enum LoadError {
    Unauthenticated,
    Backend(String),
}

pub(crate) async fn load_profiles(state: AppState) -> Result<Vec<ProfileOption>, String> {
    let mut definitions = state.definitions;
    let tree = definitions
        .get_profile_tree(tonic::Request::new(Empty {}))
        .await
        .map_err(|error| error.message().to_string())?
        .into_inner();
    Ok(tree
        .profiles
        .into_iter()
        .map(|profile| ProfileOption {
            table_count: profile.tables.len(),
            name: profile.name,
        })
        .collect())
}

pub(crate) async fn load_catalog(
    state: AppState,
    headers: &HeaderMap,
    profile_name: String,
) -> Result<CatalogView, LoadError> {
    let request = authenticated_request(headers, GetAnalyticsCatalogRequest { profile_name })
        .map_err(|_| LoadError::Unauthenticated)?;
    let mut client = state.analytics;
    let catalog = client
        .get_analytics_catalog(request)
        .await
        .map_err(|error| LoadError::Backend(error.message().to_string()))?
        .into_inner();
    Ok(catalog_view(&catalog))
}

pub(crate) async fn run_query(
    state: AppState,
    headers: &HeaderMap,
    profile_name: String,
    sql: String,
    max_rows: u32,
) -> Result<QueryOutput, LoadError> {
    let request = authenticated_request(
        headers,
        ExecuteAnalyticsQueryRequest {
            profile_name,
            sql,
            max_rows,
        },
    )
    .map_err(|_| LoadError::Unauthenticated)?;

    let mut client = state.analytics;
    let mut stream = client
        .execute_analytics_query(request)
        .await
        .map_err(|error| LoadError::Backend(error.message().to_string()))?
        .into_inner();

    let mut output = QueryOutput {
        columns: Vec::new(),
        rows: Vec::new(),
        row_count: 0,
        elapsed_ms: 0,
        truncated: false,
    };

    while let Some(batch) = stream
        .message()
        .await
        .map_err(|error| LoadError::Backend(error.message().to_string()))?
    {
        if !batch.columns.is_empty() {
            output.columns = batch
                .columns
                .into_iter()
                .map(|column| ColumnOutput {
                    name: column.name,
                    data_type: column.data_type,
                })
                .collect();
        }
        output.rows.extend(
            batch
                .rows
                .into_iter()
                .map(|row| row.values.into_iter().map(value_output).collect()),
        );
        if batch.is_final {
            output.row_count = batch.row_count;
            output.elapsed_ms = batch.elapsed_ms;
            output.truncated = batch.truncated;
        }
    }

    Ok(output)
}

fn value_output(value: crate::analytics::AnalyticsValue) -> serde_json::Value {
    use serde_json::json;
    match value.kind {
        None | Some(analytics_value::Kind::NullValue(_)) => serde_json::Value::Null,
        Some(analytics_value::Kind::BoolValue(value)) => json!(value),
        Some(analytics_value::Kind::Int64Value(value)) => json!(value),
        Some(analytics_value::Kind::Uint64Value(value)) => json!(value),
        Some(analytics_value::Kind::DoubleValue(value)) => json!(value),
        Some(analytics_value::Kind::StringValue(value)) => json!(value),
        Some(analytics_value::Kind::BytesValue(value)) => serde_json::Value::String(
            value.iter().map(|byte| format!("{byte:02x}")).collect(),
        ),
    }
}

fn catalog_view(catalog: &GetAnalyticsCatalogResponse) -> CatalogView {
    CatalogView {
        tables: catalog.tables.iter().map(catalog_table_view).collect(),
        llm_context: build_llm_context(catalog),
    }
}

fn catalog_table_view(table: &AnalyticsTable) -> CatalogTableView {
    CatalogTableView {
        name: table.name.clone(),
        base_currency: table.base_currency.clone(),
        starter_query: format!("SELECT *\nFROM {}\nLIMIT 100;", quote_identifier(&table.name)),
        columns: table
            .columns
            .iter()
            .map(|column| {
                let mut details = column.field_type.clone();
                if column.is_system {
                    details.push_str(", system");
                }
                if !column.rounding.is_empty() {
                    details.push_str(&format!(", rounding {}", column.rounding));
                }
                CatalogColumnView {
                    name: column.name.clone(),
                    insert_text: quote_identifier(&column.name),
                    details,
                }
            })
            .collect(),
        links: table
            .links
            .iter()
            .map(|link| CatalogLinkView {
                source_column: link.source_column.clone(),
                linked_table: link.linked_table.clone(),
                required: link.required,
            })
            .collect(),
    }
}

fn build_llm_context(catalog: &GetAnalyticsCatalogResponse) -> String {
    let mut text = format!(
        "Write one read-only PostgreSQL SELECT query for the komp_ac analytics API.\n\
Profile: {}\n\n\
Rules:\n\
- Return only the SQL query, without Markdown fences or explanation.\n\
- Use exactly one read-only SELECT statement.\n\
- Use only the public table and column aliases listed below.\n\
- Physical database column names are not visible through this API.\n\
- Add a reasonable LIMIT unless the query is an aggregate with a naturally small result.\n\
- For a chart, put the category or x-axis column first and numeric value columns after it.\n\n\
AVAILABLE ANALYTICS SCHEMA\n",
        catalog.profile_name,
    );

    for table in &catalog.tables {
        text.push_str(&format!("\nTABLE {}\n", quote_identifier(&table.name)));
        if !table.base_currency.is_empty() {
            text.push_str(&format!("  Base currency: {}\n", table.base_currency));
        }
        text.push_str("  Columns:\n");
        for column in &table.columns {
            text.push_str(&format!(
                "  - {}: {}",
                quote_identifier(&column.name),
                column.field_type,
            ));
            if column.is_system {
                text.push_str(" [system]");
            }
            if !column.rounding.is_empty() {
                text.push_str(&format!(" [rounding: {}]", column.rounding));
            }
            text.push('\n');
        }
        if !table.links.is_empty() {
            text.push_str("  Links:\n");
            for link in &table.links {
                text.push_str(&format!(
                    "  - {} references {} ({})\n",
                    quote_identifier(&link.source_column),
                    quote_identifier(&link.linked_table),
                    if link.required { "required" } else { "optional" },
                ));
            }
        }
    }
    text
}

fn quote_identifier(identifier: &str) -> String {
    format!("\"{}\"", identifier.replace('"', "\"\""))
}