src/pipeline.rs

//! Query pipeline — orchestrates the full preprocessing → routing → generation flow.
//!
//! Architecture:
//! ```
//! User Query
//!   ├─ 1. Hashtag extraction      → [#rust, #make, #math]
//!   ├─ 2. Language detection       → is_russian? translate to English
//!   ├─ 3. KB cache lookup          → Hit? Return cached answer (instant)
//!   ├─ 4. Router (hashtag-aware)   → Select expert adapter
//!   └─ 5. Model inference          → Generate response
//!        └─ Post-process: annotate with metadata
//! ```

use std::path::PathBuf;
use std::time::Instant;

use crate::config::Config;
use crate::engine::Engine;
use crate::hashtags;
use crate::kb::{KbLookup, KnowledgeIndex};
use crate::router::Router;
use crate::translator::{self, Translation};

/// Result of a single pipeline run
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct PipelineResult {
    /// Generated (or cached) response text
    pub response: String,
    /// Which expert was used
    pub expert: String,
    /// Extracted hashtags
    pub hashtags: Vec<String>,
    /// Whether the response came from the KB cache
    pub from_cache: bool,
    /// Translation info (if applicable)
    pub translation: Option<Translation>,
    /// Timing breakdown (milliseconds)
    pub timing: PipelineTiming,
}

#[derive(Debug, Clone, Default)]
pub struct PipelineTiming {
    pub hashtag_ms: u64,
    pub translate_ms: u64,
    pub kb_lookup_ms: u64,
    pub routing_ms: u64,
    pub generation_ms: u64,
    pub total_ms: u64,
}

/// The pipeline orchestrator
pub struct Pipeline {
    /// Knowledge base index (optional)
    kb: Option<KnowledgeIndex>,
    /// Router for expert selection
    router: Router,
    /// Configuration
    config: Config,
    /// Whether to use KB cache
    use_cache: bool,
    /// Whether to translate Russian queries
    use_translate: bool,
}

#[allow(dead_code)]
impl Pipeline {
    /// Create a new pipeline.
    pub fn new(config: Config, kb_path: Option<PathBuf>) -> anyhow::Result<Self> {
        let expert_names: Vec<String> = config.experts.iter().map(|e| e.name.clone()).collect();
        let router = Router::new(&expert_names);

        let kb = if let Some(ref path) = kb_path {
            if path.exists() {
                Some(KnowledgeIndex::load(path)?)
            } else {
                tracing::warn!(
                    "Knowledge base file not found: {}. Running without cache.",
                    path.display()
                );
                None
            }
        } else {
            // Try default location
            let default_path = PathBuf::from("knowledge_base.json");
            if default_path.exists() {
                Some(KnowledgeIndex::load(&default_path)?)
            } else {
                tracing::info!("No knowledge base found — running without cache.");
                None
            }
        };

        Ok(Self {
            kb,
            router,
            config,
            use_cache: true,
            use_translate: true,
        })
    }

    /// Enable or disable KB caching.
    pub fn set_cache(&mut self, enabled: bool) {
        self.use_cache = enabled;
    }

    /// Enable or disable Russian→English translation.
    pub fn set_translate(&mut self, enabled: bool) {
        self.use_translate = enabled;
    }

    /// Check if KB is loaded.
    pub fn has_kb(&self) -> bool {
        self.kb.as_ref().map(|kb| !kb.is_empty()).unwrap_or(false)
    }

    /// KB entry count.
    pub fn kb_len(&self) -> usize {
        self.kb.as_ref().map(|kb| kb.len()).unwrap_or(0)
    }

    /// Pre-process a query: extract hashtags, detect language, translate.
    ///
    /// Returns the processed query (potentially translated) + metadata.
    pub fn preprocess(&self, query: &str) -> PreprocessResult {
        let t0 = Instant::now();

        // 1. Extract hashtags
        let hashtags = hashtags::extract_hashtags(query);
        let tag_names = hashtags::extract_tag_names(query);
        let hashtag_ms = t0.elapsed().as_millis() as u64;

        // 2. Detect language & translate
        let t1 = Instant::now();
        let translation = if self.use_translate {
            Some(translator::translate_ru_to_en(query, false))
        } else {
            None
        };
        let translate_ms = t1.elapsed().as_millis() as u64;

        // 3. Build the effective query for the model
        let effective_query = translation
            .as_ref()
            .map(|t| t.text.clone())
            .unwrap_or_else(|| query.to_string());

        let lang_tag = translation
            .as_ref()
            .map(|t| translator::language_tag(t).to_string())
            .unwrap_or_default();

        let translation_clone = translation.clone();

        PreprocessResult {
            original_query: query.to_string(),
            effective_query,
            hashtags,
            tag_names,
            translation: translation_clone,
            lang_tag,
            timing: PreprocessTiming {
                hashtag_ms,
                translate_ms,
            },
        }
    }

    /// Look up a preprocessed query in the knowledge base.
    pub fn kb_lookup(&self, pre: &PreprocessResult) -> (KbLookup, u64) {
        if !self.use_cache {
            return (KbLookup::Miss, 0);
        }

        let t0 = Instant::now();
        let result = self
            .kb
            .as_ref()
            .map(|kb| kb.lookup(&pre.original_query, &pre.effective_query, &pre.hashtags))
            .unwrap_or(KbLookup::Miss);
        let ms = t0.elapsed().as_millis() as u64;

        (result, ms)
    }

    /// Route to the best expert using hashtags + query text.
    pub fn route(&self, pre: &PreprocessResult) -> (String, u64) {
        let t0 = Instant::now();
        // Use the original query for routing (hashtag-aware)
        let expert = self
            .router
            .classify_with_tags(&pre.original_query, &pre.tag_names);
        let ms = t0.elapsed().as_millis() as u64;
        (expert, ms)
    }

    /// Get the system prompt for an expert, potentially augmented with KB context.
    pub fn build_system_prompt(
        &self,
        expert: &str,
        pre: &PreprocessResult,
        kb_result: &KbLookup,
    ) -> String {
        let mut system = String::new();

        // Base expert system prompt
        if let Some(expert_cfg) = self.config.get_expert(expert) {
            if let Some(sp) = &expert_cfg.system_prompt {
                system.push_str(sp);
            }
        }

        // Language tag
        if !pre.lang_tag.is_empty() {
            system.push_str(&format!(
                " [Note: user's original language is {}. Respond appropriately.]",
                pre.lang_tag.trim_start_matches('[').trim_end_matches(']')
            ));
        }

        // KB context for partial hits (truncated to avoid blowing up the prompt)
        if let KbLookup::Partial { answer_hint, .. } = kb_result {
            let truncated: String = answer_hint.chars().take(300).collect();
            let ellipsis = if answer_hint.len() > 300 { "..." } else { "" };
            system.push_str(&format!(
                " [Reference answer (adapt and improve, don't copy verbatim): {truncated}{ellipsis}]",
            ));
        }

        system
    }

    /// Full pipeline: preprocess → lookup → route → (cache hit or generate).
    ///
    /// This method requires an `Engine` for generation but can also return
    /// cached results without touching the model at all.
    pub fn run(
        &self,
        query: &str,
        engine: &mut Engine,
        expert_override: Option<&str>,
        history: &[ConversationTurn],
    ) -> anyhow::Result<PipelineResult> {
        let total_start = Instant::now();

        // --- Phase 1: Preprocess ---
        let pre = self.preprocess(query);
        tracing::info!(
            "Pipeline: hashtags={:?}, lang={}, translated={}",
            pre.hashtags,
            pre.translation
                .as_ref()
                .map(|t| t.original_lang.as_str())
                .unwrap_or("en"),
            pre.translation
                .as_ref()
                .map(|t| t.was_translated)
                .unwrap_or(false),
        );

        // --- Phase 2: KB cache lookup ---
        let (kb_result, kb_lookup_ms) = self.kb_lookup(&pre);

        // --- Phase 3: Cache hit? Return immediately (model-free, instant!) ---
        if let KbLookup::Hit {
            answer,
            entry_id,
            score,
        } = &kb_result
        {
            tracing::info!(
                "KB cache HIT: {entry_id} (score={score:.2}). Returning cached answer ({} chars).",
                answer.len()
            );
            let total_ms = total_start.elapsed().as_millis() as u64;
            return Ok(PipelineResult {
                response: answer.clone(),
                expert: "cache".to_string(),
                hashtags: pre.hashtags,
                from_cache: true,
                translation: pre.translation,
                timing: PipelineTiming {
                    hashtag_ms: pre.timing.hashtag_ms,
                    translate_ms: pre.timing.translate_ms,
                    kb_lookup_ms,
                    routing_ms: 0,
                    generation_ms: 0,
                    total_ms,
                },
            });
        }

        // --- Phase 4: Route ---
        let (auto_expert, routing_ms) = self.route(&pre);
        let expert = expert_override
            .map(|s| s.to_string())
            .unwrap_or(auto_expert);
        tracing::info!("Router → expert: {expert}");

        // --- Phase 5: Build prompt ---
        let system_prompt = self.build_system_prompt(&expert, &pre, &kb_result);
        let prompt = build_chatml_prompt(&system_prompt, &pre, history);

        // --- Phase 6: Generate with adapter ---
        let t_gen = Instant::now();
        let adapter_name: Option<String> = self
            .config
            .get_expert(&expert)
            .and_then(|cfg| cfg.adapter_file.as_ref())
            .map(|f| {
                f.trim_end_matches(".gguf")
                    .trim_end_matches(".safetensors")
                    .to_string()
            });

        let temperature = self.config.temperature as f32;
        let top_p = self.config.top_p as f32;
        let max_tokens = self.config.max_gen_tokens as u32;

        let response = engine.generate_with_adapter(
            &prompt,
            max_tokens,
            temperature,
            top_p,
            adapter_name.as_deref(),
        )?;
        let generation_ms = t_gen.elapsed().as_millis() as u64;

        let total_ms = total_start.elapsed().as_millis() as u64;

        tracing::info!(
            "Pipeline complete: hashtag={}µs translate={}µs kb={}µs route={}µs gen={}ms total={}ms",
            pre.timing.hashtag_ms * 1000,
            pre.timing.translate_ms * 1000,
            kb_lookup_ms * 1000,
            routing_ms * 1000,
            generation_ms,
            total_ms,
        );

        if let KbLookup::Partial {
            entry_id, score, ..
        } = &kb_result
        {
            tracing::info!("KB partial match: {entry_id} (score={score:.2}) — used as context.");
        }

        Ok(PipelineResult {
            response,
            expert,
            hashtags: pre.hashtags,
            from_cache: false,
            translation: pre.translation,
            timing: PipelineTiming {
                hashtag_ms: pre.timing.hashtag_ms,
                translate_ms: pre.timing.translate_ms,
                kb_lookup_ms,
                routing_ms,
                generation_ms,
                total_ms,
            },
        })
    }

    /// Get routing info for display
    pub fn routing_info(&self) -> String {
        self.router.routing_info()
    }
}

/// Result of the preprocessing phase
#[derive(Debug, Clone)]
pub struct PreprocessResult {
    pub original_query: String,
    pub effective_query: String,
    pub hashtags: Vec<String>,
    pub tag_names: Vec<String>,
    pub translation: Option<Translation>,
    pub lang_tag: String,
    pub timing: PreprocessTiming,
}

#[derive(Debug, Clone, Default)]
pub struct PreprocessTiming {
    pub hashtag_ms: u64,
    pub translate_ms: u64,
}

/// A single conversation turn (user query + assistant response)
#[derive(Debug, Clone)]
pub struct ConversationTurn {
    pub user: String,
    pub assistant: String,
}

/// Build a ChatML prompt with system prompt, conversation history, and current query.
///
/// Includes up to `max_history_turns` previous conversation turns
/// so the model maintains context across the conversation.
fn build_chatml_prompt(
    system: &str,
    pre: &PreprocessResult,
    history: &[ConversationTurn],
) -> String {
    let mut prompt = String::new();

    // System
    if !system.is_empty() {
        prompt.push_str(&format!("<|im_start|>system\n{system}<|im_end|>\n"));
    }

    // Conversation history — last N turns (6 turns = plenty of context at 4K)
    let max_turns = 6;
    let start = history.len().saturating_sub(max_turns);
    for turn in history[start..].iter() {
        prompt.push_str(&format!("<|im_start|>user\n{}<|im_end|>\n", turn.user));
        prompt.push_str(&format!(
            "<|im_start|>assistant\n{}<|im_end|>\n",
            turn.assistant
        ));
    }

    // Current user message — use effective (translated) query, optionally tag with metadata
    let user_msg = if pre
        .translation
        .as_ref()
        .map(|t| t.was_translated)
        .unwrap_or(false)
    {
        format!(
            "{}  [Tags: {} | Original (ru): {}]",
            pre.effective_query,
            pre.hashtags.join(" "),
            pre.original_query
        )
    } else if !pre.hashtags.is_empty() {
        format!(
            "{}  [Tags: {}]",
            pre.effective_query,
            pre.hashtags.join(" ")
        )
    } else {
        pre.effective_query.clone()
    };

    prompt.push_str(&format!("<|im_start|>user\n{user_msg}<|im_end|>\n"));
    prompt.push_str("<|im_start|>assistant\n");

    prompt
}