relay-knowledge 1.1.9

use rusqlite::{Connection, params, params_from_iter};

use crate::{
    domain::{
        CodeParseStatus, CodeParseStatusCounts, CodeRepositoryLatencySample, CodeRepositoryReport,
        CodeRepositoryTotals,
    },
    storage::StorageError,
};

use super::code_status;

pub(in crate::storage::sqlite) fn repository_totals(
    connection: &mut Connection,
) -> Result<CodeRepositoryTotals, StorageError> {
    let symbol_counts = symbol_generation_counts(connection)?;
    Ok(CodeRepositoryTotals {
        repository_count: count_all_rows(connection, "code_repositories")?,
        indexed_file_count: count_all_rows(connection, "code_repository_files")?,
        symbol_count: count_all_rows(connection, "code_repository_symbols")?,
        handwritten_symbol_count: symbol_counts.handwritten,
        generated_symbol_count: symbol_counts.generated,
        reference_count: count_all_rows(connection, "code_repository_references")?,
        chunk_count: count_all_rows(connection, "code_repository_chunks")?,
        degraded_file_count: count_all_rows(connection, "code_repository_file_diagnostics")?,
        parse_status_counts: repository_parse_status_counts(connection)?,
    })
}

pub(in crate::storage::sqlite) fn repository_totals_excluding(
    connection: &mut Connection,
    excluded_repository_ids: &[String],
) -> Result<CodeRepositoryTotals, StorageError> {
    if excluded_repository_ids.is_empty() {
        return repository_totals(connection);
    }

    let symbol_counts = symbol_generation_counts_excluding(connection, excluded_repository_ids)?;
    Ok(CodeRepositoryTotals {
        repository_count: count_rows_excluding(
            connection,
            "code_repositories",
            excluded_repository_ids,
        )?,
        indexed_file_count: count_rows_excluding(
            connection,
            "code_repository_files",
            excluded_repository_ids,
        )?,
        symbol_count: count_rows_excluding(
            connection,
            "code_repository_symbols",
            excluded_repository_ids,
        )?,
        handwritten_symbol_count: symbol_counts.handwritten,
        generated_symbol_count: symbol_counts.generated,
        reference_count: count_rows_excluding(
            connection,
            "code_repository_references",
            excluded_repository_ids,
        )?,
        chunk_count: count_rows_excluding(
            connection,
            "code_repository_chunks",
            excluded_repository_ids,
        )?,
        degraded_file_count: count_rows_excluding(
            connection,
            "code_repository_file_diagnostics",
            excluded_repository_ids,
        )?,
        parse_status_counts: repository_parse_status_counts_excluding(
            connection,
            excluded_repository_ids,
        )?,
    })
}

pub(super) fn repository_report(
    connection: &mut Connection,
    repository: &str,
) -> Result<CodeRepositoryReport, StorageError> {
    let status = code_status::repository_status(connection, repository)?.ok_or_else(|| {
        StorageError::InvalidInput(format!("code repository '{repository}' is not registered"))
    })?;
    let scope = status.last_indexed_scope_id.as_deref().unwrap_or_default();
    let degradation_summary = repository_diagnostics(connection, scope)?;
    let degraded_file_count = repository_degraded_file_count(connection, scope)?;
    let symbol_counts = scope_symbol_generation_counts(connection, scope)?;
    let edge_counts = repository_edge_resolution_counts(connection, scope)?;
    let representative_queries = representative_queries(connection, scope)?;
    let freshness_state = if status.stale {
        "stale"
    } else {
        status.state.as_str()
    }
    .to_owned();

    Ok(CodeRepositoryReport {
        repository_id: status.repository_id,
        alias: status.alias,
        root_path: status.root_path,
        path_filters: status.path_filters,
        language_filters: status.language_filters,
        resolved_commit_sha: status.last_indexed_commit,
        tree_hash: status.tree_hash,
        indexed_file_count: status.indexed_file_count,
        symbol_count: status.symbol_count,
        handwritten_symbol_count: symbol_counts.handwritten,
        generated_symbol_count: symbol_counts.generated,
        reference_count: status.reference_count,
        chunk_count: status.chunk_count,
        degraded_file_count,
        resolved_edge_count: edge_counts.resolved,
        ambiguous_edge_count: edge_counts.ambiguous,
        unresolved_edge_count: edge_counts.unresolved,
        degradation_summary,
        representative_queries,
        latency_samples: Vec::<CodeRepositoryLatencySample>::new(),
        freshness_state,
    })
}

#[derive(Debug, Default)]
pub(in crate::storage::sqlite) struct SymbolGenerationCounts {
    pub(in crate::storage::sqlite) handwritten: usize,
    pub(in crate::storage::sqlite) generated: usize,
}

pub(in crate::storage::sqlite) fn scope_symbol_generation_counts(
    connection: &Connection,
    source_scope: &str,
) -> Result<SymbolGenerationCounts, StorageError> {
    symbol_generation_counts_with_predicate(
        connection,
        "WHERE symbol.source_scope = ?1",
        params![source_scope],
    )
}

fn symbol_generation_counts(
    connection: &Connection,
) -> Result<SymbolGenerationCounts, StorageError> {
    symbol_generation_counts_with_predicate(connection, "", [])
}

fn symbol_generation_counts_excluding(
    connection: &Connection,
    excluded_repository_ids: &[String],
) -> Result<SymbolGenerationCounts, StorageError> {
    let placeholders = placeholders(excluded_repository_ids.len());
    symbol_generation_counts_with_predicate(
        connection,
        &format!("WHERE symbol.repository_id NOT IN ({placeholders})"),
        params_from_iter(excluded_repository_ids),
    )
}

fn symbol_generation_counts_with_predicate<P>(
    connection: &Connection,
    predicate_sql: &str,
    params: P,
) -> Result<SymbolGenerationCounts, StorageError>
where
    P: rusqlite::Params,
{
    connection
        .query_row(
            &format!(
                "
                SELECT
                    SUM(CASE WHEN file.is_generated != 0 THEN 0 ELSE 1 END),
                    SUM(CASE WHEN file.is_generated != 0 THEN 1 ELSE 0 END)
                FROM code_repository_symbols symbol
                INNER JOIN code_repository_files file
                  ON file.source_scope = symbol.source_scope
                 AND file.path = symbol.path
                {predicate_sql}
                "
            ),
            params,
            |row| {
                Ok(SymbolGenerationCounts {
                    handwritten: row.get::<_, Option<usize>>(0)?.unwrap_or_default(),
                    generated: row.get::<_, Option<usize>>(1)?.unwrap_or_default(),
                })
            },
        )
        .map_err(StorageError::from)
}

fn repository_parse_status_counts(
    connection: &Connection,
) -> Result<CodeParseStatusCounts, StorageError> {
    let mut statement = connection.prepare(
        "
        SELECT parse_status, COUNT(*)
        FROM code_repository_files
        GROUP BY parse_status
        ",
    )?;
    let rows = statement.query_map([], |row| {
        Ok((row.get::<_, String>(0)?, row.get::<_, usize>(1)?))
    })?;
    parse_status_counts_from_rows(rows)
}

fn parse_status_counts_from_rows(
    rows: rusqlite::MappedRows<
        '_,
        impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<(String, usize)>,
    >,
) -> Result<CodeParseStatusCounts, StorageError> {
    let mut counts = CodeParseStatusCounts::default();
    for row in rows {
        let (status, count) = row?;
        match status.as_str() {
            value if value == CodeParseStatus::Parsed.as_str() => counts.parsed = count,
            value if value == CodeParseStatus::Partial.as_str() => counts.partial = count,
            value if value == CodeParseStatus::TextOnly.as_str() => counts.text_only = count,
            value if value == CodeParseStatus::Failed.as_str() => counts.failed = count,
            other => {
                return Err(StorageError::InvalidInput(format!(
                    "unknown code repository parse status '{other}'"
                )));
            }
        }
    }

    Ok(counts)
}

fn repository_parse_status_counts_excluding(
    connection: &Connection,
    excluded_repository_ids: &[String],
) -> Result<CodeParseStatusCounts, StorageError> {
    let placeholders = placeholders(excluded_repository_ids.len());
    let mut statement = connection.prepare(&format!(
        "
        SELECT parse_status, COUNT(*)
        FROM code_repository_files
        WHERE repository_id NOT IN ({placeholders})
        GROUP BY parse_status
        "
    ))?;
    let rows = statement.query_map(params_from_iter(excluded_repository_ids), |row| {
        Ok((row.get::<_, String>(0)?, row.get::<_, usize>(1)?))
    })?;
    parse_status_counts_from_rows(rows)
}

fn repository_degraded_file_count(
    connection: &Connection,
    source_scope: &str,
) -> Result<usize, StorageError> {
    connection
        .query_row(
            "
            SELECT COUNT(*)
            FROM code_repository_file_diagnostics
            WHERE source_scope = ?1
            ",
            params![source_scope],
            |row| row.get::<_, usize>(0),
        )
        .map_err(StorageError::from)
}

#[derive(Debug, Default)]
struct EdgeResolutionCounts {
    resolved: usize,
    ambiguous: usize,
    unresolved: usize,
}

fn repository_edge_resolution_counts(
    connection: &Connection,
    source_scope: &str,
) -> Result<EdgeResolutionCounts, StorageError> {
    let mut counts = EdgeResolutionCounts::default();
    for (table, column) in [
        ("code_repository_references", "resolution_state"),
        ("code_repository_imports", "resolution_state"),
    ] {
        let mut statement = connection.prepare(&format!(
            "
            SELECT {column}, COUNT(*)
            FROM {table}
            WHERE source_scope = ?1
            GROUP BY {column}
            "
        ))?;
        let rows = statement.query_map(params![source_scope], |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, usize>(1)?))
        })?;
        for row in rows {
            let (state, count) = row?;
            match state.as_str() {
                "resolved" => counts.resolved += count,
                "ambiguous" => counts.ambiguous += count,
                _ => counts.unresolved += count,
            }
        }
    }

    Ok(counts)
}

fn repository_diagnostics(
    connection: &Connection,
    source_scope: &str,
) -> Result<Vec<String>, StorageError> {
    let mut statement = connection.prepare(
        "
        SELECT path, message
        FROM code_repository_file_diagnostics
        WHERE source_scope = ?1
        ORDER BY path ASC, message ASC
        LIMIT 20
        ",
    )?;
    let rows = statement.query_map(params![source_scope], |row| {
        Ok(format!(
            "{}: {}",
            row.get::<_, String>(0)?,
            row.get::<_, String>(1)?
        ))
    })?;

    rows.collect::<Result<Vec<_>, _>>()
        .map_err(StorageError::from)
}

fn representative_queries(
    connection: &Connection,
    source_scope: &str,
) -> Result<Vec<String>, StorageError> {
    let mut queries = Vec::new();
    let mut statement = connection.prepare(
        "
        SELECT symbol.name
        FROM code_repository_symbols symbol
        INNER JOIN code_repository_files file
          ON file.source_scope = symbol.source_scope
         AND file.path = symbol.path
        WHERE symbol.source_scope = ?1
          AND file.is_generated = 0
        ORDER BY symbol.path ASC, symbol.line_start ASC
        LIMIT 3
        ",
    )?;
    let rows = statement.query_map(params![source_scope], |row| row.get::<_, String>(0))?;
    queries.extend(
        rows.collect::<Result<Vec<_>, _>>()
            .map_err(StorageError::from)?,
    );
    if queries.is_empty() {
        queries.push("hybrid".to_owned());
    }
    queries.sort();
    queries.dedup();

    Ok(queries)
}

fn count_all_rows(connection: &Connection, table: &'static str) -> Result<usize, StorageError> {
    connection
        .query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |row| {
            row.get(0)
        })
        .map_err(StorageError::from)
}

fn count_rows_excluding(
    connection: &Connection,
    table: &'static str,
    excluded_repository_ids: &[String],
) -> Result<usize, StorageError> {
    let placeholders = placeholders(excluded_repository_ids.len());
    connection
        .query_row(
            &format!("SELECT COUNT(*) FROM {table} WHERE repository_id NOT IN ({placeholders})"),
            params_from_iter(excluded_repository_ids),
            |row| row.get(0),
        )
        .map_err(StorageError::from)
}

fn placeholders(count: usize) -> String {
    std::iter::repeat_n("?", count)
        .collect::<Vec<_>>()
        .join(", ")
}