use anyhow::{Result, bail};
use serde_json::{Map, Value, json};

use crate::llm::protocols::shared::{assistant_parts, tool_result_text};
use crate::llm::protocols::utils::provider_options;
use crate::llm::schema::{FinishReason, LlmEvent, ToolCall, Usage};
use crate::llm::{GenerationOptions, LlmRequest, Message, MessageContent, ToolChoice, ToolSpec};

const ROUTE: &str = "gemini";

#[derive(Debug, Default)]
pub(crate) struct StreamState {
    finish_reason: Option<String>,
    has_tool_calls: bool,
    next_tool_call_id: usize,
    usage: Option<Usage>,
}

pub(crate) fn endpoint_path(model: &str) -> String {
    format!("models/{model}:streamGenerateContent")
}

pub(crate) fn request_body(request: &LlmRequest) -> Result<Value> {
    let tools_enabled =
        !request.tools.is_empty() && !matches!(request.tool_choice, Some(ToolChoice::None));
    let mut body = Map::from_iter([(
        "contents".to_string(),
        Value::Array(lower_messages(request)?),
    )]);

    if !request.system_prompt.trim().is_empty() {
        body.insert(
            "systemInstruction".to_string(),
            json!({"parts": [{"text": request.system_prompt}]}),
        );
    }
    if tools_enabled {
        body.insert(
            "tools".to_string(),
            json!([{ "functionDeclarations": lower_tools(&request.tools) }]),
        );
        if let Some(tool_config) = lower_tool_config(request.tool_choice.as_ref())? {
            body.insert("toolConfig".to_string(), tool_config);
        }
    }
    if let Some(generation_config) = lower_generation_options(request.generation.as_ref()) {
        body.insert("generationConfig".to_string(), generation_config);
    }
    provider_options::merge_json_body(ROUTE, &mut body, request.route.additional_params.as_ref())?;
    Ok(Value::Object(body))
}

pub(crate) fn parse_stream_event(state: &mut StreamState, event: &Value) -> Result<Vec<LlmEvent>> {
    let mut events = Vec::new();
    if let Some(usage) = event.get("usageMetadata") {
        state.usage = Some(Usage::from_gemini(usage));
    }
    let candidate = event
        .get("candidates")
        .and_then(Value::as_array)
        .and_then(|items| items.first());
    let Some(candidate) = candidate else {
        return Ok(events);
    };
    if let Some(reason) = candidate.get("finishReason").and_then(Value::as_str) {
        state.finish_reason = Some(reason.to_string());
    }
    let parts = candidate
        .pointer("/content/parts")
        .and_then(Value::as_array)
        .map(Vec::as_slice)
        .unwrap_or(&[]);
    for part in parts {
        if let Some(text) = part
            .get("text")
            .and_then(Value::as_str)
            .filter(|text| !text.is_empty())
        {
            if part
                .get("thought")
                .and_then(Value::as_bool)
                .unwrap_or(false)
            {
                events.push(LlmEvent::ReasoningDelta {
                    text: text.to_string(),
                });
            } else {
                events.push(LlmEvent::TextDelta {
                    text: text.to_string(),
                });
            }
        }
        if let Some(call) = part.get("functionCall") {
            let name = call
                .get("name")
                .and_then(Value::as_str)
                .unwrap_or_default()
                .to_string();
            if name.trim().is_empty() {
                bail!("Gemini tool call is missing a name");
            }
            let id = format!("tool_{}", state.next_tool_call_id);
            state.next_tool_call_id += 1;
            let arguments =
                serde_json::to_string(call.get("args").unwrap_or(&Value::Object(Map::new())))?;
            let mut tool_call = ToolCall::from_raw_input(id, name, &arguments, ROUTE)?;
            tool_call.signature = thought_signature(part);
            state.has_tool_calls = true;
            events.push(LlmEvent::ToolCall {
                call: tool_call,
                provider_executed: false,
            });
        }
    }
    Ok(events)
}

fn thought_signature(part: &Value) -> Option<String> {
    part.get("thoughtSignature")
        .and_then(Value::as_str)
        .map(str::to_string)
}

pub(crate) fn finish_stream(state: &mut StreamState) -> Result<Vec<LlmEvent>> {
    if state.finish_reason.is_none() && state.usage.is_none() {
        return Ok(Vec::new());
    }
    Ok(vec![LlmEvent::StepFinish {
        reason: map_finish_reason(state.finish_reason.as_deref(), state.has_tool_calls),
        usage: state.usage.clone(),
    }])
}

fn lower_messages(request: &LlmRequest) -> Result<Vec<Value>> {
    let mut contents = Vec::new();
    for message in &request.messages {
        match message {
            Message::System { .. } => {}
            Message::User { content } => contents.push(json!({
                "role": "user",
                "parts": lower_user_content(content)?,
            })),
            Message::Assistant { content, .. } => contents.push(json!({
                "role": "model",
                "parts": lower_assistant_content(content)?,
            })),
        }
    }
    Ok(contents)
}

fn lower_user_content(content: &[MessageContent]) -> Result<Vec<Value>> {
    let mut parts = Vec::new();
    for item in content {
        match item {
            MessageContent::Text { text, .. } => parts.push(json!({"text": text})),
            MessageContent::ToolResult {
                id,
                call_id,
                content,
                ..
            } => {
                let name = call_id.as_ref().unwrap_or(id);
                parts.push(json!({
                    "functionResponse": {
                        "name": name,
                        "response": {
                            "name": name,
                            "content": tool_result_text(content.clone())?,
                        }
                    }
                }));
            }
            MessageContent::Opaque { value, .. } => {
                parts.push(json!({"text": serde_json::to_string(value)?}))
            }
            MessageContent::ToolCall { .. } => {
                bail!("Gemini user message cannot contain a tool call")
            }
            MessageContent::Reasoning { .. } => {
                bail!("Gemini user message cannot contain reasoning")
            }
        }
    }
    Ok(parts)
}

fn lower_assistant_content(content: &[MessageContent]) -> Result<Vec<Value>> {
    let assistant = assistant_parts(content.to_vec())?;
    let mut parts = Vec::new();
    if let Some(reasoning) = assistant.reasoning_content
        && let Some(text) = reasoning.as_str()
    {
        parts.push(json!({"text": text, "thought": true}));
    }
    if !assistant.text.is_empty() {
        parts.push(json!({"text": assistant.text}));
    }
    for call in assistant.tool_calls {
        let mut part = Map::from_iter([(
            "functionCall".to_string(),
            json!({"name": call.name, "args": call.arguments_value()?}),
        )]);
        insert_thought_signature(&mut part, call.signature.as_deref());
        parts.push(Value::Object(part));
    }
    if parts.is_empty() {
        parts.push(json!({"text": ""}));
    }
    Ok(parts)
}

fn insert_thought_signature(part: &mut Map<String, Value>, signature: Option<&str>) {
    if let Some(signature) = signature.filter(|signature| !signature.is_empty()) {
        part.insert(
            "thoughtSignature".to_string(),
            Value::String(signature.to_string()),
        );
    }
}

fn lower_tools(tools: &[ToolSpec]) -> Vec<Value> {
    tools
        .iter()
        .map(|tool| {
            json!({
                "name": tool.name,
                "description": tool.description,
                "parameters": convert_tool_schema(&tool.parameters),
            })
        })
        .collect()
}

fn lower_tool_config(tool_choice: Option<&ToolChoice>) -> Result<Option<Value>> {
    Ok(match tool_choice {
        None | Some(ToolChoice::Auto) => Some(json!({"functionCallingConfig": {"mode": "AUTO"}})),
        Some(ToolChoice::None) => None,
        Some(ToolChoice::Required) => Some(json!({"functionCallingConfig": {"mode": "ANY"}})),
        Some(ToolChoice::Tool { name }) => {
            if name.trim().is_empty() {
                bail!("Gemini tool choice requires a tool name");
            }
            Some(json!({"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": [name]}}))
        }
    })
}

fn lower_generation_options(generation: Option<&GenerationOptions>) -> Option<Value> {
    let generation = generation?;
    let mut object = Map::new();
    if let Some(value) = generation.max_tokens {
        object.insert("maxOutputTokens".to_string(), json!(value));
    }
    if let Some(value) = generation.temperature {
        object.insert("temperature".to_string(), json!(value));
    }
    if let Some(value) = generation.top_p {
        object.insert("topP".to_string(), json!(value));
    }
    if let Some(value) = generation.top_k {
        object.insert("topK".to_string(), json!(value));
    }
    if let Some(stop) = generation.stop.as_ref().filter(|stop| !stop.is_empty()) {
        object.insert("stopSequences".to_string(), json!(stop));
    }
    (!object.is_empty()).then_some(Value::Object(object))
}

fn map_finish_reason(reason: Option<&str>, has_tool_calls: bool) -> FinishReason {
    match reason {
        Some("STOP") if has_tool_calls => FinishReason::ToolCalls,
        Some("STOP") => FinishReason::Stop,
        Some("MAX_TOKENS") => FinishReason::Length,
        Some(
            "IMAGE_SAFETY" | "RECITATION" | "SAFETY" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "SPII",
        ) => FinishReason::ContentFilter,
        _ => FinishReason::Unknown,
    }
}

fn convert_tool_schema(schema: &Value) -> Value {
    project_schema(&sanitize_schema(schema)).unwrap_or(Value::Null)
}

fn sanitize_schema(schema: &Value) -> Value {
    match schema {
        Value::Array(items) => Value::Array(items.iter().map(sanitize_schema).collect()),
        Value::Object(object) => {
            let mut result = object
                .iter()
                .map(|(key, value)| {
                    let value = if key == "enum" {
                        value
                            .as_array()
                            .map(|items| {
                                Value::Array(items.iter().map(stringify_schema_enum).collect())
                            })
                            .unwrap_or_else(|| sanitize_schema(value))
                    } else {
                        sanitize_schema(value)
                    };
                    (key.clone(), value)
                })
                .collect::<Map<_, _>>();
            let schema_type = result
                .get("type")
                .and_then(Value::as_str)
                .map(str::to_string);
            if matches!(schema_type.as_deref(), Some("integer" | "number"))
                && result.get("enum").is_some()
            {
                result.insert("type".to_string(), json!("string"));
            }
            if schema_type.as_deref() == Some("object")
                && let (Some(Value::Object(properties)), Some(Value::Array(required))) =
                    (result.get("properties"), result.get("required"))
            {
                let filtered = required
                    .iter()
                    .filter(|item| {
                        item.as_str()
                            .is_some_and(|field| properties.contains_key(field))
                    })
                    .cloned()
                    .collect();
                result.insert("required".to_string(), Value::Array(filtered));
            }
            if schema_type.as_deref() == Some("array")
                && !has_combiner(&Value::Object(result.clone()))
            {
                let items = result
                    .entry("items".to_string())
                    .or_insert_with(|| json!({}));
                if items
                    .as_object()
                    .is_some_and(|object| !has_schema_intent_object(object))
                {
                    *items = json!({"type": "string"});
                }
            }
            if let Some(kind) = schema_type.as_deref()
                && kind != "object"
                && !has_combiner(&Value::Object(result.clone()))
            {
                result.remove("properties");
                result.remove("required");
            }
            Value::Object(result)
        }
        value => value.clone(),
    }
}

fn stringify_schema_enum(value: &Value) -> Value {
    match value {
        Value::String(value) => Value::String(value.clone()),
        _ => Value::String(value.to_string()),
    }
}

fn project_schema(schema: &Value) -> Option<Value> {
    let object = schema.as_object()?;
    let mut out = Map::new();
    copy_if_present(&mut out, object, "description");
    copy_if_present(&mut out, object, "required");
    copy_if_present(&mut out, object, "format");
    if let Some(value) = object.get("type") {
        if let Some(types) = value.as_array() {
            if let Some(kind) = types.iter().find(|item| item.as_str() != Some("null")) {
                out.insert("type".to_string(), kind.clone());
            }
            if types.iter().any(|item| item.as_str() == Some("null")) {
                out.insert("nullable".to_string(), Value::Bool(true));
            }
        } else {
            out.insert("type".to_string(), value.clone());
        }
    }
    if let Some(value) = object.get("const") {
        out.insert("enum".to_string(), Value::Array(vec![value.clone()]));
    } else {
        copy_if_present(&mut out, object, "enum");
    }
    if let Some(properties) = object.get("properties").and_then(Value::as_object) {
        let projected = properties
            .iter()
            .filter_map(|(key, value)| project_schema(value).map(|value| (key.clone(), value)))
            .collect::<Map<_, _>>();
        out.insert("properties".to_string(), Value::Object(projected));
    }
    if let Some(items) = object.get("items") {
        let projected = if let Some(items) = items.as_array() {
            Value::Array(items.iter().filter_map(project_schema).collect())
        } else {
            project_schema(items).unwrap_or(Value::Null)
        };
        out.insert("items".to_string(), projected);
    }
    for key in ["allOf", "anyOf", "oneOf"] {
        if let Some(items) = object.get(key).and_then(Value::as_array) {
            out.insert(
                key.to_string(),
                Value::Array(items.iter().filter_map(project_schema).collect()),
            );
        }
    }
    copy_if_present(&mut out, object, "minLength");
    Some(Value::Object(out))
}

fn copy_if_present(out: &mut Map<String, Value>, object: &Map<String, Value>, key: &str) {
    if let Some(value) = object.get(key) {
        out.insert(key.to_string(), value.clone());
    }
}

fn has_combiner(schema: &Value) -> bool {
    schema.as_object().is_some_and(|object| {
        ["anyOf", "oneOf", "allOf"]
            .iter()
            .any(|key| object.get(*key).is_some_and(Value::is_array))
    })
}

fn has_schema_intent_object(object: &Map<String, Value>) -> bool {
    has_combiner(&Value::Object(object.clone()))
        || [
            "type",
            "properties",
            "items",
            "prefixItems",
            "enum",
            "const",
            "$ref",
            "additionalProperties",
            "patternProperties",
            "required",
            "not",
            "if",
            "then",
            "else",
        ]
        .iter()
        .any(|key| object.contains_key(*key))
}

#[cfg(test)]
#[path = "../test/protocols/gemini.rs"]
mod tests;
