use crate::mcp::llmclient::{LLMClient, LLMError};
use async_trait::async_trait;
use ollama_rs::Ollama;
use ollama_rs::error::OllamaError;
use ollama_rs::generation::chat::request::ChatMessageRequest;
use ollama_rs::generation::chat::{ChatMessage, ChatMessageResponse};
use ollama_rs::generation::completion::request::GenerationRequest;

impl From<OllamaError> for LLMError {
    fn from(err: OllamaError) -> LLMError {
        match err {
            OllamaError::ToolCallError(e) => LLMError::Other(e.to_string()),
            OllamaError::JsonError(e) => LLMError::Other(e.to_string()),
            OllamaError::ReqwestError(e) => LLMError::NetworkError(e.to_string()),
            OllamaError::InternalError(e) => LLMError::Other(e.message),
            OllamaError::Other(e) => LLMError::Other(e.to_string()),
        }
    }
}

#[async_trait]
impl LLMClient for Ollama {
    async fn send_chat_messages_with_history(
        &self,
        history: &mut Vec<ChatMessage>,
        request: ChatMessageRequest,
    ) -> Result<ChatMessageResponse, LLMError> {
        Ok(self
            .send_chat_messages_with_history(history, request)
            .await?)
    }

    async fn preload_model(&self, model: &str) {
        let _ = self
            .generate(GenerationRequest::new(model.to_string(), ""))
            .await;
    }
}