collet 0.1.1

use std::sync::Arc;

mod alcove;
mod bridge;

pub use alcove::AlcoveRag;
pub use bridge::BridgeRag;

use crate::config::RagConfig;

/// A single retrieved document chunk.
pub struct RagChunk {
    pub adapter: &'static str,
    pub source: String,
    pub content: String,
    pub score: f32,
}

/// Trait for RAG backends.
#[async_trait::async_trait]
pub trait RagAdapter: Send + Sync {
    async fn retrieve(&self, query: &str, project: Option<&str>) -> Vec<RagChunk>;
    /// Adapter name, used in debug logging during dispatch.
    fn name(&self) -> &'static str;
}

/// Manages one or more RAG adapters and formats results for system prompt injection.
pub struct RagManager {
    adapters: Vec<Box<dyn RagAdapter>>,
    /// Approximate token budget for injected context (rough: 4 chars ≈ 1 token).
    max_chars: usize,
}

impl RagManager {
    pub fn from_config(cfg: &RagConfig) -> Option<Arc<Self>> {
        let mut adapters: Vec<Box<dyn RagAdapter>> = Vec::new();

        #[cfg(feature = "rag-alcove")]
        if cfg.alcove.as_ref().map(|a| a.enabled).unwrap_or(true)
            && let Some(adapter) = AlcoveRag::from_config(cfg)
        {
            adapters.push(Box::new(adapter));
        }

        if let Some(bridge_cfg) = &cfg.bridge
            && bridge_cfg.enabled
        {
            adapters.push(Box::new(BridgeRag::new(bridge_cfg)));
        }

        if adapters.is_empty() {
            return None;
        }

        let max_chars = cfg.max_tokens.unwrap_or(2000) * 4;
        Some(Arc::new(Self {
            adapters,
            max_chars,
        }))
    }

    /// Keywords that indicate constraint/requirement chunks deserving a score boost.
    const CONSTRAINT_KEYWORDS: &'static [&'static str] = &[
        "must not",
        "should not",
        "must be",
        "must have",
        "required",
        "constraint",
        "important",
        "critical",
        "mandatory",
        "prohibited",
        "forbidden",
        "never",
        "always",
        "ensure that",
        "do not",
    ];

    /// Retrieve chunks from all adapters, sorted by score, limited to `max_results`.
    pub async fn retrieve(
        &self,
        query: &str,
        project: Option<&str>,
        max_results: usize,
    ) -> Vec<RagChunk> {
        let mut chunks: Vec<RagChunk> = Vec::new();

        for adapter in &self.adapters {
            let before = chunks.len();
            chunks.extend(adapter.retrieve(query, project).await);
            let found = chunks.len() - before;
            if found > 0 {
                tracing::debug!(
                    adapter = adapter.name(),
                    found,
                    "RAG adapter returned chunks"
                );
            }
        }

        // Boost chunks containing constraint/requirement keywords so implicit
        // requirements (e.g. "must not", "forbidden") survive the token budget cutoff.
        for chunk in &mut chunks {
            let lower = chunk.content.to_lowercase();
            if Self::CONSTRAINT_KEYWORDS
                .iter()
                .any(|kw| lower.contains(kw))
            {
                chunk.score += 0.15;
            }
        }

        // Sort by score descending (re-sort after keyword boost)
        chunks.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Apply token budget and max_results limit
        let mut used = 0usize;
        chunks
            .into_iter()
            .take(max_results)
            .take_while(|chunk| {
                let len = chunk.content.len() + chunk.source.len() + 50;
                if used + len > self.max_chars {
                    return false;
                }
                used += len;
                true
            })
            .collect()
    }

    /// Retrieve from all adapters and format as a prompt section.
    pub async fn retrieve_and_format(&self, query: &str, project: Option<&str>) -> String {
        let chunks = self.retrieve(query, project, usize::MAX).await;

        if chunks.is_empty() {
            return String::new();
        }

        let mut out = String::from("## RAG Context\n\n");
        for chunk in chunks {
            let entry = format!(
                "[{}] {} (score: {:.2})\n> {}\n\n",
                chunk.adapter,
                chunk.source,
                chunk.score,
                chunk.content.replace('\n', "\n> ")
            );
            out.push_str(&entry);
        }

        out
    }
}