llmshim 0.1.26

use crate::error::{Result, ShimError};
use crate::provider::{Provider, ProviderRequest};
use crate::vision;
use serde_json::{json, Value};

pub struct Xai {
    pub api_key: String,
    pub base_url: String,
}

impl Xai {
    pub fn new(api_key: String) -> Self {
        Self {
            api_key,
            base_url: "https://api.x.ai/v1".to_string(),
        }
    }

    pub fn with_base_url(mut self, url: String) -> Self {
        self.base_url = url;
        self
    }
}

/// Sanitize messages for xAI Responses API: strip cross-provider fields,
/// translate images, and convert tool call/result messages to Responses API format.
///
/// The Responses API expects:
/// - Assistant messages with tool_calls → split into the assistant message +
///   separate `function_call` items
/// - `role: "tool"` messages → `function_call_output` items
fn sanitize_messages(messages: &[Value]) -> Vec<Value> {
    let mut result = Vec::new();
    for msg in messages {
        let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");

        match role {
            "assistant" => {
                let mut out = msg.clone();
                if let Some(obj) = out.as_object_mut() {
                    obj.remove("reasoning_content");
                    obj.remove("annotations");
                    obj.remove("refusal");
                    obj.remove("tool_calls");
                }
                if let Some(content) = out.get("content").cloned() {
                    if content.is_array() {
                        let translated =
                            vision::translate_content_blocks(&content, vision::to_openai);
                        out["content"] = vision::text_blocks_to_openai(&translated);
                    }
                }
                let has_content = out
                    .get("content")
                    .map(|c| !c.is_null() && c.as_str().map(|s| !s.is_empty()).unwrap_or(true))
                    .unwrap_or(false);
                if has_content {
                    result.push(out);
                }

                // Emit function_call items for each tool call
                if let Some(tool_calls) = msg.get("tool_calls").and_then(|tc| tc.as_array()) {
                    for tc in tool_calls {
                        let call_id = tc
                            .get("id")
                            .and_then(|v| v.as_str())
                            .unwrap_or("")
                            .to_string();
                        let name = tc
                            .get("function")
                            .and_then(|f| f.get("name"))
                            .and_then(|n| n.as_str())
                            .unwrap_or("")
                            .to_string();
                        let arguments = tc
                            .get("function")
                            .and_then(|f| f.get("arguments"))
                            .and_then(|a| a.as_str())
                            .unwrap_or("{}")
                            .to_string();
                        result.push(json!({
                            "type": "function_call",
                            "call_id": call_id,
                            "name": name,
                            "arguments": arguments,
                        }));
                    }
                }
            }
            "tool" => {
                let call_id = msg
                    .get("tool_call_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let output = msg
                    .get("content")
                    .and_then(|c| c.as_str())
                    .unwrap_or("")
                    .to_string();
                result.push(json!({
                    "type": "function_call_output",
                    "call_id": call_id,
                    "output": output,
                }));
            }
            _ => {
                let mut out = msg.clone();
                if let Some(obj) = out.as_object_mut() {
                    obj.remove("reasoning_content");
                    obj.remove("annotations");
                    obj.remove("refusal");
                }
                if let Some(content) = out.get("content").cloned() {
                    if content.is_array() {
                        let translated =
                            vision::translate_content_blocks(&content, vision::to_openai);
                        out["content"] = vision::text_blocks_to_openai(&translated);
                    }
                }
                result.push(out);
            }
        }
    }
    result
}

/// Translate tools from Chat Completions nested format to Responses API flat format.
/// Chat Completions: {"type": "function", "function": {"name": ..., "parameters": ...}}
/// Responses API:    {"type": "function", "name": ..., "parameters": ...}
fn translate_tools(tools: &Value) -> Value {
    if let Some(arr) = tools.as_array() {
        let translated: Vec<Value> = arr
            .iter()
            .map(|tool| {
                if let Some(func) = tool.get("function") {
                    let mut flat = json!({"type": "function"});
                    if let Some(obj) = func.as_object() {
                        for (k, v) in obj {
                            flat[k] = v.clone();
                        }
                    }
                    if let Some(obj) = tool.as_object() {
                        for (k, v) in obj {
                            if k != "type" && k != "function" {
                                flat[k] = v.clone();
                            }
                        }
                    }
                    flat
                } else {
                    tool.clone()
                }
            })
            .collect();
        json!(translated)
    } else {
        tools.clone()
    }
}

/// Translate tool_choice from Anthropic format if needed.
fn translate_tool_choice(tc: &Value) -> Value {
    if let Some(tc_obj) = tc.as_object() {
        if let Some(tc_type) = tc_obj.get("type").and_then(|t| t.as_str()) {
            return match tc_type {
                "auto" => json!("auto"),
                "any" => json!("required"),
                "none" => json!("none"),
                "tool" => tc_obj
                    .get("name")
                    .map(|name| json!({"type": "function", "function": {"name": name}}))
                    .unwrap_or_else(|| tc.clone()),
                _ => tc.clone(),
            };
        }
    }
    tc.clone()
}

/// grok-4.20-* models encode reasoning on/off in the MODEL NAME and reject
/// any reasoning parameter with a 400 ("does not support parameter
/// reasoningEffort", verified live) — the param must be omitted entirely.
fn is_reasoning_name_locked(model: &str) -> bool {
    model.to_lowercase().contains("4.20")
}

/// grok-4.5 cannot disable reasoning: `effort: "none"` -> 400 ("does not
/// support `reasoning_effort` value `none`", verified live). Unified effort
/// "none" clamps to "low". grok-4.3 and the fast pair DO accept "none".
fn reasoning_cannot_disable(model: &str) -> bool {
    model.to_lowercase().contains("4.5")
}

impl Provider for Xai {
    fn name(&self) -> &str {
        "xai"
    }

    fn transform_request(&self, model: &str, request: &Value) -> Result<ProviderRequest> {
        let obj = request.as_object().ok_or(ShimError::MissingModel)?;

        let messages = obj
            .get("messages")
            .and_then(|m| m.as_array())
            .ok_or(ShimError::MissingModel)?;

        let clean_messages = sanitize_messages(messages);

        // Use Responses API (same as OpenAI)
        let mut body = json!({
            "model": model,
            "input": clean_messages,
        });
        let body_obj = body.as_object_mut().unwrap();

        // max_output_tokens
        if let Some(v) = obj.get("max_tokens").or(obj.get("max_completion_tokens")) {
            body_obj.insert("max_output_tokens".to_string(), v.clone());
        }

        // Stream flag
        if let Some(v) = obj.get("stream") {
            body_obj.insert("stream".to_string(), v.clone());
        }

        // Tools — translate Chat Completions format to Responses API flat format
        if let Some(tools) = obj.get("tools") {
            body_obj.insert("tools".to_string(), translate_tools(tools));
        }

        // tool_choice
        if let Some(tc) = obj.get("tool_choice") {
            body_obj.insert("tool_choice".to_string(), translate_tool_choice(tc));
        }

        // Extract system/developer messages into instructions
        let input = body_obj.get_mut("input").unwrap().as_array_mut().unwrap();
        let mut instructions: Vec<String> = Vec::new();
        input.retain(|msg| match msg.get("role").and_then(|r| r.as_str()) {
            Some("system" | "developer") => {
                if let Some(text) = msg.get("content").and_then(|c| c.as_str()) {
                    instructions.push(text.to_string());
                }
                false
            }
            _ => true,
        });
        if !instructions.is_empty() {
            body_obj.insert("instructions".to_string(), json!(instructions.join("\n\n")));
        }

        // Unified reasoning controls -> xAI's nested `reasoning: {effort}`.
        // The flat `reasoning_effort` string is silently IGNORED by xAI's API
        // (verified live) — only the nested object works. grok-4.20-* models
        // are name-locked (reasoning baked into the model name) and 400 on any
        // reasoning param, so omit it entirely for them.
        if !is_reasoning_name_locked(model) {
            if let Some(effort) = obj.get("reasoning_effort").and_then(|e| e.as_str()) {
                // mode:"pro" has no xAI analog; map to a one-tier effort bump.
                let pro = obj
                    .get("reasoning_mode")
                    .and_then(|m| m.as_str())
                    .map(|m| m == "pro")
                    .unwrap_or(false);
                let effort = match (effort, pro) {
                    // Explicit off wins — except on models that can't disable
                    // reasoning (grok-4.5), where it clamps to the floor.
                    ("none", _) if reasoning_cannot_disable(model) => "low",
                    ("none", _) => "none",
                    ("minimal" | "low", false) => "low",
                    ("minimal" | "low", true) => "medium",
                    ("medium", false) => "medium",
                    ("medium", true) => "high",
                    ("high", false) => "high",
                    // xAI's ceiling is xhigh ("max" rejected, verified live 400)
                    ("high", true) | ("xhigh", _) | ("max", _) => "xhigh",
                    _ => "low", // unknown -> xAI's default
                };
                body_obj.insert("reasoning".to_string(), json!({ "effort": effort }));
            }
        }

        // Strip provider-specific params that don't apply to xAI
        body_obj.remove("thinking");
        body_obj.remove("output_config");
        body_obj.remove("reasoning_effort");

        let url = format!("{}/responses", self.base_url);

        Ok(ProviderRequest {
            url,
            headers: vec![
                ("Authorization".into(), format!("Bearer {}", self.api_key)),
                ("Content-Type".into(), "application/json".into()),
            ],
            body,
        })
    }

    fn transform_response(&self, model: &str, response: Value) -> Result<Value> {
        // Check for error (Responses API returns "error": null on success)
        if let Some(err) = response.get("error") {
            if !err.is_null() {
                let msg = err
                    .get("message")
                    .and_then(|m| m.as_str())
                    .unwrap_or("unknown error");
                return Err(ShimError::ProviderError {
                    status: 400,
                    body: msg.to_string(),
                });
            }
        }

        let output = response
            .get("output")
            .and_then(|o| o.as_array())
            .ok_or_else(|| ShimError::ProviderError {
                status: 500,
                body: "no output in response".to_string(),
            })?;

        let mut text_content: Option<String> = None;
        let mut reasoning_content: Option<String> = None;
        let mut tool_calls: Vec<Value> = Vec::new();

        for item in output {
            match item.get("type").and_then(|t| t.as_str()) {
                Some("message") => {
                    if let Some(content) = item.get("content").and_then(|c| c.as_array()) {
                        for part in content {
                            if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
                                text_content = Some(text.to_string());
                            }
                        }
                    }
                }
                Some("function_call") => {
                    tool_calls.push(json!({
                        "id": item.get("call_id").cloned().unwrap_or(json!("")),
                        "type": "function",
                        "function": {
                            "name": item.get("name").cloned().unwrap_or(json!("")),
                            "arguments": item.get("arguments").and_then(|a| a.as_str()).unwrap_or("{}"),
                        }
                    }));
                }
                Some("reasoning") => {
                    // grok-4.3 returns a real summary here (older models emit
                    // an empty item); surface it like the OpenAI provider does.
                    if let Some(summary) = item.get("summary").and_then(|s| s.as_array()) {
                        let texts: Vec<&str> = summary
                            .iter()
                            .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
                            .collect();
                        if !texts.is_empty() {
                            reasoning_content = Some(texts.join("\n"));
                        }
                    }
                }
                _ => {}
            }
        }

        let content = text_content.map(|t| json!(t)).unwrap_or(Value::Null);
        let mut message = json!({"role": "assistant", "content": content});
        if !tool_calls.is_empty() {
            message["tool_calls"] = json!(tool_calls);
        }
        if let Some(reasoning) = reasoning_content {
            message["reasoning_content"] = json!(reasoning);
        }

        let status = response
            .get("status")
            .and_then(|s| s.as_str())
            .unwrap_or("completed");
        let finish_reason = match status {
            "completed" => "stop",
            "incomplete" => "length",
            _ => "stop",
        };

        let usage = response.get("usage").cloned().unwrap_or(json!({}));
        let reasoning_tokens = usage
            .pointer("/output_tokens_details/reasoning_tokens")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);

        let mut result = json!({
            "id": response.get("id").cloned().unwrap_or(json!("")),
            "object": "chat.completion",
            "model": model,
            "choices": [{
                "index": 0,
                "message": message,
                "finish_reason": finish_reason,
            }],
            "usage": {
                "prompt_tokens": usage.get("input_tokens").cloned().unwrap_or(json!(0)),
                "completion_tokens": usage.get("output_tokens").cloned().unwrap_or(json!(0)),
                "total_tokens": usage.get("total_tokens").cloned().unwrap_or(json!(0)),
            }
        });

        // Surface reasoning token count so users know reasoning happened
        if reasoning_tokens > 0 {
            result["usage"]["reasoning_tokens"] = json!(reasoning_tokens);
        }

        Ok(result)
    }

    fn transform_stream_chunk(&self, model: &str, chunk: &str) -> Result<Option<String>> {
        let trimmed = chunk.trim();
        if trimmed.is_empty() || trimmed == "[DONE]" {
            return Ok(None);
        }

        let parsed: Value = serde_json::from_str(trimmed)?;
        let event_type = parsed.get("type").and_then(|t| t.as_str()).unwrap_or("");

        match event_type {
            // Content text deltas
            "response.output_text.delta" => {
                let delta = parsed.get("delta").and_then(|d| d.as_str()).unwrap_or("");
                if delta.is_empty() {
                    return Ok(None);
                }
                let chunk = json!({
                    "object": "chat.completion.chunk",
                    "model": model,
                    "choices": [{
                        "index": 0,
                        "delta": {"content": delta},
                        "finish_reason": null,
                    }]
                });
                Ok(Some(serde_json::to_string(&chunk)?))
            }

            // Function call output item added — emit tool_calls start chunk
            "response.output_item.added" => {
                let empty = json!({});
                let item = parsed.get("item").unwrap_or(&empty);
                if item.get("type").and_then(|t| t.as_str()) != Some("function_call") {
                    return Ok(None);
                }
                let call_id = item.get("call_id").and_then(|v| v.as_str()).unwrap_or("");
                let name = item.get("name").and_then(|v| v.as_str()).unwrap_or("");
                let index = parsed
                    .get("output_index")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0);
                let chunk = json!({
                    "object": "chat.completion.chunk",
                    "model": model,
                    "choices": [{
                        "index": 0,
                        "delta": {
                            "tool_calls": [{
                                "index": index,
                                "id": call_id,
                                "type": "function",
                                "function": {"name": name, "arguments": ""},
                            }]
                        },
                        "finish_reason": null,
                    }]
                });
                Ok(Some(serde_json::to_string(&chunk)?))
            }

            // Function call argument deltas
            "response.function_call_arguments.delta" => {
                let delta = parsed.get("delta").and_then(|d| d.as_str()).unwrap_or("");
                if delta.is_empty() {
                    return Ok(None);
                }
                let index = parsed
                    .get("output_index")
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0);
                let chunk = json!({
                    "object": "chat.completion.chunk",
                    "model": model,
                    "choices": [{
                        "index": 0,
                        "delta": {
                            "tool_calls": [{
                                "index": index,
                                "function": {"arguments": delta},
                            }]
                        },
                        "finish_reason": null,
                    }]
                });
                Ok(Some(serde_json::to_string(&chunk)?))
            }

            // Response completed
            "response.completed" => {
                let resp = &parsed["response"];
                let status = resp
                    .get("status")
                    .and_then(|s| s.as_str())
                    .unwrap_or("completed");
                let finish_reason = match status {
                    "completed" => "stop",
                    "incomplete" => "length",
                    _ => "stop",
                };
                let usage = resp.get("usage").cloned().unwrap_or(json!({}));
                let reasoning_tokens = usage
                    .pointer("/output_tokens_details/reasoning_tokens")
                    .cloned()
                    .unwrap_or(json!(0));
                let chunk = json!({
                    "object": "chat.completion.chunk",
                    "model": model,
                    "choices": [{
                        "index": 0,
                        "delta": {},
                        "finish_reason": finish_reason,
                    }],
                    "usage": {
                        "prompt_tokens": usage.get("input_tokens").cloned().unwrap_or(json!(0)),
                        "completion_tokens": usage.get("output_tokens").cloned().unwrap_or(json!(0)),
                        "reasoning_tokens": reasoning_tokens,
                    }
                });
                Ok(Some(serde_json::to_string(&chunk)?))
            }

            // All other events: skip
            _ => Ok(None),
        }
    }
}