plexus-engine 0.3.4

use std::collections::BTreeSet;

use serde::Serialize;

use plexus_serde::{Op, Plan};

use crate::capabilities::{required_capabilities, EngineCapabilities, ExprKind, OpKind};

pub const EMBEDDED_GRAPH_RAG_PROFILE: &str = "embedded-graph-rag-v0";
pub const EMBEDDED_PHASE_1_RELEASE_CANDIDATE: &str = "embedded-phase-1";
pub const PLEXUS_CAPABILITY_REJECTION_SCENARIO: &str = "sr.plexus.capability-rejection.v1";

fn allowed_ops() -> BTreeSet<OpKind> {
    BTreeSet::from([
        OpKind::ConstRow,
        OpKind::Filter,
        OpKind::Project,
        OpKind::Limit,
        OpKind::VectorScan,
        OpKind::Rerank,
        OpKind::Return,
    ])
}

fn required_ops() -> BTreeSet<OpKind> {
    BTreeSet::from([OpKind::VectorScan, OpKind::Rerank, OpKind::Return])
}

fn allowed_exprs() -> BTreeSet<ExprKind> {
    BTreeSet::from([
        ExprKind::ColRef,
        ExprKind::PropAccess,
        ExprKind::IntLiteral,
        ExprKind::FloatLiteral,
        ExprKind::BoolLiteral,
        ExprKind::StringLiteral,
        ExprKind::NullLiteral,
        ExprKind::Cmp,
        ExprKind::And,
        ExprKind::Or,
        ExprKind::Not,
        ExprKind::IsNull,
        ExprKind::IsNotNull,
        ExprKind::StartsWith,
        ExprKind::EndsWith,
        ExprKind::Contains,
        ExprKind::ListLiteral,
        ExprKind::Arith,
        ExprKind::Param,
        ExprKind::Case,
        ExprKind::VectorSimilarity,
    ])
}

fn has_graph_ref(plan: &Plan) -> bool {
    plan.ops.iter().any(|op| match op {
        Op::ScanNodes { graph_ref, .. }
        | Op::Expand { graph_ref, .. }
        | Op::OptionalExpand { graph_ref, .. }
        | Op::ExpandVarLen { graph_ref, .. } => graph_ref.is_some(),
        _ => false,
    })
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct EmbeddedGraphRagAssessment {
    pub profile: String,
    pub supported: bool,
    pub required_ops: Vec<String>,
    pub allowed_ops: Vec<String>,
    pub allowed_exprs: Vec<String>,
    pub missing_required_ops: Vec<String>,
    pub unsupported_ops: Vec<String>,
    pub unsupported_exprs: Vec<String>,
    pub graph_ref_unsupported: bool,
}

impl EmbeddedGraphRagAssessment {
    pub fn rejection_summary(&self) -> String {
        let mut parts = Vec::new();
        if !self.missing_required_ops.is_empty() {
            parts.push(format!(
                "missing required ops [{}]",
                self.missing_required_ops.join(", ")
            ));
        }
        if !self.unsupported_ops.is_empty() {
            parts.push(format!(
                "unsupported ops [{}]",
                self.unsupported_ops.join(", ")
            ));
        }
        if !self.unsupported_exprs.is_empty() {
            parts.push(format!(
                "unsupported exprs [{}]",
                self.unsupported_exprs.join(", ")
            ));
        }
        if self.graph_ref_unsupported {
            parts.push("graph_ref is not supported".to_string());
        }
        parts.join("; ")
    }
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error(
    "plan is outside {profile}: {reason}",
    profile = .assessment.profile,
    reason = .assessment.rejection_summary()
)]
pub struct EmbeddedGraphRagProfileError {
    pub assessment: EmbeddedGraphRagAssessment,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct EmbeddedGraphRagCapabilityGap {
    pub profile: String,
    pub supported: bool,
    pub missing_ops: Vec<String>,
    pub missing_exprs: Vec<String>,
}

pub fn assess_embedded_graph_rag_plan(plan: &Plan) -> EmbeddedGraphRagAssessment {
    let required = required_capabilities(plan);
    let allowed_ops = allowed_ops();
    let allowed_exprs = allowed_exprs();
    let profile_required_ops = required_ops();

    let missing_required_ops: Vec<_> = profile_required_ops
        .difference(&required.required_ops)
        .copied()
        .map(OpKind::as_str)
        .map(str::to_string)
        .collect();
    let unsupported_ops: Vec<_> = required
        .required_ops
        .difference(&allowed_ops)
        .copied()
        .map(OpKind::as_str)
        .map(str::to_string)
        .collect();
    let unsupported_exprs: Vec<_> = required
        .required_exprs
        .difference(&allowed_exprs)
        .copied()
        .map(ExprKind::as_str)
        .map(str::to_string)
        .collect();
    let graph_ref_unsupported = has_graph_ref(plan);
    let supported = missing_required_ops.is_empty()
        && unsupported_ops.is_empty()
        && unsupported_exprs.is_empty()
        && !graph_ref_unsupported;

    EmbeddedGraphRagAssessment {
        profile: EMBEDDED_GRAPH_RAG_PROFILE.to_string(),
        supported,
        required_ops: profile_required_ops
            .iter()
            .copied()
            .map(OpKind::as_str)
            .map(str::to_string)
            .collect(),
        allowed_ops: allowed_ops
            .iter()
            .copied()
            .map(OpKind::as_str)
            .map(str::to_string)
            .collect(),
        allowed_exprs: allowed_exprs
            .iter()
            .copied()
            .map(ExprKind::as_str)
            .map(str::to_string)
            .collect(),
        missing_required_ops,
        unsupported_ops,
        unsupported_exprs,
        graph_ref_unsupported,
    }
}

#[allow(clippy::result_large_err)]
pub fn validate_embedded_graph_rag_plan(plan: &Plan) -> Result<(), EmbeddedGraphRagProfileError> {
    let assessment = assess_embedded_graph_rag_plan(plan);
    if assessment.supported {
        Ok(())
    } else {
        Err(EmbeddedGraphRagProfileError { assessment })
    }
}

pub fn assess_embedded_graph_rag_capabilities(
    capabilities: &EngineCapabilities,
) -> EmbeddedGraphRagCapabilityGap {
    let missing_ops: Vec<_> = required_ops()
        .difference(&capabilities.supported_ops)
        .copied()
        .map(OpKind::as_str)
        .map(str::to_string)
        .collect();
    let missing_exprs = if capabilities
        .supported_exprs
        .contains(&ExprKind::VectorSimilarity)
    {
        Vec::new()
    } else {
        vec![ExprKind::VectorSimilarity.as_str().to_string()]
    };

    EmbeddedGraphRagCapabilityGap {
        profile: EMBEDDED_GRAPH_RAG_PROFILE.to_string(),
        supported: missing_ops.is_empty() && missing_exprs.is_empty(),
        missing_ops,
        missing_exprs,
    }
}