ragrig 0.9.3

//! Domain types shared across the crate: document representations,
//! CLI arguments, search results, and provider enums.

use clap::Parser;
use serde::{Deserialize, Serialize};
use serde_json;
use std::borrow::Borrow;
use std::path::PathBuf;

// --- Newtype wrappers --------------------------------------------------------

/// Reciprocal Rank Fusion score produced by hybrid search.
///
/// RRF scores are rank‑based (not similarity‑based) and typically
/// fall in the 0.0–0.03 range.  This newtype prevents accidental
/// comparison with cosine similarity values (0.0–1.0).
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
pub struct RrfScore(pub f64);

impl std::fmt::Display for RrfScore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:.4}", self.0)
    }
}

impl From<f64> for RrfScore {
    fn from(v: f64) -> Self {
        Self(v)
    }
}

impl PartialEq<f64> for RrfScore {
    fn eq(&self, other: &f64) -> bool {
        self.0 == *other
    }
}

impl PartialOrd<f64> for RrfScore {
    fn partial_cmp(&self, other: &f64) -> Option<std::cmp::Ordering> {
        self.0.partial_cmp(other)
    }
}

// --- Document Types ---
/// A document source filename, used consistently across the pipeline.
///
/// Replaces bare `String` fields named `source_file` or `file_name`
/// so consumers never have to guess which field name a struct uses.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default)]
pub struct SourceFile(pub String);

impl std::fmt::Display for SourceFile {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl From<&str> for SourceFile {
    fn from(s: &str) -> Self {
        Self(s.to_string())
    }
}

impl From<String> for SourceFile {
    fn from(s: String) -> Self {
        Self(s)
    }
}

impl Borrow<str> for SourceFile {
    fn borrow(&self) -> &str {
        &self.0
    }
}

impl PartialEq<str> for SourceFile {
    fn eq(&self, other: &str) -> bool {
        self.0 == other
    }
}

impl PartialEq<&str> for SourceFile {
    fn eq(&self, other: &&str) -> bool {
        self.0 == *other
    }
}

// --- Document Types ---

/// A PDF or EPUB file on disk.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum DocumentType {
    Pdf(PathBuf),
    Epub(PathBuf),
    Html(PathBuf),
    Docx(PathBuf),
    Markdown(PathBuf),
}

impl DocumentType {
    pub fn file_name(&self) -> &str {
        self.path()
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("unknown")
    }
    pub fn path(&self) -> &PathBuf {
        match self {
            Self::Pdf(p) => p,
            Self::Epub(p) => p,
            Self::Html(p) => p,
            Self::Docx(p) => p,
            Self::Markdown(p) => p,
        }
    }
}

impl DocumentType {
    /// Create a `DocumentType` from a file extension string.
    /// Returns `None` for unsupported extensions.
    pub fn from_extension(ext: &str, path: impl Into<PathBuf>) -> Option<Self> {
        let path = path.into();
        match ext.to_lowercase().as_str() {
            "pdf" => Some(Self::Pdf(path)),
            "epub" => Some(Self::Epub(path)),
            "html" | "htm" => Some(Self::Html(path)),
            "docx" => Some(Self::Docx(path)),
            "md" | "rmd" | "qmd" => Some(Self::Markdown(path)),
            _ => None,
        }
    }
}

/// Metadata for a document file: filename + SHA-256 hash.
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
pub struct FileHashEntry {
    pub file_name: SourceFile,
    pub hash: String,
}

/// A single text chunk from a document, tagged with its source file.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct DocumentChunk {
    pub text: String,
    pub source_file: SourceFile,
}

/// Configuration for text chunking — size and overlap in tokens.
///
/// This is the library-facing config struct.  It carries no CLI baggage
/// and can be constructed directly without importing `clap`.
#[derive(Clone, Debug)]
pub struct ChunkConfig {
    pub size: usize,
    pub overlap: usize,
}

impl Default for ChunkConfig {
    fn default() -> Self {
        Self {
            size: 1024,
            overlap: 128,
        }
    }
}

/// A paper result from academic search APIs (Semantic Scholar, arXiv).
#[derive(Deserialize, Debug, Clone)]
pub struct PaperResult {
    pub title: String,
    pub authors: Vec<String>,
    pub year: Option<i32>,
    pub arxiv_id: Option<String>,
    pub doi: Option<String>,
    pub pdf_url: Option<String>,
}

impl PaperResult {
    /// Short author list for display ("Smith, et al." if > 3).
    pub fn format_authors(&self) -> String {
        if self.authors.len() > 3 {
            format!("{}, et al.", self.authors[0])
        } else {
            self.authors.join(", ")
        }
    }

    /// " (2023)" or empty.
    pub fn format_year(&self) -> String {
        self.year.map(|y| format!(" ({})", y)).unwrap_or_default()
    }

    /// Best download URL: pdf_url, or arXiv fallback, or empty.
    pub fn best_pdf_url(&self) -> String {
        if let Some(ref url) = self.pdf_url {
            url.clone()
        } else if let Some(ref id) = self.arxiv_id {
            format!("https://arxiv.org/pdf/{}.pdf", id)
        } else {
            String::new()
        }
    }
}

/// Chat backend: local Ollama or cloud DeepSeek.
#[derive(Clone, Debug, clap::ValueEnum)]
pub enum Provider {
    Ollama,
    Deepseek,
}

/// Model generation parameters shared across all providers.
///
/// Each field is optional — when `None` the provider's own default is used.
/// Applied via rig-core's `AgentBuilder` methods (`temperature`, `max_tokens`)
/// or injected through `additional_params` for provider-specific knobs
/// (`top_p`, `seed`).
#[derive(Clone, Debug, Default)]
pub struct GenerationParams {
    /// Sampling temperature (0.0 = deterministic, 1.0+ = more random).
    pub temperature: Option<f64>,
    /// Nucleus sampling cutoff.
    pub top_p: Option<f64>,
    /// Maximum tokens to generate.
    pub max_tokens: Option<usize>,
    /// Random seed for reproducible output (not supported by all providers).
    pub seed: Option<u64>,
}

impl GenerationParams {
    /// Return `true` when every field is `None` (no overrides set).
    pub fn is_empty(&self) -> bool {
        self.temperature.is_none()
            && self.top_p.is_none()
            && self.max_tokens.is_none()
            && self.seed.is_none()
    }

    /// Build the `additional_params` JSON value for rig-core's
    /// `AgentBuilder::additional_params()`.  Returns `None` when
    /// `top_p` and `seed` are both absent, since `temperature` and
    /// `max_tokens` are passed via dedicated builder methods.
    pub fn additional_json(&self) -> Option<serde_json::Value> {
        if self.top_p.is_none() && self.seed.is_none() {
            return None;
        }
        let mut extra = serde_json::Map::new();
        if let Some(p) = self.top_p {
            if let Some(n) = serde_json::Number::from_f64(p) {
                extra.insert("top_p".into(), serde_json::Value::Number(n));
            }
        }
        if let Some(s) = self.seed {
            extra.insert("seed".into(), serde_json::Value::Number(s.into()));
        }
        Some(serde_json::Value::Object(extra))
    }
}

/// Embedding backend: local Ollama or CPU-only Fastembed.
#[derive(Clone, Debug, clap::ValueEnum)]
pub enum EmbeddingProvider {
    Ollama,
    #[cfg(feature = "internal-embed")]
    Fastembed,
}

/// PDF parser backend.
#[derive(Clone, Debug, PartialEq, clap::ValueEnum)]
pub enum PdfParserBackend {
    /// kreuzberg — docling-style layout-aware Markdown extraction
    #[cfg(feature = "kreuzberg")]
    Kreuzberg,
    /// unpdf — high-performance, direct Markdown output
    Unpdf,
    /// pdfsink-rs — structured, layout-aware
    #[deprecated(since = "0.10.0", note = "performance was lousy; use Unpdf instead")]
    Sink,
    /// pdf-extract — legacy flat-text (default)
    Extract,
    /// Binary scavenger — never panics
    Internal,
    /// Vision VLM — rasterise PDF pages and extract text via a vision-language model (requires Ollama with a vision model)
    Vision,
}

/// EPUB parser backend.
#[derive(Clone, Debug, clap::ValueEnum)]
pub enum EpubParserBackend {
    /// epub-parser crate
    Epub,
}

/// Controls how context-size overflow is handled.
#[derive(Clone, Copy, Debug, PartialEq, Eq, clap::ValueEnum)]
pub enum ContextSizeMode {
    /// Auto-adjust the context budget when the model reports overflow.
    Auto,
    /// Treat overflow as a fatal error — do not retry.
    Forced,
}

#[derive(Parser, Debug)]
#[command(about = "Pure Rust local RAG — chunkedrs + rig + Ollama/DeepSeek/Fastembed")]
pub struct Args {
    /// Folder containing PDF / EPUB / DOCX / HTML documents to index.
    #[arg(short, long)]
    pub folder: PathBuf,

    /// Chat backend: `ollama` (local) or `deepseek` (cloud API).
    /// Swappable at runtime via the `/chat` REPL command.
    #[arg(long, default_value = "ollama")]
    pub provider: Provider,

    /// DeepSeek API key (required when `--provider deepseek`).
    /// Can also be set via the `DEEPSEEK_API_KEY` env var.
    #[arg(long, env = "DEEPSEEK_API_KEY")]
    pub deepseek_api_key: Option<String>,

    /// Model name for DeepSeek (ignored when `--provider ollama`).
    #[arg(long, default_value = "deepseek-v4-pro")]
    pub deepseek_model: String,

    /// Semantic Scholar API key for higher rate limits (free).
    /// See <https://www.semanticscholar.org/product/api#api-key-form> for a key.
    #[arg(long, env = "SEMANTIC_SCHOLAR_API_KEY")]
    pub semantic_scholar_api_key: Option<String>,

    /// Model name for Ollama chat (ignored when `--provider deepseek`).
    #[arg(short, long, default_value = "gemma2:latest")]
    pub model: String,

    /// Sampling temperature for the chat model.  0.0 = deterministic.
    #[arg(long)]
    pub temperature: Option<f64>,

    /// Nucleus sampling cutoff for the chat model.
    #[arg(long)]
    pub top_p: Option<f64>,

    /// Maximum tokens to generate per response.
    #[arg(long)]
    pub max_tokens: Option<usize>,

    /// Random seed for reproducible chat output.
    #[arg(long)]
    pub seed: Option<u64>,

    /// Embedding backend. `ollama` uses the local Ollama server; `fastembed` runs
    /// Nomic-Embed-Text-v1.5 directly on the CPU with zero network overhead.
    #[arg(long, default_value = "ollama")]
    pub embedding_provider: EmbeddingProvider,

    /// Model name passed to Ollama when `--embedding-provider ollama` (ignored for fastembed).
    #[arg(short, long, default_value = "nomic-embed-text")]
    pub embedding_model: String,

    /// Model used for conversational query expansion and memory (via the
    /// `Generator` trait — any backend works).  Defaults to a small local
    /// model.  Swappable at runtime via `/memory`.
    #[arg(long, default_value = "qwen2.5:1.5b")]
    pub memory_model: String,

    /// Path to a Markdown file containing a custom system prompt for the
    /// chat agent.  Use `{context}` as placeholder for retrieved documents.
    #[arg(long)]
    pub prompt_chat: Option<PathBuf>,

    /// Path to a Markdown file containing a custom system prompt for the
    /// memory / rewrite agent.  Use `{question}` as placeholder.
    #[arg(long)]
    pub prompt_rewrite: Option<PathBuf>,

    /// Use the sloppy PDF parser as fallback.  Never panics — scavenges
    /// raw text strings from the PDF binary when structured parsers fail.
    #[arg(long)]
    pub sloppy_pdf: bool,

    /// PDF parser backend: kreuzberg (docling-style, feature-gated),
    /// unpdf (Markdown-native), sink (structured), extract (legacy flat-text, default),
    /// or internal (sloppy binary scavenger, never panics).
    #[arg(long, default_value = "extract")]
    pub pdf_parser: PdfParserBackend,

    #[arg(long, default_value = "1024")]
    pub chunk_size: usize,

    #[arg(long, default_value = "128")]
    pub chunk_overlap: usize,

    #[arg(long, default_value = "50")]
    pub top_k: usize,

    /// Minimum cosine similarity for a chunk to participate in hybrid search.
    /// Can be toggled at runtime via `/embed threshold <F>`.
    #[arg(long, default_value = "0.04")]
    pub similarity_threshold: f64,

    /// Context window size of the current chat model (tokens).
    /// Used to calculate how much retrieved text fits in the prompt.
    /// Override per-model at runtime via `/chat context <N>`.
    #[arg(long, default_value = "8192")]
    pub model_ctx_tokens: usize,

    /// Controls context-overflow behaviour: `auto` retries with fewer chunks,
    /// `forced` treats overflow as a fatal error.
    #[arg(long, default_value = "auto")]
    pub context_size_forced: ContextSizeMode,
}

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

    #[test]
    fn document_type_file_name() {
        let dt = DocumentType::Pdf(PathBuf::from("/tmp/paper.pdf"));
        assert_eq!(dt.file_name(), "paper.pdf");
        assert_eq!(dt.path(), &PathBuf::from("/tmp/paper.pdf"));
    }

    #[test]
    fn document_type_file_name_unknown() {
        let dt = DocumentType::Epub(PathBuf::from("/"));
        assert_eq!(dt.file_name(), "unknown");
    }

    #[test]
    fn paper_result_format_authors_short() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec!["Smith".into(), "Jones".into()],
            year: Some(2023),
            arxiv_id: None,
            doi: None,
            pdf_url: None,
        };
        assert_eq!(p.format_authors(), "Smith, Jones");
    }

    #[test]
    fn paper_result_format_authors_long() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec!["A".into(), "B".into(), "C".into(), "D".into()],
            year: None,
            arxiv_id: None,
            doi: None,
            pdf_url: None,
        };
        assert_eq!(p.format_authors(), "A, et al.");
    }

    #[test]
    fn paper_result_format_year() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec![],
            year: Some(2023),
            arxiv_id: None,
            doi: None,
            pdf_url: None,
        };
        assert_eq!(p.format_year(), " (2023)");
    }

    #[test]
    fn paper_result_format_year_none() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec![],
            year: None,
            arxiv_id: None,
            doi: None,
            pdf_url: None,
        };
        assert!(p.format_year().is_empty());
    }

    #[test]
    fn paper_result_best_pdf_url_direct() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec![],
            year: None,
            arxiv_id: None,
            doi: None,
            pdf_url: Some("https://example.com/paper.pdf".into()),
        };
        assert_eq!(p.best_pdf_url(), "https://example.com/paper.pdf");
    }

    #[test]
    fn paper_result_best_pdf_url_arxiv_fallback() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec![],
            year: None,
            arxiv_id: Some("2301.12345".into()),
            doi: None,
            pdf_url: None,
        };
        assert_eq!(p.best_pdf_url(), "https://arxiv.org/pdf/2301.12345.pdf");
    }

    #[test]
    fn paper_result_best_pdf_url_empty() {
        let p = PaperResult {
            title: "Test".into(),
            authors: vec![],
            year: None,
            arxiv_id: None,
            doi: None,
            pdf_url: None,
        };
        assert!(p.best_pdf_url().is_empty());
    }

    #[test]
    fn args_default_model_ctx_tokens() {
        let args = Args::parse_from(["test", "--folder", "/tmp"]);
        assert_eq!(args.model_ctx_tokens, 8192);
    }

    #[test]
    fn args_default_model() {
        let args = Args::parse_from(["test", "--folder", "/tmp"]);
        assert_eq!(args.model, "gemma2:latest");
    }

    #[test]
    fn generation_params_default_all_none() {
        let p = GenerationParams::default();
        assert!(p.temperature.is_none());
        assert!(p.top_p.is_none());
        assert!(p.max_tokens.is_none());
        assert!(p.seed.is_none());
    }

    #[test]
    fn generation_params_with_values() {
        let p = GenerationParams {
            temperature: Some(0.1),
            top_p: Some(0.9),
            max_tokens: Some(2048),
            seed: Some(42),
        };
        assert_eq!(p.temperature, Some(0.1));
        assert_eq!(p.top_p, Some(0.9));
        assert_eq!(p.max_tokens, Some(2048));
        assert_eq!(p.seed, Some(42));
    }

    #[test]
    fn params_is_empty_when_all_none() {
        assert!(GenerationParams::default().is_empty());
    }

    #[test]
    fn params_is_not_empty_when_temperature_set() {
        let p = GenerationParams { temperature: Some(0.5), ..Default::default() };
        assert!(!p.is_empty());
    }

    #[test]
    fn params_is_not_empty_when_seed_set() {
        let p = GenerationParams { seed: Some(1), ..Default::default() };
        assert!(!p.is_empty());
    }

    #[test]
    fn additional_json_none_when_top_p_and_seed_absent() {
        let p = GenerationParams { temperature: Some(0.5), max_tokens: Some(100), ..Default::default() };
        assert!(p.additional_json().is_none());
    }

    #[test]
    fn additional_json_none_for_default_params() {
        assert!(GenerationParams::default().additional_json().is_none());
    }

    #[test]
    fn additional_json_includes_top_p_when_set() {
        let p = GenerationParams { top_p: Some(0.9), ..Default::default() };
        let json = p.additional_json().unwrap();
        assert_eq!(json["top_p"].as_f64(), Some(0.9));
        assert!(json.get("seed").is_none());
    }

    #[test]
    fn additional_json_includes_seed_when_set() {
        let p = GenerationParams { seed: Some(42), ..Default::default() };
        let json = p.additional_json().unwrap();
        assert_eq!(json["seed"].as_u64(), Some(42));
        assert!(json.get("top_p").is_none());
    }

    #[test]
    fn additional_json_includes_both_when_set() {
        let p = GenerationParams { top_p: Some(0.8), seed: Some(123), ..Default::default() };
        let json = p.additional_json().unwrap();
        assert_eq!(json["top_p"].as_f64(), Some(0.8));
        assert_eq!(json["seed"].as_u64(), Some(123));
    }

    #[test]
    fn additional_json_f64_precision_preserved() {
        // Regression: f64 → Number::from_f64 round-trip should preserve reasonable values.
        let p = GenerationParams { top_p: Some(0.95), ..Default::default() };
        let json = p.additional_json().unwrap();
        let returned = json["top_p"].as_f64().unwrap();
        assert!((returned - 0.95).abs() < 1e-10);
    }
}