//! Ollama [`InferenceProvider`].
//!
//! Talks to a local Ollama server over its `/api/chat` endpoint and
//! translates the neutral [`Message`] / [`ToolSpec`] types to and from
//! Ollama's OpenAI-shaped wire format.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

use crate::agent::{InferenceProvider, Message, Role, ToolCall};
use crate::error::Result;
use crate::tool::ToolSpec;

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

pub struct OllamaProvider {
    client: reqwest::Client,
    base_url: String,
    model: String,
}

impl OllamaProvider {
    pub fn new(model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            base_url: DEFAULT_OLLAMA_URL.to_string(),
            model: model.into(),
        }
    }

    #[must_use]
    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }
}

#[async_trait]
impl InferenceProvider for OllamaProvider {
    async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message> {
        let wire_messages: Vec<WireMessage> =
            messages.iter().map(WireMessage::from_neutral).collect();
        let wire_tools: Vec<WireTool> = tools
            .iter()
            .map(|s| WireTool {
                r#type: "function",
                function: WireFunction {
                    name: s.name.clone(),
                    description: s.description.clone(),
                    parameters: s.input_schema.clone(),
                },
            })
            .collect();

        let req = ChatRequest {
            model: &self.model,
            messages: &wire_messages,
            tools: &wire_tools,
            stream: false,
        };

        let resp: ChatResponse = self
            .client
            .post(format!("{}/api/chat", self.base_url))
            .json(&req)
            .send()
            .await?
            .error_for_status()?
            .json()
            .await?;

        Ok(resp.message.into_neutral())
    }
}

// ---- wire types ----

#[derive(Debug, Serialize)]
struct ChatRequest<'a> {
    model: &'a str,
    messages: &'a [WireMessage],
    tools: &'a [WireTool],
    stream: bool,
}

#[derive(Debug, Deserialize)]
struct ChatResponse {
    message: WireMessage,
}

#[derive(Debug, Serialize, Deserialize)]
struct WireMessage {
    role: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    content: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<WireToolCall>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    tool_name: Option<String>,
}

impl WireMessage {
    fn from_neutral(m: &Message) -> Self {
        let role = match m.role {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
        }
        .to_string();
        let tool_calls = if m.tool_calls.is_empty() {
            None
        } else {
            Some(
                m.tool_calls
                    .iter()
                    .map(|c| WireToolCall {
                        function: WireToolCallFunction {
                            name: c.name.clone(),
                            arguments: c.arguments.clone(),
                        },
                    })
                    .collect(),
            )
        };
        Self {
            role,
            content: m.content.clone(),
            tool_calls,
            tool_name: m.tool_name.clone(),
        }
    }

    fn into_neutral(self) -> Message {
        let role = match self.role.as_str() {
            "system" => Role::System,
            "user" => Role::User,
            "tool" => Role::Tool,
            _ => Role::Assistant,
        };
        let tool_calls = self
            .tool_calls
            .unwrap_or_default()
            .into_iter()
            .map(|c| ToolCall {
                id: String::new(),
                name: c.function.name,
                arguments: c.function.arguments,
            })
            .collect();
        Message {
            role,
            content: self.content,
            tool_calls,
            tool_name: self.tool_name,
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct WireToolCall {
    function: WireToolCallFunction,
}

#[derive(Debug, Serialize, Deserialize)]
struct WireToolCallFunction {
    name: String,
    #[serde(default = "empty_object")]
    arguments: Value,
}

fn empty_object() -> Value {
    json!({})
}

#[derive(Debug, Serialize)]
struct WireTool {
    r#type: &'static str,
    function: WireFunction,
}

#[derive(Debug, Serialize)]
struct WireFunction {
    name: String,
    description: String,
    parameters: Value,
}