gitlore/src/timeline/expand.rs

use std::collections::{HashSet, VecDeque};

use rusqlite::Connection;

use super::types::{EntityRef, ExpandedEntityRef, UnresolvedRef, resolve_entity_ref};
use crate::core::error::Result;

/// Result of the expand phase.
pub struct ExpandResult {
    pub expanded_entities: Vec<ExpandedEntityRef>,
    pub unresolved_references: Vec<UnresolvedRef>,
}

/// Run the EXPAND phase of the timeline pipeline (BFS over entity_references).
///
/// Starting from seed entities, traverses cross-references (both outgoing and incoming)
/// to discover related entities. Collects provenance (who referenced whom, how).
pub fn expand_timeline(
    conn: &Connection,
    seeds: &[EntityRef],
    depth: u32,
    include_mentions: bool,
    max_entities: usize,
) -> Result<ExpandResult> {
    if depth == 0 || seeds.is_empty() {
        return Ok(ExpandResult {
            expanded_entities: Vec::new(),
            unresolved_references: Vec::new(),
        });
    }

    let edge_types = if include_mentions {
        vec!["closes", "related", "mentioned"]
    } else {
        vec!["closes", "related"]
    };

    let mut visited: HashSet<(String, i64)> = seeds
        .iter()
        .map(|s| (s.entity_type.clone(), s.entity_id))
        .collect();

    let mut queue: VecDeque<(EntityRef, u32)> = seeds.iter().map(|s| (s.clone(), 0)).collect();

    let mut expanded = Vec::new();
    let mut unresolved = Vec::new();

    while let Some((current, current_depth)) = queue.pop_front() {
        if expanded.len() >= max_entities {
            break;
        }

        let neighbors = find_neighbors(conn, &current, &edge_types)?;

        for neighbor in neighbors {
            match neighbor {
                Neighbor::Resolved {
                    entity_ref,
                    reference_type,
                    source_method,
                } => {
                    let key = (entity_ref.entity_type.clone(), entity_ref.entity_id);
                    if !visited.insert(key) {
                        continue;
                    }

                    expanded.push(ExpandedEntityRef {
                        entity_ref: entity_ref.clone(),
                        depth: current_depth + 1,
                        via_from: current.clone(),
                        via_reference_type: reference_type,
                        via_source_method: source_method,
                    });

                    if expanded.len() >= max_entities {
                        break;
                    }

                    if current_depth + 1 < depth {
                        queue.push_back((entity_ref, current_depth + 1));
                    }
                }
                Neighbor::Unresolved(unresolved_ref) => {
                    unresolved.push(unresolved_ref);
                }
            }
        }
    }

    Ok(ExpandResult {
        expanded_entities: expanded,
        unresolved_references: unresolved,
    })
}

enum Neighbor {
    Resolved {
        entity_ref: EntityRef,
        reference_type: String,
        source_method: String,
    },
    Unresolved(UnresolvedRef),
}

/// Find all neighbors (outgoing + incoming) for an entity in entity_references.
fn find_neighbors(
    conn: &Connection,
    entity: &EntityRef,
    edge_types: &[&str],
) -> Result<Vec<Neighbor>> {
    let mut neighbors = Vec::new();

    find_outgoing(conn, entity, edge_types, &mut neighbors)?;
    find_incoming(conn, entity, edge_types, &mut neighbors)?;

    Ok(neighbors)
}

/// Find outgoing references: current entity is the source.
fn find_outgoing(
    conn: &Connection,
    entity: &EntityRef,
    edge_types: &[&str],
    neighbors: &mut Vec<Neighbor>,
) -> Result<()> {
    let placeholders: String = edge_types
        .iter()
        .enumerate()
        .map(|(i, _)| format!("?{}", i + 3))
        .collect::<Vec<_>>()
        .join(", ");

    let sql = format!(
        "SELECT target_entity_type, target_entity_id, target_project_path, target_entity_iid,
                reference_type, source_method
         FROM entity_references
         WHERE source_entity_type = ?1
           AND source_entity_id = ?2
           AND reference_type IN ({placeholders})"
    );

    let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
        Box::new(entity.entity_type.clone()),
        Box::new(entity.entity_id),
    ];
    for et in edge_types {
        params.push(Box::new(et.to_string()));
    }

    let params_refs: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();

    let mut stmt = conn.prepare(&sql)?;
    let rows = stmt.query_map(params_refs.as_slice(), |row| {
        Ok((
            row.get::<_, String>(0)?,         // target_entity_type
            row.get::<_, Option<i64>>(1)?,    // target_entity_id
            row.get::<_, Option<String>>(2)?, // target_project_path
            row.get::<_, Option<i64>>(3)?,    // target_entity_iid
            row.get::<_, String>(4)?,         // reference_type
            row.get::<_, String>(5)?,         // source_method
        ))
    })?;

    for row_result in rows {
        let (target_type, target_id, target_project_path, target_iid, ref_type, source_method) =
            row_result?;

        match target_id {
            Some(tid) => {
                if let Some(resolved) = resolve_entity_ref(conn, &target_type, tid, None)? {
                    neighbors.push(Neighbor::Resolved {
                        entity_ref: resolved,
                        reference_type: ref_type,
                        source_method,
                    });
                }
            }
            None => {
                neighbors.push(Neighbor::Unresolved(UnresolvedRef {
                    source: entity.clone(),
                    target_project: target_project_path,
                    target_type,
                    target_iid,
                    reference_type: ref_type,
                }));
            }
        }
    }

    Ok(())
}

/// Find incoming references: current entity is the target.
fn find_incoming(
    conn: &Connection,
    entity: &EntityRef,
    edge_types: &[&str],
    neighbors: &mut Vec<Neighbor>,
) -> Result<()> {
    let placeholders: String = edge_types
        .iter()
        .enumerate()
        .map(|(i, _)| format!("?{}", i + 3))
        .collect::<Vec<_>>()
        .join(", ");

    let sql = format!(
        "SELECT source_entity_type, source_entity_id, reference_type, source_method
         FROM entity_references
         WHERE target_entity_type = ?1
           AND target_entity_id = ?2
           AND reference_type IN ({placeholders})"
    );

    let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
        Box::new(entity.entity_type.clone()),
        Box::new(entity.entity_id),
    ];
    for et in edge_types {
        params.push(Box::new(et.to_string()));
    }

    let params_refs: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();

    let mut stmt = conn.prepare(&sql)?;
    let rows = stmt.query_map(params_refs.as_slice(), |row| {
        Ok((
            row.get::<_, String>(0)?, // source_entity_type
            row.get::<_, i64>(1)?,    // source_entity_id
            row.get::<_, String>(2)?, // reference_type
            row.get::<_, String>(3)?, // source_method
        ))
    })?;

    for row_result in rows {
        let (source_type, source_id, ref_type, source_method) = row_result?;

        if let Some(resolved) = resolve_entity_ref(conn, &source_type, source_id, None)? {
            neighbors.push(Neighbor::Resolved {
                entity_ref: resolved,
                reference_type: ref_type,
                source_method,
            });
        }
    }

    Ok(())
}

#[cfg(test)]
#[path = "timeline_expand_tests.rs"]
mod tests;