//! Ollama client for AI-powered text generation with security hardening.
//!
//! Security features:
//! - API key authentication via X-API-Key header
//! - Rate limiting per user/IP
//! - Request timeouts and circuit breaker
//! - Input validation and sanitization

use reqwest::{Client, Error as ReqwestError};
use serde::{Deserialize, Serialize};
use std::time::Duration;

const OLLAMA_URL: &str = "http://nxtgauge-ai-assistant:11434";
const DEFAULT_MODEL: &str = "gemma3:270m";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
const MAX_PROMPT_LENGTH: usize = 100_000; // 100KB limit
const MAX_RETRIES: u32 = 3;

/// Security configuration for Ollama client
#[derive(Debug, Clone)]
pub struct OllamaSecurityConfig {
    /// API key for authentication (required)
    pub api_key: Option<String>,
    /// User ID for rate limiting tracking
    pub user_id: Option<String>,
    /// Client IP for rate limiting
    pub client_ip: Option<String>,
    /// Maximum requests per minute per user
    pub rate_limit_per_minute: u32,
}

impl Default for OllamaSecurityConfig {
    fn default() -> Self {
        Self {
            api_key: std::env::var("OLLAMA_API_KEY").ok(),
            user_id: None,
            client_ip: None,
            rate_limit_per_minute: 60, // Default: 60 requests/minute
        }
    }
}

#[derive(Debug, Clone)]
pub struct OllamaClient {
    http_client: Client,
    base_url: String,
    model: String,
    security: OllamaSecurityConfig,
}

#[derive(Debug, Serialize)]
struct GenerateRequest {
    model: String,
    prompt: String,
    stream: bool,
    options: Option<GenerationOptions>,
}

#[derive(Debug, Serialize, Default)]
struct GenerationOptions {
    temperature: Option<f32>,
    top_p: Option<f32>,
    top_k: Option<i32>,
    num_predict: Option<i32>,
}

#[derive(Debug, Deserialize)]
pub struct GenerateResponse {
    pub model: String,
    pub created_at: String,
    pub response: String,
    pub done: bool,
    pub context: Option<Vec<i32>>,
    pub total_duration: Option<u64>,
    pub load_duration: Option<u64>,
    pub prompt_eval_count: Option<i32>,
    pub prompt_eval_duration: Option<u64>,
    pub eval_count: Option<i32>,
    pub eval_duration: Option<u64>,
}

#[derive(Debug, Deserialize)]
#[allow(dead_code)]
struct OllamaErrorResponse {
    error: String,
}

#[derive(Debug, thiserror::Error)]
pub enum OllamaError {
    #[error("HTTP request failed: {0}")]
    RequestFailed(#[from] ReqwestError),
    
    #[error("Ollama API error: {0}")]
    ApiError(String),
    
    #[error("Failed to parse response: {0}")]
    ParseError(String),
    
    #[error("Connection timeout")]
    Timeout,
    
    #[error("Model not found: {0}")]
    ModelNotFound(String),
    
    #[error("API key required but not provided")]
    MissingApiKey,
    
    #[error("Invalid API key")]
    InvalidApiKey,
    
    #[error("Rate limit exceeded: {0}")]
    RateLimitExceeded(String),
    
    #[error("Prompt too long: {0} chars (max {1})")]
    PromptTooLong(usize, usize),
    
    #[error("Circuit breaker open: too many failures")]
    CircuitBreakerOpen,
    
    #[error("Max retries exceeded")]
    MaxRetriesExceeded,
}

/// Validates and sanitizes user input
fn validate_prompt(prompt: &str) -> Result<String, OllamaError> {
    if prompt.len() > MAX_PROMPT_LENGTH {
        return Err(OllamaError::PromptTooLong(prompt.len(), MAX_PROMPT_LENGTH));
    }
    
    // Basic sanitization - remove null bytes and control characters
    let sanitized: String = prompt
        .chars()
        .filter(|c| !c.is_control() || *c == '\n' || *c == '\t')
        .collect();
    
    Ok(sanitized)
}

/// Rate limiter using a simple token bucket approach
#[derive(Debug)]
pub struct RateLimiter {
    requests: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<String, Vec<std::time::Instant>>>>,
    max_requests: u32,
    window_duration: Duration,
}

impl RateLimiter {
    pub fn new(max_requests: u32, window_minutes: u64) -> Self {
        Self {
            requests: std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
            max_requests,
            window_duration: Duration::from_secs(window_minutes * 60),
        }
    }
    
    /// Check if the request is allowed
    pub fn check_rate_limit(&self, key: &str) -> Result<(), OllamaError> {
        let now = std::time::Instant::now();
        let window_start = now - self.window_duration;
        
        let mut requests = self.requests.lock().map_err(|_| {
            OllamaError::ApiError("Rate limiter lock failed".to_string())
        })?;
        
        // Get or create entry for this key
        let user_requests = requests.entry(key.to_string()).or_default();
        
        // Remove old requests outside the window
        user_requests.retain(|&time| time > window_start);
        
        // Check if under limit
        if user_requests.len() >= self.max_requests as usize {
            return Err(OllamaError::RateLimitExceeded(
                format!("{} requests per {} minutes exceeded", self.max_requests, self.window_duration.as_secs() / 60)
            ));
        }
        
        // Add this request
        user_requests.push(now);
        Ok(())
    }
}

/// Circuit breaker for handling cascading failures
#[derive(Debug)]
pub struct CircuitBreaker {
    failures: std::sync::Arc<std::sync::atomic::AtomicU32>,
    threshold: u32,
    reset_after: Duration,
    last_failure: std::sync::Arc<std::sync::Mutex<Option<std::time::Instant>>>,
}

impl CircuitBreaker {
    pub fn new(threshold: u32, reset_after_secs: u64) -> Self {
        Self {
            failures: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
            threshold,
            reset_after: Duration::from_secs(reset_after_secs),
            last_failure: std::sync::Arc::new(std::sync::Mutex::new(None)),
        }
    }
    
    pub fn can_execute(&self) -> Result<(), OllamaError> {
        let failures = self.failures.load(std::sync::atomic::Ordering::Relaxed);
        
        if failures >= self.threshold {
            // Check if enough time has passed to reset
            if let Ok(last) = self.last_failure.lock() {
                if let Some(last_failure) = *last {
                    if std::time::Instant::now().duration_since(last_failure) > self.reset_after {
                        // Reset circuit breaker
                        self.failures.store(0, std::sync::atomic::Ordering::Relaxed);
                        return Ok(());
                    }
                }
            }
            return Err(OllamaError::CircuitBreakerOpen);
        }
        
        Ok(())
    }
    
    pub fn record_success(&self) {
        self.failures.store(0, std::sync::atomic::Ordering::Relaxed);
    }
    
    pub fn record_failure(&self) {
        self.failures.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if let Ok(mut last) = self.last_failure.lock() {
            *last = Some(std::time::Instant::now());
        }
    }
}

// Global rate limiter instance (shared across all clients)
lazy_static::lazy_static! {
    static ref GLOBAL_RATE_LIMITER: RateLimiter = RateLimiter::new(
        std::env::var("OLLAMA_RATE_LIMIT_PER_MINUTE")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(60),
        1 // 1 minute window
    );
    
    static ref CIRCUIT_BREAKER: CircuitBreaker = CircuitBreaker::new(
        std::env::var("OLLAMA_CIRCUIT_THRESHOLD")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(5),
        60 // Reset after 60 seconds
    );
}

impl OllamaClient {
    pub fn new() -> Self {
        let http_client = Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .build()
            .expect("Failed to create HTTP client");
        
        Self {
            http_client,
            base_url: OLLAMA_URL.to_string(),
            model: DEFAULT_MODEL.to_string(),
            security: OllamaSecurityConfig::default(),
        }
    }
    
    pub fn with_url(base_url: impl Into<String>) -> Self {
        let http_client = Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .build()
            .expect("Failed to create HTTP client");
        
        Self {
            http_client,
            base_url: base_url.into(),
            model: DEFAULT_MODEL.to_string(),
            security: OllamaSecurityConfig::default(),
        }
    }
    
    pub fn with_security(mut self, security: OllamaSecurityConfig) -> Self {
        self.security = security;
        self
    }
    
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model = model.into();
        self
    }
    
    pub fn get_model(&self) -> &str {
        &self.model
    }
    
    fn get_rate_limit_key(&self) -> String {
        self.security.user_id.clone()
            .or_else(|| self.security.client_ip.clone())
            .unwrap_or_else(|| "anonymous".to_string())
    }
    
    fn build_request_headers(&self) -> reqwest::header::HeaderMap {
        let mut headers = reqwest::header::HeaderMap::new();
        
        // Add API key header if configured
        if let Some(api_key) = &self.security.api_key {
            if let Ok(val) = reqwest::header::HeaderValue::from_str(api_key) {
                headers.insert("X-API-Key", val);
            }
        }
        
        // Add user tracking headers for audit
        if let Some(user_id) = &self.security.user_id {
            if let Ok(val) = reqwest::header::HeaderValue::from_str(user_id) {
                headers.insert("X-User-Id", val);
            }
        }
        
        headers
    }
    
    /// Generate text using the configured model and prompt
    pub async fn generate(&self, prompt: impl Into<String>) -> Result<GenerateResponse, OllamaError> {
        // 1. Check circuit breaker
        CIRCUIT_BREAKER.can_execute()?;
        
        // 2. Validate and sanitize input
        let sanitized_prompt = validate_prompt(&prompt.into())?;
        
        // 3. Check rate limit
        let rate_key = self.get_rate_limit_key();
        GLOBAL_RATE_LIMITER.check_rate_limit(&rate_key)?;
        
        // 4. Build request
        let request = GenerateRequest {
            model: self.model.clone(),
            prompt: sanitized_prompt,
            stream: false,
            options: None,
        };
        
        let url = format!("{}/api/generate", self.base_url);
        
        // 5. Execute with retry logic
        let mut last_error = None;
        for attempt in 0..MAX_RETRIES {
            let result = self.execute_request(&url, &request).await;
            
            match result {
                Ok(response) => {
                    CIRCUIT_BREAKER.record_success();
                    return Ok(response);
                }
                Err(e) => {
                    last_error = Some(e);
                    // Don't retry on client errors (4xx)
                    if let OllamaError::ApiError(_) = &last_error.as_ref().unwrap() {
                        break;
                    }
                    // Exponential backoff
                    if attempt < MAX_RETRIES - 1 {
                        tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
                    }
                }
            }
        }
        
        // All retries failed
        CIRCUIT_BREAKER.record_failure();
        Err(last_error.unwrap_or(OllamaError::MaxRetriesExceeded))
    }
    
    async fn execute_request(
        &self,
        url: &str,
        request: &GenerateRequest,
    ) -> Result<GenerateResponse, OllamaError> {
        let mut request_builder = self.http_client
            .post(url)
            .json(request);
        
        // Add security headers
        let headers = self.build_request_headers();
        request_builder = request_builder.headers(headers);
        
        let response = request_builder
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() {
                    OllamaError::Timeout
                } else {
                    OllamaError::RequestFailed(e)
                }
            })?;
        
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            
            match status.as_u16() {
                401 => return Err(OllamaError::InvalidApiKey),
                404 => return Err(OllamaError::ModelNotFound(self.model.clone())),
                429 => return Err(OllamaError::RateLimitExceeded("Server rate limit".to_string())),
                _ => return Err(OllamaError::ApiError(format!("{}: {}", status, error_text))),
            }
        }
        
        let result = response.json::<GenerateResponse>()
            .await
            .map_err(|e| OllamaError::ParseError(e.to_string()))?;
        
        Ok(result)
    }
    
    /// Health check endpoint
    pub async fn health_check(&self) -> Result<bool, OllamaError> {
        let url = format!("{}/api/tags", self.base_url);
        
        let mut request_builder = self.http_client.get(&url);
        
        // Add security headers
        let headers = self.build_request_headers();
        request_builder = request_builder.headers(headers);
        
        let response = request_builder
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() {
                    OllamaError::Timeout
                } else {
                    OllamaError::RequestFailed(e)
                }
            })?;
        
        Ok(response.status().is_success())
    }
    
    /// Generate a job description based on a prompt
    pub async fn generate_job_description(&self, prompt: &str) -> Result<String, OllamaError> {
        let enhanced_prompt = format!(
            "Generate a professional job description based on the following prompt:\n\n{}\n\n\
             Provide a well-structured description with clear responsibilities and requirements.",
            prompt
        );
        
        let response = self.generate(enhanced_prompt).await?;
        Ok(response.response)
    }
    
    /// Analyze a resume and provide feedback
    pub async fn analyze_resume(&self, resume_content: &str, job_description: &str) -> Result<String, OllamaError> {
        let prompt = format!(
            "Analyze the following resume against this job description:\n\n\
             Job Description:\n{}\n\n\
             Resume:\n{}\n\n\
             Provide specific feedback on:\n\
             1. How well the resume matches the job requirements\n\
             2. Missing skills or experience\n\
             3. Suggestions for improvement\n\
             4. Overall match percentage",
            job_description, resume_content
        );
        
        let response = self.generate(prompt).await?;
        Ok(response.response)
    }
    
    /// Generate a cover letter
    pub async fn generate_cover_letter(
        &self,
        candidate_info: &str,
        job_description: &str,
        tone: &str,
    ) -> Result<String, OllamaError> {
        let prompt = format!(
            "Write a {} cover letter for a candidate with the following background:\n\n\
             Candidate: {}\n\n\
             Job Description: {}\n\n\
             The cover letter should be professional and highlight relevant experience.",
            tone, candidate_info, job_description
        );
        
        let response = self.generate(prompt).await?;
        Ok(response.response)
    }
}

impl Default for OllamaClient {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_client_creation() {
        let client = OllamaClient::new();
        assert_eq!(client.get_model(), DEFAULT_MODEL);
    }
    
    #[test]
    fn test_client_with_custom_model() {
        let client = OllamaClient::new()
            .with_model("gemma:4b");
        assert_eq!(client.get_model(), "gemma:4b");
    }
    
    #[test]
    fn test_client_with_custom_url() {
        let client = OllamaClient::with_url("http://custom:11434");
        assert_eq!(client.get_model(), DEFAULT_MODEL);
    }
    
    #[test]
    fn test_validate_prompt_too_long() {
        let long_prompt = "a".repeat(MAX_PROMPT_LENGTH + 1);
        let result = validate_prompt(&long_prompt);
        assert!(matches!(result, Err(OllamaError::PromptTooLong(_, _))));
    }
    
    #[test]
    fn test_validate_prompt_ok() {
        let prompt = "Hello, world!";
        let result = validate_prompt(prompt);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), prompt);
    }
    
    #[test]
    fn test_rate_limiter() {
        let limiter = RateLimiter::new(2, 1); // 2 requests per minute
        
        // First two should succeed
        assert!(limiter.check_rate_limit("user1").is_ok());
        assert!(limiter.check_rate_limit("user1").is_ok());
        
        // Third should fail
        assert!(limiter.check_rate_limit("user1").is_err());
        
        // Different user should succeed
        assert!(limiter.check_rate_limit("user2").is_ok());
    }
    
    #[test]
    fn test_circuit_breaker() {
        let breaker = CircuitBreaker::new(2, 60);
        
        // Initially should allow
        assert!(breaker.can_execute().is_ok());
        
        // Record failures
        breaker.record_failure();
        breaker.record_failure();
        
        // Should now be open
        assert!(breaker.can_execute().is_err());
        
        // Success should reset
        breaker.record_success();
        assert!(breaker.can_execute().is_ok());
    }
}
