srz 0.7.3

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ServerCapabilities {
    pub endpoint: String,
    pub model_id: String,
    pub n_ctx: usize,
    pub total_slots: usize,
    pub supports_tokenize: bool,
    pub tokenize_endpoint: Option<String>,
    pub backend_kind: String,
}

impl Default for ServerCapabilities {
    fn default() -> Self {
        Self {
            endpoint: String::new(),
            model_id: String::new(),
            n_ctx: 16_384,
            total_slots: 1,
            supports_tokenize: false,
            tokenize_endpoint: None,
            backend_kind: "default".to_string(),
        }
    }
}

/// Minimum context window size in tokens.
pub const MIN_CONTEXT_WINDOW: usize = 2_048;

/// Maximum context window size in tokens (2M tokens ceiling to prevent arithmetic overflow and OOM).
pub const MAX_CONTEXT_WINDOW: usize = 2_097_152;

/// Maximum concurrent map-reduce worker count (caps untrusted server `total_slots` or CLI `--concurrency`).
pub const MAX_CONCURRENCY: usize = 64;

/// Sizing parameters dynamically derived from the active LLM server context window and slot capacity.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct DynamicBudget {
    /// Total context window in tokens (n_ctx)
    pub n_ctx: usize,
    /// Maximum tokens the model is instructed to generate in completion
    pub max_completion: u32,
    /// Parallel map-reduce worker count (autoscaled to total_slots unless CLI overrides)
    pub concurrency: usize,
    /// Maximum input tokens allowable in a direct single-pass summarization
    pub single_pass_threshold: usize,
    /// Maximum tokens per chunk when map-reducing long documents
    pub chunk_tokens: usize,
    /// Token overlap between consecutive chunks
    pub chunk_overlap: usize,
    /// Grounding document context budget for grounded Q&A (srz ask)
    pub ask_doc_budget: usize,
    /// Sliding conversational history budget for multi-turn Q&A (srz ask)
    pub ask_hist_budget: usize,
}

impl DynamicBudget {
    /// Derives dynamic token sizing and concurrency limits from the detected server context and slots.
    pub fn new(
        n_ctx: usize,
        total_slots: usize,
        cli_max_tokens: Option<u32>,
        cli_concurrency: Option<usize>,
    ) -> Self {
        let n_ctx = n_ctx.clamp(MIN_CONTEXT_WINDOW, MAX_CONTEXT_WINDOW);
        let max_completion_cap = (n_ctx / 3).clamp(512, u32::MAX as usize) as u32;
        let max_completion = match cli_max_tokens {
            Some(m) if m > 0 => m.min(max_completion_cap),
            _ => (n_ctx / 4).clamp(1024, 8192) as u32,
        };
        let concurrency = cli_concurrency
            .unwrap_or(total_slots)
            .clamp(1, MAX_CONCURRENCY);
        let system_margin = 500usize;
        let cushion = (n_ctx / 20).clamp(200, 1600); // 5% headroom
        let overhead = (max_completion as usize)
            .saturating_add(system_margin)
            .saturating_add(cushion);
        let usable = n_ctx.saturating_sub(overhead);
        let single_pass_threshold = usable;
        let chunk_tokens = (usable.saturating_mul(7) / 10).max(1500);
        let chunk_overlap = (chunk_tokens / 15).clamp(100, 1000);
        let ask_doc_budget = usable / 2;
        let ask_hist_budget = usable.saturating_mul(4) / 10;

        Self {
            n_ctx,
            max_completion,
            concurrency,
            single_pass_threshold,
            chunk_tokens,
            chunk_overlap,
            ask_doc_budget,
            ask_hist_budget,
        }
    }

    /// Default baseline budget assuming a standard 32k context window and 1 slot.
    pub fn default_32k() -> Self {
        Self::new(32_768, 1, None, None)
    }
}

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

    #[test]
    fn test_dynamic_budget_small_context_4k() {
        let budget = DynamicBudget::new(4096, 1, None, None);
        assert_eq!(budget.n_ctx, 4096);
        assert_eq!(budget.max_completion, 1024);
        assert_eq!(budget.concurrency, 1);
        assert!(budget.single_pass_threshold <= 2500);
        assert!(budget.chunk_tokens <= 2000);
        assert!(budget.chunk_tokens + budget.max_completion as usize <= 4096);
    }

    #[test]
    fn test_dynamic_budget_standard_32k_with_multi_slots() {
        let budget = DynamicBudget::new(32_768, 4, None, None);
        assert_eq!(budget.n_ctx, 32_768);
        assert_eq!(budget.max_completion, 8192);
        assert_eq!(budget.concurrency, 4); // Autoscaled to 4 slots!
        assert!(budget.single_pass_threshold > 20_000);
        assert!(budget.chunk_tokens >= 14_000);
    }

    #[test]
    fn test_dynamic_budget_large_128k() {
        let budget = DynamicBudget::new(131_072, 8, None, None);
        assert_eq!(budget.n_ctx, 131_072);
        assert_eq!(budget.max_completion, 8192);
        assert_eq!(budget.concurrency, 8);
        assert!(budget.single_pass_threshold > 100_000);
    }

    #[test]
    fn test_dynamic_budget_cli_overrides() {
        let budget = DynamicBudget::new(32_768, 4, Some(2048), Some(2));
        assert_eq!(budget.max_completion, 2048);
        assert_eq!(budget.concurrency, 2);
    }

    #[test]
    fn test_dynamic_budget_clamping_and_overflow_safety() {
        // Extreme values from rogue server or misconfigured CLI: usize::MAX should clamp without panic
        let budget = DynamicBudget::new(usize::MAX, 100_000, Some(u32::MAX), Some(1_000_000));
        assert_eq!(budget.n_ctx, MAX_CONTEXT_WINDOW);
        assert_eq!(budget.concurrency, MAX_CONCURRENCY);
        assert!(budget.max_completion <= (MAX_CONTEXT_WINDOW / 3) as u32);
        assert!(budget.single_pass_threshold < MAX_CONTEXT_WINDOW);

        // Lower bound clamping
        let small_budget = DynamicBudget::new(10, 0, None, None);
        assert_eq!(small_budget.n_ctx, MIN_CONTEXT_WINDOW);
        assert_eq!(small_budget.concurrency, 1);
    }
}