coderlib 0.1.0

//! Local model provider implementation for CoderLib
//!
//! This module provides support for local models running on:
//! - LM Studio
//! - Ollama
//! - Text Generation WebUI
//! - Any OpenAI-compatible local server
//!
//! The provider automatically discovers available models and provides
//! a seamless interface for local AI model usage.

use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};

use crate::core::error::ProviderError;
use crate::llm::{Provider, ProviderEvent, ProviderResponse, Model, ModelCapabilities};
use crate::storage::Message;
use crate::tools::Tool;
use super::openai_compatible::{OpenAICompatibleProvider, CompatibleProviderConfig};

/// Local model provider that supports various local AI servers
#[derive(Debug, Clone)]
pub struct LocalProvider {
    inner: OpenAICompatibleProvider,
    endpoint: String,
    discovered_models: Vec<LocalModel>,
}

/// Local model information from discovery
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalModel {
    pub id: String,
    pub object: String,
    #[serde(rename = "type")]
    pub model_type: Option<String>,
    pub publisher: Option<String>,
    pub arch: Option<String>,
    pub compatibility_type: Option<String>,
    pub quantization: Option<String>,
    pub state: Option<String>,
    pub max_context_length: Option<u64>,
    pub loaded_context_length: Option<u64>,
}

/// Response from local model discovery endpoints
#[derive(Debug, Deserialize)]
struct LocalModelList {
    data: Vec<LocalModel>,
}

/// Configuration for local model providers
#[derive(Debug, Clone)]
pub struct LocalProviderConfig {
    pub endpoint: String,
    pub api_key: Option<String>,
    pub timeout_seconds: u64,
    pub discovery_paths: Vec<String>,
}

impl Default for LocalProviderConfig {
    fn default() -> Self {
        Self {
            endpoint: env::var("LOCAL_ENDPOINT")
                .unwrap_or_else(|_| "http://localhost:1234".to_string()),
            api_key: env::var("LOCAL_API_KEY").ok(),
            timeout_seconds: 60,
            discovery_paths: vec![
                "v1/models".to_string(),           // Standard OpenAI format
                "api/v0/models".to_string(),       // LM Studio beta
                "api/v1/models".to_string(),       // Ollama
                "v1/internal/model/list".to_string(), // Text Generation WebUI
            ],
        }
    }
}

impl LocalProvider {
    /// Create a new local provider with automatic model discovery
    pub async fn new(config: LocalProviderConfig) -> Result<Self, ProviderError> {
        info!("Initializing local provider with endpoint: {}", config.endpoint);
        
        // Discover available models
        let discovered_models = Self::discover_models(&config).await?;
        
        if discovered_models.is_empty() {
            warn!("No local models discovered at endpoint: {}", config.endpoint);
            return Err(ProviderError::Configuration(
                "No local models found. Ensure your local AI server is running.".to_string()
            ));
        }
        
        info!("Discovered {} local models", discovered_models.len());
        for model in &discovered_models {
            debug!("Found model: {} (state: {:?})", model.id, model.state);
        }
        
        // Use the first available model or the first loaded model
        let selected_model = discovered_models
            .iter()
            .find(|m| m.state.as_deref() == Some("loaded"))
            .or_else(|| discovered_models.first())
            .ok_or_else(|| ProviderError::Configuration("No models available".to_string()))?;
        
        let model = Self::convert_local_model(selected_model);
        
        // Create OpenAI-compatible provider configuration
        let compat_config = CompatibleProviderConfig {
            name: "local".to_string(),
            base_url: format!("{}/v1", config.endpoint.trim_end_matches('/')),
            supports_tools: true,
            supports_streaming: true,
            requires_auth: config.api_key.is_some(),
            custom_headers: vec![],
        };
        
        let api_key = config.api_key.unwrap_or_else(|| "dummy".to_string());
        let inner = OpenAICompatibleProvider::new(compat_config, api_key, model, None, None)?;
        
        Ok(Self {
            inner,
            endpoint: config.endpoint,
            discovered_models,
        })
    }
    
    /// Create a local provider with a specific model
    pub async fn with_model(config: LocalProviderConfig, model_id: &str) -> Result<Self, ProviderError> {
        let discovered_models = Self::discover_models(&config).await?;
        
        let selected_model = discovered_models
            .iter()
            .find(|m| m.id == model_id)
            .ok_or_else(|| ProviderError::Configuration(
                format!("Model '{}' not found in discovered models", model_id)
            ))?;
        
        let model = Self::convert_local_model(selected_model);
        
        let compat_config = CompatibleProviderConfig {
            name: "local".to_string(),
            base_url: format!("{}/v1", config.endpoint.trim_end_matches('/')),
            supports_tools: true,
            supports_streaming: true,
            requires_auth: config.api_key.is_some(),
            custom_headers: vec![],
        };
        
        let api_key = config.api_key.unwrap_or_else(|| "dummy".to_string());
        let inner = OpenAICompatibleProvider::new(compat_config, api_key, model, None, None)?;
        
        Ok(Self {
            inner,
            endpoint: config.endpoint,
            discovered_models,
        })
    }
    
    /// Discover available models from the local endpoint
    async fn discover_models(config: &LocalProviderConfig) -> Result<Vec<LocalModel>, ProviderError> {
        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(config.timeout_seconds))
            .build()
            .map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;
        
        for path in &config.discovery_paths {
            let url = format!("{}/{}", config.endpoint.trim_end_matches('/'), path);
            debug!("Trying model discovery at: {}", url);
            
            match Self::try_discover_at_endpoint(&client, &url, config.api_key.as_deref()).await {
                Ok(models) if !models.is_empty() => {
                    info!("Successfully discovered {} models at {}", models.len(), url);
                    return Ok(models);
                }
                Ok(_) => {
                    debug!("No models found at {}", url);
                }
                Err(e) => {
                    debug!("Failed to discover models at {}: {}", url, e);
                }
            }
        }
        
        Err(ProviderError::Configuration(
            format!("No models discovered at any endpoint. Tried: {:?}", config.discovery_paths)
        ))
    }
    
    /// Try to discover models at a specific endpoint
    async fn try_discover_at_endpoint(
        client: &Client,
        url: &str,
        api_key: Option<&str>,
    ) -> Result<Vec<LocalModel>, ProviderError> {
        let mut request = client.get(url);
        
        if let Some(key) = api_key {
            request = request.header("Authorization", format!("Bearer {}", key));
        }
        
        let response = request
            .send()
            .await
            .map_err(|e| ProviderError::NetworkError(e.to_string()))?;
        
        if !response.status().is_success() {
            return Err(ProviderError::ApiError(format!(
                "Model discovery failed with status: {}",
                response.status()
            )));
        }
        
        let model_list: LocalModelList = response
            .json()
            .await
            .map_err(|e| ProviderError::ResponseParsing(e.to_string()))?;
        
        // Filter models based on the endpoint type
        let filtered_models = if url.contains("api/v0/models") {
            // LM Studio beta endpoint - filter for LLM models
            model_list.data.into_iter()
                .filter(|m| m.object == "model" && m.model_type.as_deref() == Some("llm"))
                .collect()
        } else {
            // Standard endpoints - accept all models
            model_list.data
        };
        
        Ok(filtered_models)
    }
    
    /// Convert a local model to CoderLib Model format
    fn convert_local_model(local_model: &LocalModel) -> Model {
        let context_length = local_model.loaded_context_length
            .or(local_model.max_context_length)
            .unwrap_or(4096) as u32;

        Model {
            id: format!("local.{}", local_model.id),
            name: Self::friendly_model_name(&local_model.id),
            provider: "local".to_string(),
            context_length,
            max_output_tokens: context_length / 2, // Conservative estimate
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false, // Most local models don't support vision yet
            cost_per_input_token: 0.0, // Local models are free
            cost_per_output_token: 0.0,
            capabilities: ModelCapabilities::default(),
        }
    }
    
    /// Generate a friendly name for a model ID
    fn friendly_model_name(model_id: &str) -> String {
        // Remove common prefixes and suffixes
        let mut name = model_id.to_string();
        
        // Remove file path if present
        if let Some(last_slash) = name.rfind('/') {
            name = name[last_slash + 1..].to_string();
        }
        
        // Remove common suffixes
        for suffix in &[".gguf", ".bin", ".safetensors"] {
            if name.ends_with(suffix) {
                name = name[..name.len() - suffix.len()].to_string();
                break;
            }
        }
        
        // Convert underscores and hyphens to spaces
        name = name.replace(['_', '-'], " ");
        
        // Capitalize words
        name.split_whitespace()
            .map(|word| {
                let mut chars = word.chars();
                match chars.next() {
                    None => String::new(),
                    Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
                }
            })
            .collect::<Vec<_>>()
            .join(" ")
    }
    
    /// Get list of discovered models
    pub fn discovered_models(&self) -> &[LocalModel] {
        &self.discovered_models
    }
    
    /// Get the endpoint URL
    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }
}

#[async_trait]
impl Provider for LocalProvider {
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError> {
        self.inner.send_messages(messages, tools).await
    }
    
    async fn stream_response(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError> {
        self.inner.stream_response(messages, tools).await
    }
    
    fn model(&self) -> &Model {
        self.inner.model()
    }
    
    fn name(&self) -> &str {
        "local"
    }
    
    async fn is_available(&self) -> bool {
        // Check if the endpoint is reachable
        let client = Client::new();
        let health_url = format!("{}/v1/models", self.endpoint);
        
        match client.get(&health_url).send().await {
            Ok(response) => response.status().is_success(),
            Err(_) => false,
        }
    }
}