komp_ac/graphs/src/lib.rs

use std::{env, net::SocketAddr};

use axum::{
    Form, Router,
    extract::State,
    http::{HeaderMap, HeaderValue, header},
    response::{Html, IntoResponse, Response},
    routing::{get, post},
};
mod pages;
mod services;
mod analytics {
    include!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../common/src/proto/komp_ac.analytics.rs"
    ));
}
mod auth {
    include!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../common/src/proto/komp_ac.auth.rs"
    ));
}
mod definitions {
    #[allow(dead_code)]
    pub mod common {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.common.rs"
        ));
    }
    pub mod table_definition {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.table_definition.rs"
        ));
    }
    pub mod table_structure {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.table_structure.rs"
        ));
    }
    pub mod table_script {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.table_script.rs"
        ));
    }
    pub mod table_validation {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.table_validation.rs"
        ));
    }
    pub mod tables_data {
        include!(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../common/src/proto/komp_ac.tables_data.rs"
        ));
    }
}

use analytics::{
    AnalyticsResultColumn, AnalyticsTable, AnalyticsValue, ExecuteAnalyticsQueryRequest,
    GetAnalyticsCatalogRequest, GetAnalyticsCatalogResponse,
    analytics_service_client::AnalyticsServiceClient, analytics_value,
};
use auth::{LoginRequest, auth_service_client::AuthServiceClient};
use definitions::{
    common::Empty,
    table_definition::table_definition_client::TableDefinitionClient,
    table_script::table_script_client::TableScriptClient,
    table_structure::table_structure_service_client::TableStructureServiceClient,
    table_validation::table_validation_service_client::TableValidationServiceClient,
    tables_data::tables_data_client::TablesDataClient,
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tonic::{Request, metadata::MetadataValue, transport::Channel};

const INDEX_HTML: &str = include_str!("../static/index.html");
const LOGIN_HTML: &str = include_str!("../static/login.html");

#[derive(Clone)]
pub(crate) struct AppState {
    analytics: AnalyticsServiceClient<Channel>,
    auth: AuthServiceClient<Channel>,
    definitions: TableDefinitionClient<Channel>,
    scripts: TableScriptClient<Channel>,
    structures: TableStructureServiceClient<Channel>,
    validations: TableValidationServiceClient<Channel>,
    tables_data: TablesDataClient<Channel>,
}

#[derive(Deserialize)]
struct QueryInput {
    profile_name: String,
    sql: String,
    #[serde(default = "default_max_rows")]
    max_rows: u32,
    chart_type: String,
}

#[derive(Deserialize)]
struct LoginInput {
    identifier: String,
    #[serde(default)]
    password: String,
}

#[derive(Deserialize)]
struct CatalogInput {
    profile_name: String,
}

fn default_max_rows() -> u32 {
    1_000
}

#[derive(Serialize)]
struct QueryOutput {
    columns: Vec<ColumnOutput>,
    rows: Vec<Vec<Value>>,
    row_count: u64,
    elapsed_ms: u64,
    truncated: bool,
}

#[derive(Serialize)]
struct ColumnOutput {
    name: String,
    data_type: String,
}

/// Starts the analytics web UI as a detached task on the current Tokio runtime.
///
/// The web UI still uses the existing gRPC endpoints for now. Keeping startup
/// here limits the server integration to a single call until the application
/// is switched to the in-process analytics runtime.
pub fn spawn() -> tokio::task::JoinHandle<()> {
    tokio::spawn(async {
        if let Err(error) = serve().await {
            eprintln!("Analytics graphs stopped: {error}");
        }
    })
}

/// Serves the analytics web UI until the task is cancelled or the listener fails.
pub async fn serve() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let grpc_endpoint =
        env::var("ANALYTICS_GRPC_ENDPOINT").unwrap_or_else(|_| "http://[::1]:50051".into());
    let listen_address = env::var("LISTEN_ADDRESS")
        .unwrap_or_else(|_| "127.0.0.1:3000".into())
        .parse::<SocketAddr>()?;

    let channel = Channel::from_shared(grpc_endpoint.clone())?.connect_lazy();
    let state = AppState {
        analytics: AnalyticsServiceClient::new(channel.clone()),
        auth: AuthServiceClient::new(channel.clone()),
        definitions: TableDefinitionClient::new(channel.clone()),
        scripts: TableScriptClient::new(channel.clone()),
        validations: TableValidationServiceClient::new(channel.clone()),
        tables_data: TablesDataClient::new(channel.clone()),
        structures: TableStructureServiceClient::new(channel),
    };

    let app = Router::new()
        .route("/", get(index))
        .route("/login", get(login_page).post(login))
        .route("/api/profiles", get(load_profiles))
        .route("/api/catalog", post(load_catalog))
        .route("/api/query", post(run_query))
        .merge(pages::admin::admin::router())
        .merge(pages::add_table::router())
        .merge(pages::add_logic::router())
        .merge(pages::add_validation::router())
        .merge(pages::import_export::router())
        .with_state(state);

    let listener = tokio::net::TcpListener::bind(listen_address).await?;
    println!("Analytics graphs: http://{listen_address}");
    println!("Analytics gRPC endpoint: {grpc_endpoint}");
    axum::serve(listener, app).await?;
    Ok(())
}

async fn index() -> Html<&'static str> {
    Html(INDEX_HTML)
}

async fn login_page() -> Html<&'static str> {
    Html(LOGIN_HTML)
}

async fn login(State(state): State<AppState>, Form(input): Form<LoginInput>) -> Response {
    if input.identifier.trim().is_empty() {
        return error_fragment("Username or email is required").into_response();
    }

    let mut client = state.auth;
    match client
        .login(Request::new(LoginRequest {
            identifier: input.identifier,
            password: input.password,
        }))
        .await
    {
        Ok(response) => {
            let login = response.into_inner();
            let cookie = format!(
                "analytics_token={}; Path=/; HttpOnly; SameSite=Strict; Max-Age={}",
                login.access_token, login.expires_in,
            );
            let Ok(cookie) = HeaderValue::try_from(cookie) else {
                return error_fragment("The server returned an invalid access token").into_response();
            };
            let mut response = Html(String::new()).into_response();
            response.headers_mut().insert(header::SET_COOKIE, cookie);
            response
                .headers_mut()
                .insert("hx-redirect", HeaderValue::from_static("/admin"));
            response
        }
        Err(error) => error_fragment(error.message()).into_response(),
    }
}

async fn run_query(
    State(state): State<AppState>,
    headers: HeaderMap,
    Form(input): Form<QueryInput>,
) -> Html<String> {
    if input.profile_name.trim().is_empty() || input.sql.trim().is_empty() {
        return error_fragment("Profile name and SQL are required");
    }

    let mut request = Request::new(ExecuteAnalyticsQueryRequest {
        profile_name: input.profile_name,
        sql: input.sql,
        max_rows: input.max_rows,
    });

    let Some(token) = cookie_value(&headers, "analytics_token") else {
        return Html("<p class=\"error\">Please <a href=\"/login\">log in</a> before running a query.</p>".into());
    };
    let Ok(value) = MetadataValue::try_from(format!("Bearer {token}")) else {
        return error_fragment("The stored bearer token is invalid");
    };
    request.metadata_mut().insert("authorization", value);

    let mut client = state.analytics;
    let mut stream = match client.execute_analytics_query(request).await {
        Ok(response) => response.into_inner(),
        Err(error) => return error_fragment(error.message()),
    };

    let mut columns = Vec::new();
    let mut rows = Vec::new();
    let mut row_count = 0;
    let mut elapsed_ms = 0;
    let mut truncated = false;

    loop {
        let batch = match stream.message().await {
            Ok(Some(batch)) => batch,
            Ok(None) => break,
            Err(error) => return error_fragment(error.message()),
        };
        if !batch.columns.is_empty() {
            columns = batch.columns.into_iter().map(column_output).collect();
        }
        rows.extend(batch.rows.into_iter().map(|row| {
            row.values.into_iter().map(value_output).collect::<Vec<_>>()
        }));

        if batch.is_final {
            row_count = batch.row_count;
            elapsed_ms = batch.elapsed_ms;
            truncated = batch.truncated;
        }
    }

    let output = QueryOutput {
        columns,
        rows,
        row_count,
        elapsed_ms,
        truncated,
    };
    render_result(&output, &input.chart_type)
}

async fn load_catalog(
    State(state): State<AppState>,
    headers: HeaderMap,
    Form(input): Form<CatalogInput>,
) -> Html<String> {
    if input.profile_name.trim().is_empty() {
        return error_fragment("Enter a profile name to load its schema");
    }
    let Some(token) = cookie_value(&headers, "analytics_token") else {
        return Html("<p class=\"error\">Please <a href=\"/login\">log in</a> before loading the schema.</p>".into());
    };
    let Ok(value) = MetadataValue::try_from(format!("Bearer {token}")) else {
        return error_fragment("The stored bearer token is invalid");
    };

    let mut request = Request::new(GetAnalyticsCatalogRequest {
        profile_name: input.profile_name,
    });
    request.metadata_mut().insert("authorization", value);

    let mut client = state.analytics;
    match client.get_analytics_catalog(request).await {
        Ok(response) => render_catalog(&response.into_inner()),
        Err(error) => error_fragment(error.message()),
    }
}

async fn load_profiles(State(state): State<AppState>) -> Html<String> {
    let mut client = state.definitions;
    match client.get_profile_tree(Request::new(Empty {})).await {
        Ok(response) => {
            let profiles = response.into_inner().profiles;
            if profiles.is_empty() {
                return Html("<option value=\"\">No profiles available</option>".into());
            }
            let options = profiles
                .iter()
                .map(|profile| {
                    format!(
                        "<option value=\"{}\">{} ({} tables)</option>",
                        escape_html(&profile.name),
                        escape_html(&profile.name),
                        profile.tables.len(),
                    )
                })
                .collect::<String>();
            Html(format!(
                "<option value=\"\">Choose a profile…</option>{options}"
            ))
        }
        Err(error) => Html(format!(
            "<option value=\"\">Could not load profiles: {}</option>",
            escape_html(error.message()),
        )),
    }
}

fn column_output(column: AnalyticsResultColumn) -> ColumnOutput {
    ColumnOutput {
        name: column.name,
        data_type: column.data_type,
    }
}

fn value_output(value: AnalyticsValue) -> Value {
    match value.kind {
        None | Some(analytics_value::Kind::NullValue(_)) => 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)) => {
            Value::String(value.iter().map(|byte| format!("{byte:02x}")).collect())
        }
    }
}

fn render_result(result: &QueryOutput, chart_type: &str) -> Html<String> {
    let metadata = format!(
        "{} rows in {} ms{}",
        result.row_count,
        result.elapsed_ms,
        if result.truncated { " (truncated)" } else { "" },
    );

    if chart_type == "table" {
        let headings = result
            .columns
            .iter()
            .map(|column| format!("<th>{}</th>", escape_html(&column.name)))
            .collect::<String>();
        let rows = result
            .rows
            .iter()
            .map(|row| {
                let cells = row
                    .iter()
                    .map(|value| format!("<td>{}</td>", escape_html(&display_value(value))))
                    .collect::<String>();
                format!("<tr>{cells}</tr>")
            })
            .collect::<String>();
        return Html(format!(
            "<p class=\"result-meta\">{metadata}</p><div class=\"table-wrap\"><table><thead><tr>{headings}</tr></thead><tbody>{rows}</tbody></table></div>"
        ));
    }

    if result.columns.len() < 2 {
        return error_fragment("Charts require at least two result columns");
    }

    let names = result
        .columns
        .iter()
        .map(|column| column.name.clone())
        .collect::<Vec<_>>();
    let option = match chart_type {
        "pie" => json!({
            "tooltip": { "trigger": "item" },
            "legend": { "type": "scroll", "bottom": 0 },
            "series": [{
                "type": "pie",
                "radius": ["25%", "68%"],
                "data": result.rows.iter().map(|row| json!({
                    "name": row.first().cloned().unwrap_or(Value::Null),
                    "value": row.get(1).cloned().unwrap_or(Value::Null),
                })).collect::<Vec<_>>()
            }]
        }),
        "scatter" => json!({
            "tooltip": { "trigger": "item" },
            "xAxis": { "name": names[0], "type": "value" },
            "yAxis": { "name": names[1], "type": "value" },
            "series": [{
                "name": names[1],
                "type": "scatter",
                "data": result.rows.iter().map(|row| json!([
                    row.first().cloned().unwrap_or(Value::Null),
                    row.get(1).cloned().unwrap_or(Value::Null),
                ])).collect::<Vec<_>>()
            }]
        }),
        "bar" | "line" => {
            let series = names
                .iter()
                .enumerate()
                .skip(1)
                .map(|(index, name)| json!({
                    "name": name,
                    "type": chart_type,
                    "smooth": chart_type == "line",
                    "data": result.rows.iter().map(|row| row.get(index).cloned().unwrap_or(Value::Null)).collect::<Vec<_>>(),
                }))
                .collect::<Vec<_>>();
            json!({
                "tooltip": { "trigger": "axis" },
                "legend": { "type": "scroll", "bottom": 0 },
                "grid": { "left": 55, "right": 25, "top": 35, "bottom": 70, "containLabel": true },
                "xAxis": {
                    "type": "category",
                    "data": result.rows.iter().map(|row| row.first().cloned().unwrap_or(Value::Null)).collect::<Vec<_>>()
                },
                "yAxis": { "type": "value" },
                "series": series,
            })
        }
        _ => return error_fragment("Unknown chart type"),
    };
    let option = escape_html(&option.to_string());

    Html(format!(
        "<p class=\"result-meta\">{metadata}</p>\
         <div class=\"chart\" data-option=\"{option}\" \
              x-init=\"echarts.init($el).setOption(JSON.parse($el.dataset.option))\"></div>"
    ))
}

fn render_catalog(catalog: &GetAnalyticsCatalogResponse) -> Html<String> {
    if catalog.tables.is_empty() {
        return Html("<p class=\"hint\">This profile has no analytics tables.</p>".into());
    }

    let tables = catalog
        .tables
        .iter()
        .map(render_catalog_table)
        .collect::<String>();
    let llm_context = build_llm_context(catalog);
    Html(format!(
        "<p class=\"schema-summary\"><strong>{}</strong> tables available</p>\
         <div class=\"catalog-tables\">{tables}</div>\
         <details class=\"llm-context\">\
           <summary>LLM schema context</summary>\
           <p class=\"hint\">Copy this entire text and include it when asking an LLM to write a query.</p>\
           <textarea x-ref=\"llmSchema\" readonly>{}</textarea>\
           <button type=\"button\" class=\"secondary\" \
                   x-on:click=\"navigator.clipboard.writeText($refs.llmSchema.value); copied = true\">Copy context</button>\
           <span class=\"hint\" x-show=\"copied\" x-cloak>Copied</span>\
         </details>",
        catalog.tables.len(),
        escape_html(&llm_context),
    ))
}

fn render_catalog_table(table: &AnalyticsTable) -> String {
    let table_name = quote_identifier(&table.name);
    let starter_query = format!("SELECT *\nFROM {table_name}\nLIMIT 100;");
    let 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));
            }
            format!(
                "<li><button type=\"button\" class=\"column-name\" data-insert=\"{}\" \
                    x-on:click=\"$refs.sql.setRangeText($el.dataset.insert, $refs.sql.selectionStart, $refs.sql.selectionEnd, 'end'); sql = $refs.sql.value; $refs.sql.focus()\">{}</button><span>{}</span></li>",
                escape_html(&quote_identifier(&column.name)),
                escape_html(&column.name),
                escape_html(&details),
            )
        })
        .collect::<String>();
    let links = if table.links.is_empty() {
        String::new()
    } else {
        let items = table
            .links
            .iter()
            .map(|link| {
                format!(
                    "<li><code>{}</code> → <code>{}</code>{}</li>",
                    escape_html(&link.source_column),
                    escape_html(&link.linked_table),
                    if link.required { " (required)" } else { "" },
                )
            })
            .collect::<String>();
        format!("<div class=\"links\"><span>Links</span><ul>{items}</ul></div>")
    };
    let currency = if table.base_currency.is_empty() {
        String::new()
    } else {
        format!(
            "<span class=\"currency\">{}</span>",
            escape_html(&table.base_currency)
        )
    };

    format!(
        "<details class=\"catalog-table\">\
           <summary><code>{}</code>{currency}</summary>\
           <button type=\"button\" class=\"starter-query\" data-sql=\"{}\" \
                   x-on:click=\"sql = $el.dataset.sql; $refs.sql.focus()\">Use starter query</button>\
           <ul class=\"columns\">{columns}</ul>{links}\
         </details>",
        escape_html(&table.name),
        escape_html(&starter_query),
    )
}

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('"', "\"\""))
}

fn display_value(value: &Value) -> String {
    match value {
        Value::Null => String::new(),
        Value::String(value) => value.clone(),
        value => value.to_string(),
    }
}

fn error_fragment(message: &str) -> Html<String> {
    Html(format!("<p class=\"error\">{}</p>", escape_html(message)))
}

fn escape_html(value: &str) -> String {
    value
        .replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&#39;")
}

fn cookie_value<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
    headers
        .get(header::COOKIE)?
        .to_str()
        .ok()?
        .split(';')
        .map(str::trim)
        .find_map(|cookie| cookie.strip_prefix(name)?.strip_prefix('='))
}