use async_trait::async_trait;
use motosan_index_core::core::error::Result;
use motosan_index_core::core::node::NodeWithScore;
use motosan_index_core::traits::post_processor::PostProcessor;

/// Keep only the top N nodes by score. Input is assumed to be sorted
/// by descending score (as returned by vector stores).
pub struct TopN {
    n: usize,
}

impl TopN {
    pub fn new(n: usize) -> Self {
        Self { n }
    }
}

#[async_trait]
impl PostProcessor for TopN {
    async fn process(&self, mut nodes: Vec<NodeWithScore>) -> Result<Vec<NodeWithScore>> {
        nodes.truncate(self.n);
        Ok(nodes)
    }
}
