use std::str::FromStr as _;

use async_trait::async_trait;
use base64::Engine as _;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use jp_attachment::AttachmentContent;
use jp_config::{
    assistant::tool_choice::ToolChoice,
    model::{
        id::{ModelIdConfig, Name, ProviderId},
        parameters::ReasoningConfig,
    },
    providers::llm::ollama::OllamaConfig,
};
use jp_conversation::{
    ConversationStream,
    event::{ChatResponse, EventKind},
    thread::text_attachments_to_xml,
};
use ollama_rs::{
    Ollama as Client,
    error::OllamaError,
    generation::{
        chat::{ChatMessage, ChatMessageResponse, MessageRole, request::ChatMessageRequest},
        images::Image,
        parameters::{FormatType, JsonStructure, KeepAlive, TimeUnit},
        tools::{ToolCall, ToolCallFunction, ToolFunctionInfo, ToolInfo, ToolType},
    },
    models::{LocalModel, ModelOptions},
};
use serde_json::Value;
use tracing::{debug, trace, warn};
use url::Url;

use super::{EventStream, ModelDetails, Provider, trace_to_tmpfile};
use crate::{
    error::{Error, Result, StreamError},
    event::{Event, FinishReason},
    query::ChatQuery,
    tool::ToolDefinition,
};

static PROVIDER: ProviderId = ProviderId::Ollama;

#[derive(Debug, Clone)]
pub struct Ollama {
    client: Client,
}

#[async_trait]
impl Provider for Ollama {
    async fn model_details(&self, name: &Name) -> Result<ModelDetails> {
        let id: ModelIdConfig = (PROVIDER, name.as_ref()).try_into()?;

        Ok(self
            .models()
            .await?
            .into_iter()
            .find(|m| m.id == id)
            .unwrap_or(ModelDetails::empty(id)))
    }

    async fn models(&self) -> Result<Vec<ModelDetails>> {
        let models = self.client.list_local_models().await?;

        models.into_iter().map(map_model).collect::<Result<_>>()
    }

    async fn chat_completion_stream(
        &self,
        model: &ModelDetails,
        query: ChatQuery,
    ) -> Result<EventStream> {
        debug!(
            model = %model.id.name,
            "Starting Ollama chat completion stream."
        );

        let (request, is_structured) = create_request(model, query)?;

        trace!(
            request = %trace_to_tmpfile("jp-ollama-request", &request),
            "Sending request to Ollama."
        );

        Ok(self
            .client
            .send_chat_messages_stream(request)
            .await?
            .map(|v| v.map_err(|()| StreamError::other("Ollama stream error")))
            .map_ok({
                let mut reasoning_flushed = false;
                move |v| {
                    stream::iter(
                        map_event(v, is_structured, &mut reasoning_flushed)
                            .into_iter()
                            .map(Ok),
                    )
                }
            })
            .try_flatten()
            .boxed())
    }
}

fn map_model(model: LocalModel) -> Result<ModelDetails> {
    Ok(ModelDetails {
        id: (PROVIDER, &model.name).try_into()?,
        display_name: Some(model.name),
        context_window: None,
        max_output_tokens: None,
        reasoning: None,
        knowledge_cutoff: None,
        deprecated: None,
        structured_output: None,
        features: vec![],
    })
}

/// Map an Ollama streaming chunk into provider-agnostic events.
///
/// Index convention: 0 = reasoning, 1 = message content, 2+ = tool calls.
///
/// Ollama guarantees that thinking tokens arrive before content/tool call
/// tokens.
/// We exploit this by flushing the reasoning stream (index 0) as soon as the
/// first content or tool call chunk appears, ensuring the reasoning event
/// precedes content and tool call events in the history.
fn map_event(
    event: ChatMessageResponse,
    is_structured: bool,
    reasoning_flushed: &mut bool,
) -> Vec<Event> {
    let ChatMessageResponse { message, done, .. } = event;

    trace!(
        content = message.content,
        thinking = message.thinking,
        tool_calls = message.tool_calls.len(),
        done,
        "Ollama stream chunk."
    );

    let mut events = Vec::new();

    if let Some(thinking) = message.thinking
        && !thinking.is_empty()
    {
        events.push(Event::reasoning(0, thinking));
    }

    let has_content = !message.content.is_empty();
    let has_tool_calls = !message.tool_calls.is_empty();

    // Flush reasoning before emitting content or tool calls so the reasoning
    // event always precedes them in the conversation history.
    if !*reasoning_flushed && (has_content || has_tool_calls) {
        events.push(Event::flush(0));
        *reasoning_flushed = true;
    }

    if has_content {
        if is_structured {
            events.push(Event::structured(1, message.content));
        } else {
            events.push(Event::message(1, message.content));
        }
    }

    for (
        idx,
        ToolCall {
            function: ToolCallFunction { name, arguments },
        },
    ) in message.tool_calls.into_iter().enumerate()
    {
        let index = idx + 2;
        let id = format!("{name}_{index}");
        let args_json = match arguments {
            Value::Object(map) => serde_json::to_string(&map).unwrap_or_default(),
            v => serde_json::to_string(&v).unwrap_or_default(),
        };
        events.push(Event::tool_call_start(index, id, &name));
        events.push(Event::tool_call_args(index, args_json));
        events.push(Event::flush(index));
    }

    if done {
        if !*reasoning_flushed {
            events.push(Event::flush(0));
        }
        events.push(Event::flush(1));
        // `done` is Ollama's terminal signal. Emit `Finished` from it so the
        // stream carries its own completion: a stream that ends without it (a
        // dropped connection) is treated as incomplete by the consumer and
        // retried, rather than being mistaken for a clean finish.
        events.push(Event::Finished(FinishReason::Completed));
    }

    events
}

#[expect(clippy::too_many_lines)]
fn create_request(model: &ModelDetails, query: ChatQuery) -> Result<(ChatMessageRequest, bool)> {
    let ChatQuery {
        thread,
        tools,
        tool_choice,
    } = query;

    let structured_schema = thread
        .events
        .schema()
        .map(|schema| JsonStructure::new_for_schema(schema.into()))
        .map(|schema| FormatType::StructuredJson(Box::new(schema)));

    let is_structured = structured_schema.is_some();

    let config = thread.events.config()?;
    let parameters = &config.assistant.model.parameters;

    let parts = thread.into_parts();

    let mut system_parts = parts.system_parts;
    if let Some(xml) = text_attachments_to_xml(&parts.attachments)? {
        system_parts.push(xml);
    }

    let mut messages: Vec<ChatMessage> = to_system_messages(system_parts).collect();

    // Prepend binary image attachments as a user message with base64 images.
    let images: Vec<_> = parts
        .attachments
        .iter()
        .filter_map(|a| match &a.content {
            AttachmentContent::Binary { data, media_type } if media_type.starts_with("image/") => {
                Some(Image::from_base64(
                    base64::engine::general_purpose::STANDARD.encode(data),
                ))
            }
            AttachmentContent::Binary { media_type, .. } => {
                warn!(
                    source = %a.source,
                    media_type,
                    "Unsupported binary attachment media type for Ollama, skipping."
                );
                None
            }
            AttachmentContent::Text(_) => None,
        })
        .collect();

    if !images.is_empty() {
        messages.push(ChatMessage {
            role: MessageRole::User,
            content: String::new(),
            tool_calls: vec![],
            images: Some(images),
            thinking: None,
        });
    }

    messages.extend(convert_events(parts.events));

    if let Some(tool_choice) = tool_choice_to_system_message(&tool_choice) {
        messages.push(tool_choice);
    }

    let mut request = ChatMessageRequest::new(model.id.name.to_string(), messages);

    let tools = convert_tools(tools)?;
    if !tools.is_empty() {
        request = request.tools(tools);
    }

    let mut options = ModelOptions::default();

    if let Some(temperature) = parameters.temperature {
        options = options.temperature(temperature);
    }

    if let Some(top_p) = parameters.top_p {
        options = options.top_p(top_p);
    }

    if let Some(top_k) = parameters.top_k {
        options = options.top_k(top_k);
    }

    // Set the context window for the model.
    //
    // This can be used to force Ollama to use a larger context window then the
    // one determined based on the machine's resources.
    if let Some(context_window) = parameters
        .other
        .get("context_window")
        .and_then(|v| v.as_u64())
    {
        options = options.num_ctx(context_window);
    }

    if let Some(keep_alive) = parameters.other.get("keep_alive").and_then(|v| v.as_str()) {
        let unit = keep_alive
            .chars()
            .last()
            .filter(char::is_ascii_alphabetic)
            .unwrap_or('m');

        let value = keep_alive
            .chars()
            .take_while(char::is_ascii_alphanumeric)
            .collect::<String>();

        request = request.keep_alive(KeepAlive::Until {
            time: value.parse::<u64>().unwrap_or(5),
            unit: match unit {
                's' => TimeUnit::Seconds,
                'h' => TimeUnit::Hours,
                _ => TimeUnit::Minutes,
            },
        });
    }

    request = request.options(options);

    // Ollama models may default to thinking-on, so we must explicitly set the
    // flag in both directions. `None` is treated as off because we can't know
    // whether an arbitrary local model supports reasoning.
    request = request.think(!matches!(
        parameters.reasoning,
        None | Some(ReasoningConfig::Off)
    ));

    if let Some(schema) = structured_schema {
        request = request.format(schema);
    }

    Ok((request, is_structured))
}

impl TryFrom<&OllamaConfig> for Ollama {
    type Error = Error;

    fn try_from(config: &OllamaConfig) -> Result<Self> {
        let url = Url::from_str(&config.base_url)?;
        let port = url.port().unwrap_or(11434);
        let client = reqwest::Client::new();

        Ok(Ollama {
            client: Client::new_with_client(url, port, client),
        })
    }
}

fn convert_tools(tools: Vec<ToolDefinition>) -> Result<Vec<ToolInfo>> {
    tools
        .into_iter()
        .map(|tool| {
            Ok(ToolInfo {
                tool_type: ToolType::Function,
                function: ToolFunctionInfo {
                    parameters: tool.to_parameters_map().into(),
                    name: tool.name,
                    description: tool
                        .docs
                        .schema_description()
                        .unwrap_or_default()
                        .to_owned(),
                },
            })
        })
        .collect::<Result<Vec<_>>>()
}

/// Poor-man's version of API-based tool choice.
/// Needed until Ollama has first-class support for tool choice.
fn tool_choice_to_system_message(choice: &ToolChoice) -> Option<ChatMessage> {
    let (ToolChoice::Function(_) | ToolChoice::Required) = choice else {
        return None;
    };

    let msg = if let Some(tool) = choice.function_name() {
        format!("You MUST use the function named '{tool}' available to you.")
    } else {
        "You MUST use AT LEAST ONE tool available to you.".to_string()
    };

    let content = format!(
        "IMPORTANT: {msg} DO NOT QUESTION THIS DIRECTIVE. DO NOT PROMPT FOR MORE CONTEXT OR \
         DETAILS. JUST RUN IT."
    );

    Some(ChatMessage {
        role: MessageRole::System,
        content,
        tool_calls: vec![],
        images: None,
        thinking: None,
    })
}

/// Convert some content into a system message.
fn to_system_messages(parts: Vec<String>) -> impl Iterator<Item = ChatMessage> {
    parts.into_iter().map(|content| ChatMessage {
        role: MessageRole::System,
        content,
        tool_calls: vec![],
        images: None,
        thinking: None,
    })
}

/// Convert a conversation stream into a list of messages.
fn convert_events(events: ConversationStream) -> Vec<ChatMessage> {
    events
        .into_iter()
        .filter_map(|event| match event.into_kind() {
            EventKind::ChatRequest(request) => Some(ChatMessage::user(request.content)),
            EventKind::ChatResponse(response) => match response {
                ChatResponse::Message { message } => Some(ChatMessage::assistant(message)),
                ChatResponse::Reasoning { reasoning, .. } => Some(ChatMessage {
                    role: MessageRole::Assistant,
                    content: String::new(),
                    tool_calls: vec![],
                    images: None,
                    thinking: Some(reasoning),
                }),
                ChatResponse::Structured { data } => Some(ChatMessage::assistant(data.to_string())),
            },
            EventKind::ToolCallRequest(request) => Some(ChatMessage {
                role: MessageRole::Assistant,
                content: String::new(),
                tool_calls: vec![ToolCall {
                    function: ToolCallFunction {
                        name: request.name,
                        arguments: Value::Object(request.arguments),
                    },
                }],
                images: None,
                thinking: None,
            }),
            EventKind::ToolCallResponse(response) => {
                Some(ChatMessage::tool(match response.result {
                    Ok(content) => content,
                    Err(error) => error,
                }))
            }
            _ => None,
        })
        .fold(vec![], |mut messages, message| match messages.last_mut() {
            Some(last)
                if last.role == message.role
                    && message.thinking.is_some()
                    && last.thinking.is_none() =>
            {
                last.thinking = message.thinking;
                messages
            }
            _ => {
                messages.push(message);
                messages
            }
        })
}

impl From<OllamaError> for StreamError {
    fn from(err: OllamaError) -> Self {
        use ollama_rs::error::InternalOllamaError;

        match err {
            OllamaError::ReqwestError(error) => Self::from(error),
            OllamaError::ToolCallError(error) => {
                StreamError::other("tool-call error").with_source(error)
            }
            OllamaError::JsonError(error) => StreamError::other("json error").with_source(error),
            OllamaError::InternalError(InternalOllamaError { message }) => {
                StreamError::other(message)
            }
            OllamaError::Other(message) => StreamError::other(message),
        }
    }
}
