query-engine/src/frontend/session.rs

//! Session state and command execution shared by the REPL and GUI.

use std::collections::HashMap;
use std::fmt;

use crate::catalog::{CatalogError, PredicateCatalog};
use crate::chase::{
    Atom, Instance, MaterializedState, Rule, Substitution, find_matches, materialize,
};
use crate::execution::execute;
use crate::planner::sql::plan_select;
use crate::relational::{DataType, Field, ResultSet, Schema};

use super::language::{Command, parse_script};
use super::provenance::explain_atom;

#[derive(Debug, Default)]
pub struct Session {
    base_instance: Instance,
    rules: Vec<Rule>,
    column_names: HashMap<String, Vec<String>>,
    materialized: Option<MaterializedState>,
}

impl Session {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn execute_script(&mut self, script: &str) -> Result<String, String> {
        let commands = parse_script(script)?;
        let mut output = Vec::new();

        for command in commands {
            let message = self.execute(command)?;
            if !message.is_empty() {
                output.push(message);
            }
        }

        if output.is_empty() {
            Ok("No commands executed.".to_string())
        } else {
            Ok(output.join("\n"))
        }
    }

    pub fn execute(&mut self, command: Command) -> Result<String, String> {
        match command {
            Command::Fact(atom) => {
                self.materialized = None;
                let inserted = self
                    .base_instance
                    .try_add(atom.clone())
                    .map_err(|err| err.to_string())?;
                let action = if inserted {
                    "Added"
                } else {
                    "Skipped duplicate"
                };
                Ok(format!("{} fact: {}", action, atom))
            }
            Command::Rule(rule) => {
                self.materialized = None;
                self.rules.push(rule.clone());
                Ok(format!("Added rule #{}: {}", self.rules.len(), rule))
            }
            Command::Schema { table, columns } => {
                self.column_names.insert(table.clone(), columns.clone());
                Ok(format!(
                    "Registered schema for {}: {}",
                    table,
                    columns.join(", ")
                ))
            }
            Command::Sql(select) => self.run_sql(&select),
            Command::Run => Ok(self.run_chase()),
            Command::Query(query) => Ok(self.run_query(&query)),
            Command::Explain(query) => Ok(self.explain_query(&query)),
            Command::ShowFacts => Ok(self.show_facts()),
            Command::ShowRules => Ok(self.show_rules()),
            Command::Reset => {
                *self = Self::default();
                Ok("Session reset.".to_string())
            }
            Command::Help => Ok(help_text().to_string()),
        }
    }

    pub fn reset(&mut self) {
        *self = Self::default();
    }

    fn run_chase(&mut self) -> String {
        let state = materialize(self.base_instance.clone(), &self.rules);
        let message = if state.result.terminated {
            format!(
                "Chase completed in {} step(s); {} fact(s) available.",
                state.result.steps,
                state.result.instance.len()
            )
        } else {
            format!(
                "Chase stopped after {} step(s); result may be incomplete.",
                state.result.steps
            )
        };
        self.materialized = Some(state);
        message
    }

    fn run_query(&self, query: &[Atom]) -> String {
        let instance = self.active_instance();
        let matches = find_matches(instance, query);
        let variables = query_variables(query);

        if variables.is_empty() {
            return if matches.is_empty() {
                "false".to_string()
            } else {
                "true".to_string()
            };
        }

        if matches.is_empty() {
            return "0 rows".to_string();
        }

        let mut rows = matches
            .iter()
            .map(|subst| format_substitution(subst, &variables))
            .collect::<Vec<_>>();
        rows.sort();

        let mut rendered = Vec::with_capacity(rows.len() + 1);
        rendered.push(format!("{} row(s)", rows.len()));
        rendered.extend(rows);
        rendered.join("\n")
    }

    fn run_sql(&self, select: &crate::sql::ast::Select) -> Result<String, String> {
        let instance = self.active_instance();
        let mut catalog =
            PredicateCatalog::from_instance(instance).map_err(|err| err.to_string())?;
        for (table, columns) in &self.column_names {
            match catalog.schema_for(table) {
                Ok(_) => {
                    catalog
                        .rename_columns(table, columns.clone())
                        .map_err(|err| err.to_string())?;
                }
                Err(CatalogError::UnknownTable(_)) => {
                    let schema = Schema::new(
                        columns
                            .iter()
                            .cloned()
                            .map(|name| Field::new(name, DataType::Text, false))
                            .collect(),
                    );
                    catalog.register_table(table.clone(), schema);
                }
                Err(err) => return Err(err.to_string()),
            }
        }
        let plan = plan_select(select, &catalog).map_err(|err| err.to_string())?;
        let result = execute(&plan, instance).map_err(|err| err.to_string())?;
        Ok(render_result_set(&result))
    }

    fn explain_query(&self, query: &[Atom]) -> String {
        let instance = self.active_instance();
        let matches = find_matches(instance, query);
        if matches.is_empty() {
            return if self.materialized.is_none() && !self.rules.is_empty() {
                "0 explanations. Run `run.` first to trace derived answers.".to_string()
            } else {
                "0 explanations".to_string()
            };
        }

        let variables = query_variables(query);
        let mut sections = Vec::new();

        for (index, subst) in matches.iter().enumerate() {
            let mut lines = Vec::new();
            lines.push(format!("match {}", index + 1));
            if !variables.is_empty() {
                lines.push(format!("  {}", format_substitution(subst, &variables)));
            }

            for atom in query.iter().map(|atom| subst.apply_atom(atom)) {
                lines.push(format!("  answer atom: {}", atom));
                if let Some(state) = &self.materialized {
                    for detail in explain_atom(&atom, state).lines().skip(1) {
                        lines.push(format!("  {}", detail));
                    }
                } else {
                    lines.push("    input fact".to_string());
                }
            }

            sections.push(lines.join("\n"));
        }

        let mut output = vec![format!("{} explanation(s)", sections.len())];
        output.extend(sections);
        output.join("\n")
    }

    fn show_facts(&self) -> String {
        let facts = sorted_render(self.active_instance().iter());
        if facts.is_empty() {
            return "No facts loaded.".to_string();
        }

        facts.join("\n")
    }

    fn show_rules(&self) -> String {
        if self.rules.is_empty() {
            return "No rules loaded.".to_string();
        }

        self.rules
            .iter()
            .enumerate()
            .map(|(index, rule)| format!("{}: {}", index + 1, rule))
            .collect::<Vec<_>>()
            .join("\n")
    }

    fn active_instance(&self) -> &Instance {
        if let Some(result) = &self.materialized {
            &result.result.instance
        } else {
            &self.base_instance
        }
    }
}

fn help_text() -> &'static str {
    "Commands:
fact Parent(alice, bob).
rule Parent(?X, ?Y) -> Ancestor(?X, ?Y).
schema Parent(parent, child).
sql SELECT * FROM Parent;
run.
query Ancestor(?X, ?Y)?
explain Ancestor(alice, bob)?
show facts
show rules
reset
help"
}

fn sorted_render<'a, T>(items: impl Iterator<Item = &'a T>) -> Vec<String>
where
    T: fmt::Display + 'a,
{
    let mut rendered = items.map(ToString::to_string).collect::<Vec<_>>();
    rendered.sort();
    rendered
}

fn render_result_set(result: &ResultSet) -> String {
    let mut lines = Vec::new();
    lines.push(format!("{} row(s)", result.rows().len()));

    if result.schema().is_empty() {
        return lines.join("\n");
    }

    let header = result
        .schema()
        .fields()
        .iter()
        .map(|field| field.name().to_string())
        .collect::<Vec<_>>()
        .join(" | ");
    lines.push(header);

    let rows = result
        .rows()
        .iter()
        .map(|row| {
            row.values()
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(" | ")
        })
        .collect::<Vec<_>>();
    lines.extend(rows);
    lines.join("\n")
}

fn query_variables(query: &[Atom]) -> Vec<String> {
    let mut variables = query
        .iter()
        .flat_map(|atom| atom.variables())
        .cloned()
        .collect::<Vec<_>>();
    variables.sort();
    variables.dedup();
    variables
}

fn format_substitution(subst: &Substitution, variables: &[String]) -> String {
    variables
        .iter()
        .filter_map(|var| subst.get(var).map(|term| format!("?{} = {}", var, term)))
        .collect::<Vec<_>>()
        .join(", ")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn session_runs_chase_and_query() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 rule Parent(?X, ?Y) -> Ancestor(?X, ?Y).\n\
                 rule Ancestor(?X, ?Y), Parent(?Y, ?Z) -> Ancestor(?X, ?Z).\n\
                 run.\n\
                 query Ancestor(?X, ?Y)?",
            )
            .unwrap();

        assert!(output.contains("Chase completed"));
        assert!(output.contains("?X = alice, ?Y = bob"));
        assert!(output.contains("?X = alice, ?Y = carol"));
    }

    #[test]
    fn boolean_query_returns_truth_value() {
        let mut session = Session::new();
        let output = session
            .execute_script("fact Parent(alice, bob).\nquery Parent(alice, bob)?")
            .unwrap();
        assert!(output.ends_with("true"));
    }

    #[test]
    fn explain_query_shows_rule_trace() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 rule Parent(?X, ?Y) -> Ancestor(?X, ?Y).\n\
                 rule Ancestor(?X, ?Y), Parent(?Y, ?Z) -> Ancestor(?X, ?Z).\n\
                 run.\n\
                 explain Ancestor(alice, carol)?",
            )
            .unwrap();

        assert!(output.contains("explanation(s)"));
        assert!(output.contains("derived by rule #2"));
        assert!(output.contains("premise: Ancestor(alice, bob)"));
        assert!(output.contains("input fact"));
    }

    #[test]
    fn session_runs_sql_query() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 sql SELECT c0 FROM Parent WHERE c1 = 'bob';",
            )
            .unwrap();

        assert!(output.contains("1 row(s)"));
        assert!(output.contains("c0"));
        assert!(output.contains("alice"));
    }

    #[test]
    fn session_runs_sql_query_with_alias_and_literal_projection() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 sql SELECT c0 AS parent_name, 'seed' AS label FROM Parent;",
            )
            .unwrap();

        assert!(output.contains("2 row(s)"));
        assert!(output.contains("parent_name | label"));
        assert!(output.contains("alice | seed"));
        assert!(output.contains("bob | seed"));
    }

    #[test]
    fn session_runs_sql_query_with_named_columns() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 schema Parent(parent, child).\n\
                 sql SELECT parent FROM Parent WHERE child = 'bob';",
            )
            .unwrap();

        assert!(output.contains("Registered schema for Parent: parent, child"));
        assert!(output.contains("1 row(s)"));
        assert!(output.contains("parent"));
        assert!(output.contains("alice"));
    }

    #[test]
    fn session_runs_sql_query_with_declared_empty_table() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "schema Parent(parent, child).\n\
                 sql SELECT parent FROM Parent;",
            )
            .unwrap();

        assert!(output.contains("Registered schema for Parent: parent, child"));
        assert!(output.contains("0 row(s)"));
        assert!(output.contains("parent"));
    }

    #[test]
    fn session_runs_sql_join_query() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 fact Ancestor(bob, carol).\n\
                 fact Ancestor(carol, dave).\n\
                 schema Parent(parent, child).\n\
                 schema Ancestor(parent, child).\n\
                 sql SELECT Parent.parent, Ancestor.child FROM Parent, Ancestor \
                     WHERE Parent.child = Ancestor.parent;",
            )
            .unwrap();

        assert!(output.contains("2 row(s)"));
        assert!(output.contains("Parent.parent | Ancestor.child"));
        assert!(output.contains("alice | carol"));
        assert!(output.contains("bob | dave"));
    }

    #[test]
    fn session_runs_sql_self_join_with_aliases() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 fact Parent(carol, dave).\n\
                 schema Parent(parent, child).\n\
                 sql SELECT p.parent, q.child FROM Parent AS p, Parent AS q \
                     WHERE p.child = q.parent;",
            )
            .unwrap();

        assert!(output.contains("2 row(s)"));
        assert!(output.contains("p.parent | q.child"));
        assert!(output.contains("alice | carol"));
        assert!(output.contains("bob | dave"));
    }

    #[test]
    fn session_runs_sql_query_with_single_table_alias() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 schema Parent(parent, child).\n\
                 sql SELECT p.parent FROM Parent AS p WHERE p.child = 'bob';",
            )
            .unwrap();

        assert!(output.contains("1 row(s)"));
        assert!(output.contains("p.parent"));
        assert!(output.contains("alice"));
    }

    #[test]
    fn session_runs_sql_query_with_and_filter() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 sql SELECT c0 FROM Parent WHERE c1 = 'bob' AND c0 = 'alice';",
            )
            .unwrap();

        assert!(output.contains("1 row(s)"));
        assert!(output.contains("alice"));
    }

    #[test]
    fn session_runs_sql_query_with_order_by() {
        let mut session = Session::new();
        let output = session
            .execute_script(
                "fact Parent(alice, bob).\n\
                 fact Parent(bob, carol).\n\
                 fact Parent(carol, dave).\n\
                 sql SELECT c0 FROM Parent ORDER BY c0 DESC;",
            )
            .unwrap();

        let lines = output.lines().collect::<Vec<_>>();
        let tail = &lines[lines.len() - 5..];
        assert_eq!(tail[0], "3 row(s)");
        assert_eq!(tail[1], "c0");
        assert_eq!(tail[2], "carol");
        assert_eq!(tail[3], "bob");
        assert_eq!(tail[4], "alice");
    }
}