//! Ollama API client implementation
//!
//! Ollama provides local LLM execution with an OpenAI-compatible API.
//! Default endpoint: http://localhost:11434

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::env;

use super::provider::{
    LlmProvider, LlmRequest, LlmResponse, LlmStream, MessageRole, StreamChunk, StreamMetadata,
    TokenUsage,
};
use crate::error::{Error, Result};
use futures::stream::StreamExt;

const DEFAULT_OLLAMA_URL: &str = "http://localhost:11434";

/// Ollama API client
pub struct OllamaProvider {
    base_url: String,
    client: reqwest::Client,
}

/// Ollama API request format (OpenAI-compatible)
#[derive(Debug, Serialize)]
struct OllamaRequest {
    model: String,
    messages: Vec<OllamaMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    options: Option<OllamaOptions>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
}

#[derive(Debug, Serialize)]
struct OllamaOptions {
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    num_predict: Option<i32>, // Ollama's equivalent of max_tokens
}

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

/// Ollama API response format
#[derive(Debug, Deserialize)]
struct OllamaResponse {
    model: String,
    #[allow(dead_code)]
    created_at: String,
    message: OllamaMessage,
    #[allow(dead_code)]
    done: bool,
    #[serde(default)]
    prompt_eval_count: Option<i64>,
    #[serde(default)]
    eval_count: Option<i64>,
}

impl OllamaProvider {
    /// Create a new Ollama provider with default localhost URL
    pub fn new() -> Result<Self> {
        let base_url =
            env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| DEFAULT_OLLAMA_URL.to_string());

        let client = reqwest::Client::new();

        Ok(Self { base_url, client })
    }

    /// Create a provider with custom base URL
    pub fn with_base_url(base_url: String) -> Self {
        let client = reqwest::Client::new();
        Self { base_url, client }
    }

    /// Build API URL for chat endpoint
    fn build_url(&self) -> String {
        format!("{}/api/chat", self.base_url)
    }
}

#[async_trait]
impl LlmProvider for OllamaProvider {
    async fn complete(&self, request: LlmRequest) -> Result<LlmResponse> {
        // Convert generic messages to Ollama format
        let mut messages: Vec<OllamaMessage> = request
            .messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System)
            .map(|m| OllamaMessage {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "assistant".to_string(),
                    MessageRole::System => "system".to_string(),
                },
                content: m.content,
            })
            .collect();

        // Add system message at the beginning if provided
        if let Some(system_prompt) = request.system {
            messages.insert(
                0,
                OllamaMessage {
                    role: "system".to_string(),
                    content: system_prompt,
                },
            );
        }

        // Build options
        let options = if request.max_tokens.is_some() || request.temperature.is_some() {
            Some(OllamaOptions {
                temperature: request.temperature,
                num_predict: request.max_tokens.map(|t| t as i32),
            })
        } else {
            None
        };

        // Build Ollama request
        let ollama_request = OllamaRequest {
            model: request.model.clone(),
            messages,
            options,
            stream: Some(false),
        };

        // Make API call
        let url = self.build_url();
        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&ollama_request)
            .send()
            .await
            .map_err(|e| {
                Error::execution(
                    "Failed to call Ollama API - is Ollama running?",
                    format!("{} (URL: {})", e, url),
                )
            })?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Ollama API error: {}", status),
                error_text,
            ));
        }

        // Parse response
        let ollama_response: OllamaResponse = response
            .json()
            .await
            .map_err(|e| Error::execution("Failed to parse Ollama API response", e.to_string()))?;

        // Calculate usage
        let usage = if ollama_response.prompt_eval_count.is_some()
            || ollama_response.eval_count.is_some()
        {
            let input_tokens = ollama_response.prompt_eval_count.unwrap_or(0) as usize;
            let output_tokens = ollama_response.eval_count.unwrap_or(0) as usize;
            Some(TokenUsage {
                input_tokens,
                output_tokens,
                total_tokens: input_tokens + output_tokens,
            })
        } else {
            None
        };

        Ok(LlmResponse {
            content: ollama_response.message.content,
            model: ollama_response.model,
            usage,
        })
    }

    async fn stream(&self, request: LlmRequest) -> Result<LlmStream> {
        use futures::stream;

        // Convert generic messages to Ollama format
        let mut messages: Vec<OllamaMessage> = request
            .messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System)
            .map(|m| OllamaMessage {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "assistant".to_string(),
                    MessageRole::System => "system".to_string(),
                },
                content: m.content,
            })
            .collect();

        // Add system message at the beginning if provided
        if let Some(system_prompt) = request.system {
            messages.insert(
                0,
                OllamaMessage {
                    role: "system".to_string(),
                    content: system_prompt,
                },
            );
        }

        // Build options
        let options = if request.max_tokens.is_some() || request.temperature.is_some() {
            Some(OllamaOptions {
                temperature: request.temperature,
                num_predict: request.max_tokens.map(|t| t as i32),
            })
        } else {
            None
        };

        // Build Ollama streaming request
        let ollama_request = OllamaRequest {
            model: request.model.clone(),
            messages,
            options,
            stream: Some(true),
        };

        // Make streaming API call
        let url = self.build_url();
        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&ollama_request)
            .send()
            .await
            .map_err(|e| {
                Error::execution(
                    "Failed to call Ollama API - is Ollama running?",
                    format!("{} (URL: {})", e, url),
                )
            })?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Ollama API error: {}", status),
                error_text,
            ));
        }

        // Convert response to JSON stream
        use futures::TryStreamExt;
        let byte_stream = response.bytes_stream();
        let model = request.model.clone();

        let json_stream = byte_stream
            .map_err(|e| Error::execution("Stream error", e.to_string()))
            .map(move |chunk_result| {
                match chunk_result {
                    Ok(bytes) => {
                        // Parse JSON events (newline-delimited like Gemini)
                        let text = String::from_utf8_lossy(&bytes);
                        parse_json_events(&text, &model)
                    }
                    Err(e) => vec![Err(e)],
                }
            })
            .flat_map(stream::iter);

        Ok(Box::pin(json_stream))
    }

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

/// Parse JSON events from Ollama streaming response
fn parse_json_events(text: &str, model: &str) -> Vec<Result<StreamChunk>> {
    let mut chunks = Vec::new();

    for line in text.lines() {
        if line.trim().is_empty() {
            continue;
        }

        // Parse JSON event
        if let Ok(event) = serde_json::from_str::<serde_json::Value>(line) {
            // Extract message content
            if let Some(message) = event.get("message") {
                if let Some(content) = message.get("content").and_then(|c| c.as_str()) {
                    if !content.is_empty() {
                        chunks.push(Ok(StreamChunk::Content(content.to_string())));
                    }
                }
            }

            // Check for completion (done: true)
            if let Some(true) = event.get("done").and_then(|d| d.as_bool()) {
                // Extract usage metadata
                let usage = Some(TokenUsage {
                    input_tokens: event
                        .get("prompt_eval_count")
                        .and_then(|v| v.as_i64())
                        .unwrap_or(0) as usize,
                    output_tokens: event
                        .get("eval_count")
                        .and_then(|v| v.as_i64())
                        .unwrap_or(0) as usize,
                    total_tokens: (event
                        .get("prompt_eval_count")
                        .and_then(|v| v.as_i64())
                        .unwrap_or(0)
                        + event
                            .get("eval_count")
                            .and_then(|v| v.as_i64())
                            .unwrap_or(0)) as usize,
                });

                chunks.push(Ok(StreamChunk::Done(StreamMetadata {
                    model: model.to_string(),
                    usage,
                })));
            }

            // Check for errors
            if let Some(error) = event.get("error") {
                let error_msg = error.as_str().unwrap_or("Unknown error");
                chunks.push(Ok(StreamChunk::Error(error_msg.to_string())));
            }
        }
    }

    chunks
}

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

    #[test]
    fn test_create_provider_with_default_url() {
        let provider = OllamaProvider::new().expect("Should create provider");
        assert_eq!(provider.base_url, DEFAULT_OLLAMA_URL);
        assert_eq!(provider.name(), "Ollama");
    }

    #[test]
    fn test_create_provider_with_custom_url() {
        let custom_url = "http://custom-host:8080";
        let provider = OllamaProvider::with_base_url(custom_url.to_string());
        assert_eq!(provider.base_url, custom_url);
        assert_eq!(provider.name(), "Ollama");
    }

    #[test]
    fn test_build_url() {
        let provider = OllamaProvider::with_base_url("http://localhost:11434".to_string());
        assert_eq!(provider.build_url(), "http://localhost:11434/api/chat");
    }

    #[test]
    fn test_parse_json_content() {
        let json_line = r#"{"model":"llama2","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":"Hello"},"done":false}"#;
        let model = "llama2";

        let chunks = parse_json_events(json_line, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected Content chunk"),
        }
    }

    #[test]
    fn test_parse_json_done() {
        let json_line = r#"{"model":"llama2","created_at":"2024-01-01T00:00:00Z","message":{"role":"assistant","content":""},"done":true,"prompt_eval_count":10,"eval_count":20}"#;
        let model = "llama2";

        let chunks = parse_json_events(json_line, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Done(metadata)) => {
                assert_eq!(metadata.model, model);
                let usage = metadata.usage.as_ref().unwrap();
                assert_eq!(usage.input_tokens, 10);
                assert_eq!(usage.output_tokens, 20);
                assert_eq!(usage.total_tokens, 30);
            }
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_json_error() {
        let json_line = r#"{"error":"model not found"}"#;
        let model = "llama2";

        let chunks = parse_json_events(json_line, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Error(msg)) => assert_eq!(msg, "model not found"),
            _ => panic!("Expected Error chunk"),
        }
    }

    #[test]
    fn test_parse_json_multiple_chunks() {
        let json_lines = r#"{"model":"llama2","message":{"role":"assistant","content":"Hello"},"done":false}
{"model":"llama2","message":{"role":"assistant","content":" world"},"done":false}
{"model":"llama2","message":{"role":"assistant","content":""},"done":true,"prompt_eval_count":5,"eval_count":10}"#;
        let model = "llama2";

        let chunks = parse_json_events(json_lines, model);

        assert_eq!(chunks.len(), 3);

        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected first Content chunk"),
        }

        match &chunks[1] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, " world"),
            _ => panic!("Expected second Content chunk"),
        }

        match &chunks[2] {
            Ok(StreamChunk::Done(metadata)) => {
                assert!(metadata.usage.is_some());
            }
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_json_empty_content_ignored() {
        let json_line =
            r#"{"model":"llama2","message":{"role":"assistant","content":""},"done":false}"#;
        let model = "llama2";

        let chunks = parse_json_events(json_line, model);

        // Empty content should not produce a chunk
        assert_eq!(chunks.len(), 0);
    }

    #[test]
    fn test_parse_json_empty_input() {
        let json_lines = "";
        let model = "llama2";

        let chunks = parse_json_events(json_lines, model);

        assert_eq!(chunks.len(), 0);
    }

    #[tokio::test]
    #[ignore] // Only run when Ollama is running locally
    async fn test_ollama_api_call() {
        let provider = OllamaProvider::new().expect("Should create provider");

        let request = LlmRequest {
            model: "llama2".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Say hello in exactly 3 words.".to_string(),
            }],
            max_tokens: Some(50),
            temperature: Some(0.7),
            system: None,
        };

        let response = provider.complete(request).await;
        assert!(response.is_ok());

        let response = response.unwrap();
        assert!(!response.content.is_empty());

        println!("Response: {}", response.content);
        println!("Usage: {:?}", response.usage);
    }
}
