//! Ollama local LLM provider

use super::super::error::{LlmError, Result};
use super::super::messages::{Message, Role};
use super::super::options::GenerationOptions;
use super::super::{LlmProvider, StreamResponse};
use async_trait::async_trait;
use futures::stream::StreamExt;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// Ollama API request
#[derive(Debug, Serialize)]
struct OllamaRequest {
    model: String,
    messages: Vec<OllamaMessage>,
    temperature: f32,
    stream: bool,
}

/// Ollama message format
#[derive(Debug, Serialize, Deserialize)]
struct OllamaMessage {
    role: String,
    content: String,
}

/// Ollama API response
#[derive(Debug, Deserialize)]
struct OllamaResponse {
    message: OllamaMessage,
}

/// Ollama streaming response
#[derive(Debug, Deserialize)]
struct OllamaStreamResponse {
    message: OllamaMessage,
}

/// Ollama provider
pub struct OllamaProvider {
    model: String,
    client: Arc<Client>,
    base_url: String,
}

impl OllamaProvider {
    /// Create new Ollama provider
    ///
    /// # Arguments
    /// * `model` - Model name (e.g., "llama2", "mistral", "neural-chat")
    /// * `base_url` - Ollama server URL (default: http://localhost:11434)
    ///
    /// # Errors
    /// Returns error if model is empty
    pub fn new(model: impl Into<String>) -> Result<Self> {
        Self::with_url(model, "http://localhost:11434")
    }

    /// Create with custom Ollama server URL
    pub fn with_url(model: impl Into<String>, base_url: impl Into<String>) -> Result<Self> {
        let model = model.into();
        if model.is_empty() {
            return Err(LlmError::ConfigError("Ollama model name is empty".into()));
        }

        Ok(OllamaProvider {
            model,
            client: Arc::new(Client::new()),
            base_url: base_url.into(),
        })
    }

    /// Create from environment variable `OLLAMA_MODEL`
    pub fn from_env() -> Result<Self> {
        let model = std::env::var("OLLAMA_MODEL")
            .map_err(|_| LlmError::ConfigError("OLLAMA_MODEL not set".into()))?;
        Self::new(model)
    }

    fn build_request(&self, messages: &[Message], options: &GenerationOptions) -> OllamaRequest {
        let ollama_messages = messages
            .iter()
            .map(|msg| OllamaMessage {
                role: match msg.role {
                    Role::System => "system",
                    Role::User => "user",
                    Role::Assistant => "assistant",
                }
                .to_string(),
                content: msg.content.clone(),
            })
            .collect();

        OllamaRequest {
            model: self.model.clone(),
            messages: ollama_messages,
            temperature: options.temperature.clamp(0.0, 2.0),
            stream: false,
        }
    }

    fn build_stream_request(
        &self,
        messages: &[Message],
        options: &GenerationOptions,
    ) -> OllamaRequest {
        let mut req = self.build_request(messages, options);
        req.stream = true;
        req
    }
}

#[async_trait]
impl LlmProvider for OllamaProvider {
    async fn generate(&self, messages: &[Message], options: &GenerationOptions) -> Result<String> {
        let req = self.build_request(messages, options);
        let url = format!("{}/api/chat", self.base_url);

        let response = self
            .client
            .post(&url)
            .json(&req)
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() {
                    LlmError::Timeout("Ollama request timed out".into())
                } else {
                    LlmError::NetworkError(e)
                }
            })?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(LlmError::ApiError(format!(
                "Ollama error ({}): {}",
                status, error_text
            )));
        }

        let body: OllamaResponse = response
            .json()
            .await
            .map_err(|e| LlmError::ApiError(format!("Failed to parse response: {}", e)))?;

        Ok(body.message.content)
    }

    async fn stream(
        &self,
        messages: &[Message],
        options: &GenerationOptions,
    ) -> Result<StreamResponse> {
        let req = self.build_stream_request(messages, options);
        let url = format!("{}/api/chat", self.base_url);

        let response = self
            .client
            .post(&url)
            .json(&req)
            .send()
            .await
            .map_err(|e| {
                if e.is_timeout() {
                    LlmError::Timeout("Ollama request timed out".into())
                } else {
                    LlmError::NetworkError(e)
                }
            })?;

        let status = response.status();
        if !status.is_success() {
            return Err(LlmError::ApiError(format!("Ollama error: {}", status)));
        }

        let stream = response
            .bytes_stream()
            .filter_map(|result| async move {
                match result {
                    Ok(bytes) => {
                        let text = String::from_utf8_lossy(&bytes);
                        let text = text.trim();
                        if text.is_empty() {
                            return None;
                        }

                        if let Ok(chunk) = serde_json::from_str::<OllamaStreamResponse>(text) {
                            if !chunk.message.content.is_empty() {
                                return Some(Ok(chunk.message.content));
                            }
                        }
                        None
                    }
                    Err(e) => Some(Err(LlmError::StreamError(e.to_string()))),
                }
            })
            .boxed();

        Ok(Box::pin(stream))
    }

    fn name(&self) -> &str {
        "Ollama"
    }

    fn model(&self) -> &str {
        &self.model
    }

    async fn is_available(&self) -> bool {
        self.client
            .get(format!("{}/api/tags", self.base_url))
            .send()
            .await
            .map(|r| r.status().is_success())
            .unwrap_or(false)
    }
}

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

    #[test]
    fn test_ollama_new() {
        let provider = OllamaProvider::new("llama2");
        assert!(provider.is_ok());
    }

    #[test]
    fn test_ollama_empty_model() {
        let provider = OllamaProvider::new("");
        assert!(provider.is_err());
    }

    #[test]
    fn test_ollama_with_url() {
        let provider = OllamaProvider::with_url("llama2", "http://127.0.0.1:11434");
        assert!(provider.is_ok());
    }

    #[test]
    fn test_provider_name() {
        let provider = OllamaProvider::new("llama2").unwrap();
        assert_eq!(provider.name(), "Ollama");
        assert_eq!(provider.model(), "llama2");
    }

    #[test]
    fn test_build_request() {
        let provider = OllamaProvider::new("llama2").unwrap();
        let messages = vec![Message::user("Hello")];
        let options = GenerationOptions::default();

        let req = provider.build_request(&messages, &options);
        assert_eq!(req.model, "llama2");
        assert_eq!(req.messages.len(), 1);
        assert!(!req.stream);
    }
}
