chatdelta 0.8.2

//! Prompt Optimization Engine
//! 
//! Advanced prompt engineering and optimization for better AI responses

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Prompt optimizer for enhancing query effectiveness
pub struct PromptOptimizer {
    /// Optimization strategies
    strategies: Vec<Box<dyn OptimizationStrategy>>,
    /// Template library
    templates: TemplateLibrary,
    /// Performance history
    history: PerformanceHistory,
}

/// Optimization strategies
trait OptimizationStrategy: Send + Sync {
    fn optimize(&self, prompt: &str, context: &OptimizationContext) -> String;
    fn name(&self) -> &str;
}

/// Context for optimization decisions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationContext {
    pub task_type: TaskCategory,
    pub target_model: Option<String>,
    pub desired_length: Option<usize>,
    pub tone: Option<Tone>,
    pub expertise_level: ExpertiseLevel,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TaskCategory {
    Analysis,
    Generation,
    Summarization,
    Translation,
    QuestionAnswering,
    Reasoning,
    Creative,
    Technical,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Tone {
    Professional,
    Casual,
    Academic,
    Creative,
    Technical,
    Friendly,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ExpertiseLevel {
    Beginner,
    Intermediate,
    Advanced,
    Expert,
}

impl PromptOptimizer {
    pub fn new() -> Self {
        Self {
            strategies: Self::default_strategies(),
            templates: TemplateLibrary::default(),
            history: PerformanceHistory::new(),
        }
    }
    
    /// Optimize a prompt for better AI response
    pub fn optimize(&self, prompt: &str) -> OptimizedPrompt {
        let context = self.analyze_context(prompt);
        
        // Apply optimization strategies
        let mut optimized = prompt.to_string();
        let mut techniques_applied = Vec::new();
        
        for strategy in &self.strategies {
            if self.should_apply_strategy(&strategy, &context) {
                optimized = strategy.optimize(&optimized, &context);
                techniques_applied.push(strategy.name().to_string());
            }
        }
        
        // Apply template if applicable
        if let Some(template) = self.templates.find_best_template(&context) {
            optimized = template.apply(&optimized);
            techniques_applied.push(format!("Template: {}", template.name));
        }
        
        // Generate variations
        let variations = self.generate_variations(&optimized, &context);
        
        let confidence = self.calculate_confidence(&techniques_applied);
        
        OptimizedPrompt {
            original: prompt.to_string(),
            optimized,
            variations,
            techniques_applied,
            context,
            confidence,
        }
    }
    
    fn analyze_context(&self, prompt: &str) -> OptimizationContext {
        let task_type = self.detect_task_type(prompt);
        let expertise = self.detect_expertise_level(prompt);
        
        OptimizationContext {
            task_type,
            target_model: None,
            desired_length: self.estimate_desired_length(prompt),
            tone: self.detect_tone(prompt),
            expertise_level: expertise,
        }
    }
    
    fn detect_task_type(&self, prompt: &str) -> TaskCategory {
        let lower = prompt.to_lowercase();
        
        if lower.contains("analyze") || lower.contains("explain") {
            TaskCategory::Analysis
        } else if lower.contains("create") || lower.contains("generate") || lower.contains("write") {
            TaskCategory::Generation
        } else if lower.contains("summarize") || lower.contains("tldr") {
            TaskCategory::Summarization
        } else if lower.contains("translate") {
            TaskCategory::Translation
        } else if lower.starts_with("what") || lower.starts_with("how") || lower.starts_with("why") {
            TaskCategory::QuestionAnswering
        } else if lower.contains("reason") || lower.contains("think") {
            TaskCategory::Reasoning
        } else {
            TaskCategory::Technical
        }
    }
    
    fn detect_expertise_level(&self, prompt: &str) -> ExpertiseLevel {
        let technical_terms = ["algorithm", "implementation", "architecture", "optimization", "complexity"];
        let count = technical_terms.iter().filter(|t| prompt.contains(*t)).count();
        
        match count {
            0 => ExpertiseLevel::Beginner,
            1 => ExpertiseLevel::Intermediate,
            2 => ExpertiseLevel::Advanced,
            _ => ExpertiseLevel::Expert,
        }
    }
    
    fn detect_tone(&self, prompt: &str) -> Option<Tone> {
        if prompt.contains("please") || prompt.contains("could you") {
            Some(Tone::Friendly)
        } else if prompt.contains("technical") || prompt.contains("detailed") {
            Some(Tone::Technical)
        } else if prompt.contains("academic") || prompt.contains("research") {
            Some(Tone::Academic)
        } else {
            None
        }
    }
    
    fn estimate_desired_length(&self, prompt: &str) -> Option<usize> {
        if prompt.contains("brief") || prompt.contains("short") {
            Some(200)
        } else if prompt.contains("detailed") || prompt.contains("comprehensive") {
            Some(1000)
        } else if prompt.contains("concise") {
            Some(300)
        } else {
            None
        }
    }
    
    fn should_apply_strategy(&self, _strategy: &Box<dyn OptimizationStrategy>, _context: &OptimizationContext) -> bool {
        // Decide whether to apply a strategy based on context
        true // Simplified for now
    }
    
    fn generate_variations(&self, optimized: &str, _context: &OptimizationContext) -> Vec<PromptVariation> {
        let mut variations = Vec::new();
        
        // Variation 1: More specific
        variations.push(PromptVariation {
            prompt: format!("Specifically regarding the following: {}", optimized),
            strategy: "Specificity Enhancement".to_string(),
            expected_improvement: 15.0,
        });
        
        // Variation 2: With examples
        variations.push(PromptVariation {
            prompt: format!("{}\n\nProvide concrete examples in your response.", optimized),
            strategy: "Example Request".to_string(),
            expected_improvement: 20.0,
        });
        
        // Variation 3: Structured output
        variations.push(PromptVariation {
            prompt: format!("{}\n\nStructure your response with clear sections and bullet points where appropriate.", optimized),
            strategy: "Structure Enhancement".to_string(),
            expected_improvement: 25.0,
        });
        
        variations
    }
    
    fn calculate_confidence(&self, techniques: &[String]) -> f64 {
        // Base confidence
        let mut confidence = 0.7;
        
        // Add confidence for each technique applied
        confidence += techniques.len() as f64 * 0.05;
        
        confidence.min(0.95)
    }
    
    fn default_strategies() -> Vec<Box<dyn OptimizationStrategy>> {
        vec![
            Box::new(ClarityEnhancer),
            Box::new(ContextInjector),
            Box::new(ChainOfThought),
            Box::new(FewShotLearning),
            Box::new(RoleSpecification),
        ]
    }
}

/// Optimized prompt result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizedPrompt {
    pub original: String,
    pub optimized: String,
    pub variations: Vec<PromptVariation>,
    pub techniques_applied: Vec<String>,
    pub context: OptimizationContext,
    pub confidence: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptVariation {
    pub prompt: String,
    pub strategy: String,
    pub expected_improvement: f64,
}

// Optimization Strategies

struct ClarityEnhancer;
impl OptimizationStrategy for ClarityEnhancer {
    fn optimize(&self, prompt: &str, _context: &OptimizationContext) -> String {
        // Add clarity markers
        if !prompt.contains('?') && !prompt.contains('.') {
            format!("{}. Please provide a clear and detailed response.", prompt)
        } else {
            prompt.to_string()
        }
    }
    
    fn name(&self) -> &str {
        "Clarity Enhancement"
    }
}

struct ContextInjector;
impl OptimizationStrategy for ContextInjector {
    fn optimize(&self, prompt: &str, context: &OptimizationContext) -> String {
        match context.task_type {
            TaskCategory::Analysis => {
                format!("Analyze the following in detail: {}", prompt)
            }
            TaskCategory::Generation => {
                format!("Generate a comprehensive response for: {}", prompt)
            }
            _ => prompt.to_string(),
        }
    }
    
    fn name(&self) -> &str {
        "Context Injection"
    }
}

struct ChainOfThought;
impl OptimizationStrategy for ChainOfThought {
    fn optimize(&self, prompt: &str, context: &OptimizationContext) -> String {
        match context.task_type {
            TaskCategory::Reasoning | TaskCategory::Analysis => {
                format!("Let's think step by step about this: {}", prompt)
            }
            _ => prompt.to_string(),
        }
    }
    
    fn name(&self) -> &str {
        "Chain of Thought"
    }
}

struct FewShotLearning;
impl OptimizationStrategy for FewShotLearning {
    fn optimize(&self, prompt: &str, context: &OptimizationContext) -> String {
        if matches!(context.expertise_level, ExpertiseLevel::Expert) {
            prompt.to_string()
        } else {
            format!("{}\n\nConsider similar examples if helpful.", prompt)
        }
    }
    
    fn name(&self) -> &str {
        "Few-Shot Learning"
    }
}

struct RoleSpecification;
impl OptimizationStrategy for RoleSpecification {
    fn optimize(&self, prompt: &str, context: &OptimizationContext) -> String {
        let role = match context.task_type {
            TaskCategory::Technical => "You are a technical expert. ",
            TaskCategory::Creative => "You are a creative professional. ",
            TaskCategory::Generation => "You are a content creator. ",
            _ => "",
        };
        
        if !role.is_empty() {
            format!("{}{}", role, prompt)
        } else {
            prompt.to_string()
        }
    }
    
    fn name(&self) -> &str {
        "Role Specification"
    }
}

/// Template library for common patterns
#[derive(Default)]
struct TemplateLibrary {
    templates: HashMap<String, PromptTemplate>,
}

struct PromptTemplate {
    name: String,
    pattern: String,
}

impl PromptTemplate {
    fn apply(&self, prompt: &str) -> String {
        self.pattern.replace("{PROMPT}", prompt)
    }
}

impl TemplateLibrary {
    fn find_best_template(&self, _context: &OptimizationContext) -> Option<&PromptTemplate> {
        None // Simplified for now
    }
}

/// Performance history for learning
struct PerformanceHistory {
    history: Vec<HistoryEntry>,
}

struct HistoryEntry {
    prompt: String,
    optimized: String,
    performance_score: f64,
}

impl PerformanceHistory {
    fn new() -> Self {
        Self {
            history: Vec::new(),
        }
    }
    
    fn record(&mut self, prompt: String, optimized: String, score: f64) {
        self.history.push(HistoryEntry {
            prompt,
            optimized,
            performance_score: score,
        });
    }
}