use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use swarm_core::error::SwarmResult;
use swarm_core::message::ToolCall;
use tracing::{debug, trace};

use crate::provider::{
    ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, StopReason, ToolDefinition,
    Usage,
};

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

impl OllamaProvider {
    pub fn new(base_url: String) -> Self {
        Self {
            client: Client::new(),
            base_url,
        }
    }
}

#[derive(Serialize)]
struct OllamaRequest {
    model: String,
    messages: Vec<OllamaMessage>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    options: Option<OllamaOptions>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tools: Vec<OllamaTool>,
}

#[derive(Serialize)]
struct OllamaOptions {
    temperature: f32,
    num_predict: u32,
}

#[derive(Serialize, Deserialize)]
struct OllamaMessage {
    role: String,
    content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OllamaToolCall>>,
}

#[derive(Serialize, Deserialize)]
struct OllamaToolCall {
    function: OllamaFunction,
}

#[derive(Serialize, Deserialize)]
struct OllamaFunction {
    name: String,
    arguments: serde_json::Value,
}

#[derive(Serialize)]
struct OllamaTool {
    #[serde(rename = "type")]
    kind: String,
    function: OllamaToolDef,
}

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

#[derive(Deserialize)]
struct OllamaResponse {
    message: OllamaMessage,
    model: String,
    #[serde(default)]
    prompt_eval_count: Option<u32>,
    #[serde(default)]
    eval_count: Option<u32>,
}

fn convert_messages(messages: &[ChatMessage], system: &Option<String>) -> Vec<OllamaMessage> {
    let mut result = Vec::new();

    if let Some(sys) = system {
        result.push(OllamaMessage {
            role: "system".to_string(),
            content: sys.clone(),
            tool_calls: None,
        });
    }

    for msg in messages {
        if !msg.tool_results.is_empty() {
            for tr in &msg.tool_results {
                result.push(OllamaMessage {
                    role: "tool".to_string(),
                    content: tr.content.clone(),
                    tool_calls: None,
                });
            }
        } else {
            result.push(OllamaMessage {
                role: msg.role.clone(),
                content: msg.content.clone(),
                tool_calls: None,
            });
        }
    }

    result
}

fn convert_tools(tools: &[ToolDefinition]) -> Vec<OllamaTool> {
    tools
        .iter()
        .map(|t| OllamaTool {
            kind: "function".to_string(),
            function: OllamaToolDef {
                name: t.name.clone(),
                description: t.description.clone(),
                parameters: t.parameters.clone(),
            },
        })
        .collect()
}

#[async_trait]
impl LlmProvider for OllamaProvider {
    fn name(&self) -> &str {
        "ollama"
    }

    async fn complete(&self, request: CompletionRequest) -> SwarmResult<CompletionResponse> {
        let tool_defs = request.tools.clone();
        let api_request = OllamaRequest {
            model: request.model,
            messages: convert_messages(&request.messages, &request.system),
            stream: false,
            options: Some(OllamaOptions {
                temperature: request.temperature,
                num_predict: request.max_tokens,
            }),
            tools: convert_tools(&request.tools),
        };

        let response = self
            .client
            .post(format!("{}/api/chat", self.base_url))
            .json(&api_request)
            .send()
            .await
            .map_err(|e| swarm_core::error::SwarmError::LlmError(e.to_string()))?;

        let status = response.status();
        let body = response
            .text()
            .await
            .map_err(|e| swarm_core::error::SwarmError::LlmError(e.to_string()))?;

        if !status.is_success() {
            return Err(swarm_core::error::SwarmError::LlmError(format!(
                "Ollama API {}: {}",
                status, body
            )));
        }

        debug!(model = %api_request.model, tools = api_request.tools.len(), "Ollama request sent");
        trace!(raw_response = %body, "Ollama raw response");

        let api_response: OllamaResponse = serde_json::from_str(&body)
            .map_err(|e| swarm_core::error::SwarmError::LlmError(e.to_string()))?;

        let mut tool_calls = Vec::new();
        if let Some(calls) = &api_response.message.tool_calls {
            for (i, call) in calls.iter().enumerate() {
                tool_calls.push(ToolCall {
                    id: format!("tool_{i}"),
                    name: call.function.name.clone(),
                    arguments: call.function.arguments.clone(),
                });
            }
        }

        // Fallback: some models emit tool calls as JSON text instead of structured tool_calls.
        // Try to extract them from the content.
        if tool_calls.is_empty()
            && !tool_defs.is_empty()
            && let Some(extracted) = extract_tool_call_from_text(&api_response.message.content, &tool_defs)
        {
            debug!(tool = %extracted.name, "Extracted tool call from text content");
            tool_calls.push(extracted);
        }

        let stop_reason = if !tool_calls.is_empty() {
            StopReason::ToolUse
        } else {
            StopReason::EndTurn
        };

        debug!(
            content_len = api_response.message.content.len(),
            tool_calls = tool_calls.len(),
            stop = ?stop_reason,
            "Ollama response parsed"
        );

        Ok(CompletionResponse {
            content: if !tool_calls.is_empty() { String::new() } else { api_response.message.content },
            tool_calls,
            model: api_response.model,
            usage: Usage {
                input_tokens: api_response.prompt_eval_count.unwrap_or(0),
                output_tokens: api_response.eval_count.unwrap_or(0),
            },
            stop_reason,
        })
    }
}

/// Try to extract a tool call from text content when the model doesn't use structured tool calls.
/// Looks for JSON objects with "name" and "arguments" fields, or just raw JSON matching a known tool.
fn extract_tool_call_from_text(content: &str, tools: &[ToolDefinition]) -> Option<ToolCall> {
    let trimmed = content.trim();

    // Try to find a JSON object in the content
    let json_str = if let Some(start) = trimmed.find('{') {
        if let Some(end) = trimmed.rfind('}') {
            &trimmed[start..=end]
        } else {
            return None;
        }
    } else {
        return None;
    };

    // Try parsing as {"name": "...", "arguments": {...}}
    if let Ok(obj) = serde_json::from_str::<serde_json::Value>(json_str) {
        if let (Some(name), Some(args)) = (obj.get("name").and_then(|n| n.as_str()), obj.get("arguments")) {
            // Verify it's a known tool
            if tools.iter().any(|t| t.name == name) {
                return Some(ToolCall {
                    id: "text_extract_0".to_string(),
                    name: name.to_string(),
                    arguments: args.clone(),
                });
            }
        }

        // Try as {"tool": "...", "args": {...}} variant
        if let (Some(name), Some(args)) = (obj.get("tool").and_then(|n| n.as_str()), obj.get("args"))
            && tools.iter().any(|t| t.name == name)
        {
            return Some(ToolCall {
                id: "text_extract_0".to_string(),
                name: name.to_string(),
                arguments: args.clone(),
            });
        }

        // Try as direct tool arguments: if only one tool and the JSON looks like its params
        if tools.len() == 1 {
            return Some(ToolCall {
                id: "text_extract_0".to_string(),
                name: tools[0].name.clone(),
                arguments: obj,
            });
        }

        // Try matching by checking if the JSON has fields matching a tool's param schema
        for tool in tools {
            if let Some(props) = tool.parameters.get("properties")
                && let Some(props_obj) = props.as_object()
            {
                let matching = props_obj.keys().filter(|k| obj.get(*k).is_some()).count();
                let required_count = tool.parameters.get("required")
                    .and_then(|r| r.as_array())
                    .map(|a| a.len())
                    .unwrap_or(0);
                if matching > 0 && matching >= required_count {
                    return Some(ToolCall {
                        id: "text_extract_0".to_string(),
                        name: tool.name.clone(),
                        arguments: obj,
                    });
                }
            }
        }
    }

    None
}
