//! DeepSeek provider adapter.
//!
//! DeepSeek exposes an OpenAI-compatible chat completions API, so this
//! provider builds standard OpenAI-format requests and parses the
//! corresponding responses.

use async_trait::async_trait;
use futures::StreamExt;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tracing::debug;

use athen_core::error::{AthenError, Result};
use athen_core::llm::*;
use athen_core::traits::llm::LlmProvider;

use crate::providers::openai::{
    parse_sse_chunks, parse_tool_arguments, take_complete_lines, ToolCallAccumulator,
};
use crate::quirks::{self, seed, ModelQuirks};

const DEFAULT_BASE_URL: &str = "https://api.deepseek.com";
const DEFAULT_MODEL: &str = "deepseek-chat";

/// DeepSeek LLM provider.
pub struct DeepSeekProvider {
    api_key: String,
    default_model: String,
    client: Client,
    base_url: String,
    /// Resolved at construction from the user-selected `ModelFamily`.
    /// Defaults to the `DeepSeekV4Chat` quirks (control-char repair, no
    /// reasoning); switch to `DeepSeekR1` to enable reasoning_content
    /// promotion + echo-on-tool-turn.
    quirks: ModelQuirks,
}

impl DeepSeekProvider {
    /// Create a new DeepSeek provider with the default model (`deepseek-chat`).
    pub fn new(api_key: String) -> Self {
        Self {
            api_key,
            default_model: DEFAULT_MODEL.to_string(),
            client: Client::new(),
            base_url: DEFAULT_BASE_URL.to_string(),
            // DeepSeek-V4 chat is the typical default for this provider.
            // Callers can override via `with_family(ModelFamily::DeepSeekR1)`
            // if they're pointed at the reasoner endpoint.
            quirks: seed::quirks_for_family(ModelFamily::DeepSeekV4Chat),
        }
    }

    /// Set the model family. The default constructor seeds
    /// `DeepSeekV4Chat`; switch to `DeepSeekR1` when pointing at the
    /// reasoner endpoint so reasoning_content is promoted and echoed.
    pub fn with_family(mut self, family: ModelFamily) -> Self {
        self.quirks = seed::quirks_for_family(family);
        self
    }

    /// Override the base URL (useful for testing or proxies).
    pub fn with_base_url(mut self, url: String) -> Self {
        self.base_url = url;
        self
    }

    /// Override the default model.
    pub fn with_model(mut self, model: String) -> Self {
        self.default_model = model;
        self
    }

    /// Override the HTTP client.
    pub fn with_client(mut self, client: Client) -> Self {
        self.client = client;
        self
    }

    /// Build the OpenAI-compatible request body from our generic request.
    fn build_request_body(&self, request: &LlmRequest) -> OpenAiRequestOut {
        let mut messages: Vec<OpenAiMessageOut> = Vec::new();

        // Prepend system prompt as a system message if provided.
        if let Some(ref system) = request.system_prompt {
            messages.push(OpenAiMessageOut {
                role: "system".to_string(),
                content: Some(serde_json::Value::String(system.clone())),
                tool_call_id: None,
                tool_calls: None,
            });
        }

        for m in &request.messages {
            let role = match m.role {
                Role::System => "system",
                Role::User => "user",
                Role::Assistant => "assistant",
                Role::Tool => "tool",
            };

            match (&m.role, &m.content) {
                // Assistant messages with embedded tool_calls metadata.
                (Role::Assistant, MessageContent::Structured(v))
                    if v.get("tool_calls").is_some() =>
                {
                    let text = v.get("text").and_then(|t| t.as_str()).unwrap_or_default();
                    let content = if text.is_empty() {
                        None
                    } else {
                        Some(serde_json::Value::String(text.to_string()))
                    };

                    // Convert our ToolCall structs into OpenAI wire format.
                    let tool_calls: Option<Vec<OpenAiToolCallOut>> = v
                        .get("tool_calls")
                        .and_then(|tc| serde_json::from_value::<Vec<ToolCallWire>>(tc.clone()).ok())
                        .map(|calls| {
                            calls
                                .into_iter()
                                .map(|tc| OpenAiToolCallOut {
                                    id: tc.id,
                                    call_type: "function".to_string(),
                                    function: OpenAiToolCallFunctionOut {
                                        name: tc.name,
                                        arguments: if tc.arguments.is_string() {
                                            tc.arguments.as_str().unwrap_or_default().to_string()
                                        } else {
                                            serde_json::to_string(&tc.arguments).unwrap_or_default()
                                        },
                                    },
                                })
                                .collect()
                        });

                    messages.push(OpenAiMessageOut {
                        role: role.to_string(),
                        content,
                        tool_call_id: None,
                        tool_calls,
                    });
                }
                // Tool result messages with tool_call_id.
                (Role::Tool, MessageContent::Structured(v)) if v.get("tool_call_id").is_some() => {
                    let tool_call_id = v
                        .get("tool_call_id")
                        .and_then(|id| id.as_str())
                        .unwrap_or_default()
                        .to_string();
                    let content_str = v.get("content").and_then(|c| c.as_str()).unwrap_or("{}");

                    messages.push(OpenAiMessageOut {
                        role: role.to_string(),
                        content: Some(serde_json::Value::String(content_str.to_string())),
                        tool_call_id: Some(tool_call_id),
                        tool_calls: None,
                    });
                }
                // All other messages: plain text or structured. Multimodal
                // is rejected upstream by `reject_multimodal` before we ever
                // get here; we degrade to the text-only fallback if it
                // somehow slips through, never silently dropping images
                // (the upstream rejection is the contract).
                (_, content) => {
                    let content_value = match content {
                        MessageContent::Text(t) => serde_json::Value::String(t.clone()),
                        MessageContent::Structured(v) => v.clone(),
                        MessageContent::Multimodal { text, .. } => {
                            serde_json::Value::String(text.clone())
                        }
                    };

                    messages.push(OpenAiMessageOut {
                        role: role.to_string(),
                        content: Some(content_value),
                        tool_call_id: None,
                        tool_calls: None,
                    });
                }
            }
        }

        // Map tools to OpenAI function-calling format if present.
        let tools = request.tools.as_ref().map(|tool_defs| {
            tool_defs
                .iter()
                .map(|td| OpenAiTool {
                    tool_type: "function".to_string(),
                    function: OpenAiFunction {
                        name: td.name.clone(),
                        description: td.description.clone(),
                        parameters: td.parameters.clone(),
                    },
                })
                .collect()
        });

        OpenAiRequestOut {
            model: self.default_model.clone(),
            messages,
            max_tokens: request.max_tokens.unwrap_or(4096),
            temperature: request.temperature.unwrap_or(0.7),
            tools,
            stream: false,
            reasoning_effort: map_deepseek_reasoning_effort(request.reasoning_effort),
        }
    }

    /// Map HTTP error responses to `AthenError`.
    fn map_error(&self, status: reqwest::StatusCode, body: &str) -> AthenError {
        let message = if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
            format!("rate_limit: {}", body)
        } else if status == reqwest::StatusCode::UNAUTHORIZED {
            format!("auth error: {}", body)
        } else {
            format!("HTTP {}: {}", status, body)
        };

        AthenError::LlmProvider {
            provider: "deepseek".into(),
            message,
        }
    }
}

#[async_trait]
impl LlmProvider for DeepSeekProvider {
    fn provider_id(&self) -> &str {
        "deepseek"
    }

    async fn complete(&self, request: &LlmRequest) -> Result<LlmResponse> {
        crate::providers::reject_multimodal("deepseek", request)?;
        let body = self.build_request_body(request);
        let url = format!("{}/v1/chat/completions", self.base_url);

        debug!(model = %body.model, "sending DeepSeek completion request");

        let http_response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| AthenError::LlmProvider {
                provider: "deepseek".into(),
                message: format!("request failed: {}", e),
            })?;

        let status = http_response.status();
        if !status.is_success() {
            let error_body = http_response.text().await.unwrap_or_default();
            return Err(self.map_error(status, &error_body));
        }

        let api_response: OpenAiResponse =
            http_response
                .json()
                .await
                .map_err(|e| AthenError::LlmProvider {
                    provider: "deepseek".into(),
                    message: format!("failed to parse response: {}", e),
                })?;

        let choice = api_response
            .choices
            .first()
            .ok_or_else(|| AthenError::LlmProvider {
                provider: "deepseek".into(),
                message: "response contained no choices".into(),
            })?;

        let content = choice.message.content.clone().unwrap_or_default();
        let reasoning_content = choice.message.reasoning_content.clone();

        // Extract tool calls if present.
        let tool_calls: Vec<ToolCall> = choice
            .message
            .tool_calls
            .as_ref()
            .map(|calls| {
                calls
                    .iter()
                    .map(|tc| ToolCall {
                        id: tc.id.clone(),
                        name: tc.function.name.clone(),
                        arguments: parse_tool_arguments(&tc.function.arguments),
                        thought_signature: None,
                    })
                    .collect()
            })
            .unwrap_or_default();

        let finish_reason = match choice.finish_reason.as_deref() {
            Some("stop") => FinishReason::Stop,
            Some("tool_calls") => FinishReason::ToolUse,
            Some("length") => FinishReason::MaxTokens,
            _ => FinishReason::Stop,
        };

        let usage = if let Some(ref u) = api_response.usage {
            TokenUsage {
                prompt_tokens: u.prompt_tokens,
                completion_tokens: u.completion_tokens,
                total_tokens: u.total_tokens,
                estimated_cost_usd: Some(estimate_deepseek_cost(
                    &api_response.model,
                    u.prompt_tokens,
                    u.completion_tokens,
                )),
            }
        } else {
            TokenUsage {
                prompt_tokens: 0,
                completion_tokens: 0,
                total_tokens: 0,
                estimated_cost_usd: None,
            }
        };

        let mut response = LlmResponse {
            content,
            reasoning_content,
            model_used: api_response.model,
            provider: "deepseek".into(),
            usage,
            tool_calls,
            finish_reason,
        };
        quirks::apply_to_response(&self.quirks, &mut response);
        Ok(response)
    }

    async fn complete_streaming(&self, request: &LlmRequest) -> Result<LlmStream> {
        crate::providers::reject_multimodal("deepseek", request)?;
        let mut body = self.build_request_body(request);
        body.stream = true;
        let url = format!("{}/v1/chat/completions", self.base_url);

        debug!(model = %body.model, "sending DeepSeek streaming request");

        let http_response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| AthenError::LlmProvider {
                provider: "deepseek".into(),
                message: format!("streaming request failed: {}", e),
            })?;

        let status = http_response.status();
        if !status.is_success() {
            let error_body = http_response.text().await.unwrap_or_default();
            return Err(self.map_error(status, &error_body));
        }

        let byte_stream = http_response.bytes_stream();
        let quirks_snapshot = self.quirks;

        // DeepSeek is OpenAI-compatible — reuse the stateful SSE parser so
        // fragmented tool-call deltas are correctly assembled across chunks,
        // and buffer raw bytes across HTTP chunks so SSE lines and multi-byte
        // UTF-8 codepoints split at chunk boundaries reassemble correctly.
        // The `content_buffer` + `saw_structured` fields drive the
        // end-of-stream inline tool-call extraction (a no-op for the
        // baseline `Structured` strategy DeepSeek's chat-class models use).
        let chunk_stream = futures::stream::unfold(
            DeepSeekStreamState {
                byte_stream,
                acc: ToolCallAccumulator::default(),
                pending: Vec::<u8>::new(),
                done: false,
                content_buffer: String::new(),
                saw_structured: false,
                quirks: quirks_snapshot,
            },
            |mut s| async move {
                if s.done {
                    return None;
                }
                let parsed = match s.byte_stream.next().await {
                    Some(Ok(bytes)) => {
                        s.pending.extend_from_slice(&bytes);
                        let complete = take_complete_lines(&mut s.pending);
                        if complete.is_empty() {
                            Vec::new()
                        } else {
                            let text = String::from_utf8_lossy(&complete);
                            let chunks = parse_sse_chunks(&text, "deepseek", &mut s.acc);
                            observe_deepseek_chunks(&chunks, &mut s);
                            chunks
                        }
                    }
                    Some(Err(e)) => vec![Err(AthenError::LlmProvider {
                        provider: "deepseek".into(),
                        message: format!("stream error: {}", e),
                    })],
                    None => {
                        s.done = true;
                        let mut out = Vec::new();
                        if !s.pending.is_empty() {
                            let tail = std::mem::take(&mut s.pending);
                            let text = String::from_utf8_lossy(&tail);
                            let chunks = parse_sse_chunks(&text, "deepseek", &mut s.acc);
                            observe_deepseek_chunks(&chunks, &mut s);
                            out.extend(chunks);
                        }
                        if !s.acc.is_empty() {
                            let tool_calls = s.acc.drain();
                            if !tool_calls.is_empty() {
                                s.saw_structured = true;
                                out.push(Ok(LlmChunk {
                                    delta: String::new(),
                                    is_final: true,
                                    is_thinking: false,
                                    tool_calls,
                                }));
                            }
                        }
                        if let Some(tail) = quirks::extract_streaming_tail(
                            &s.quirks,
                            &s.content_buffer,
                            s.saw_structured,
                        ) {
                            out.push(Ok(tail));
                        }
                        out
                    }
                };
                Some((futures::stream::iter(parsed), s))
            },
        )
        .flatten();

        Ok(Box::pin(chunk_stream))
    }

    async fn is_available(&self) -> bool {
        true
    }
}

/// Streaming state for DeepSeek's `complete_streaming` unfold. Mirrors
/// `OpenAiCompatibleProvider::StreamState` but flattens the byte-stream
/// generic so the unfold closure stays inferable.
struct DeepSeekStreamState<S> {
    byte_stream: S,
    acc: ToolCallAccumulator,
    pending: Vec<u8>,
    done: bool,
    content_buffer: String,
    saw_structured: bool,
    quirks: ModelQuirks,
}

/// Per-batch chunk observer: appends visible deltas to the content buffer
/// (only when extraction will be needed at end-of-stream) and flips the
/// `saw_structured` flag on any chunk that surfaced a structured tool call.
fn observe_deepseek_chunks<S>(
    chunks: &[Result<athen_core::llm::LlmChunk>],
    state: &mut DeepSeekStreamState<S>,
) {
    let needs_buffer = !matches!(
        state.quirks.tool_extraction,
        crate::quirks::ToolExtractionStrategy::Structured
    );
    for c in chunks.iter().flatten() {
        if !c.tool_calls.is_empty() {
            state.saw_structured = true;
        }
        if needs_buffer && !c.delta.is_empty() && !c.is_thinking {
            state.content_buffer.push_str(&c.delta);
        }
    }
}

/// Cost estimation for DeepSeek models.
///
/// DeepSeek pricing is very competitive (as of 2025):
/// - deepseek-chat: ~$0.14/M input, ~$0.28/M output
/// - deepseek-reasoner: ~$0.55/M input, ~$2.19/M output
fn estimate_deepseek_cost(model: &str, input_tokens: u32, output_tokens: u32) -> f64 {
    let (input_per_m, output_per_m) = if model.contains("reasoner") {
        (0.55, 2.19)
    } else {
        // deepseek-chat and other models
        (0.14, 0.28)
    };

    let input_cost = (input_tokens as f64 / 1_000_000.0) * input_per_m;
    let output_cost = (output_tokens as f64 / 1_000_000.0) * output_per_m;
    input_cost + output_cost
}

// ---------------------------------------------------------------------------
// OpenAI-compatible API types (used by DeepSeek)
// ---------------------------------------------------------------------------

/// Helper for deserializing tool calls from our structured ChatMessage content.
#[derive(Debug, Deserialize)]
struct ToolCallWire {
    id: String,
    name: String,
    arguments: serde_json::Value,
}

// ---------------------------------------------------------------------------
// Outbound (serialized) API types
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize)]
struct OpenAiRequestOut {
    model: String,
    messages: Vec<OpenAiMessageOut>,
    max_tokens: u32,
    temperature: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<OpenAiTool>>,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    reasoning_effort: Option<&'static str>,
}

/// DeepSeek V4 chat exposes `reasoning_effort` with only two valid
/// values: `"high"` and `"max"`. Omission disables reasoning entirely.
/// Per the design doc:
/// - Default / Off → omit (DeepSeek defaults to reasoning off)
/// - Minimal / Low / Medium / High → `"high"` (the lowest-on value)
/// - Max → `"max"`
///
/// Models with reasoning baked in (DeepSeek-R1) silently ignore the
/// field on the wire, so this mapping is safe to apply per-request
/// without family branching.
fn map_deepseek_reasoning_effort(effort: ReasoningEffort) -> Option<&'static str> {
    match effort {
        ReasoningEffort::Default | ReasoningEffort::Off => None,
        ReasoningEffort::Minimal
        | ReasoningEffort::Low
        | ReasoningEffort::Medium
        | ReasoningEffort::High => Some("high"),
        ReasoningEffort::Max => Some("max"),
    }
}

#[derive(Debug, Serialize)]
struct OpenAiMessageOut {
    role: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    content: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<Vec<OpenAiToolCallOut>>,
}

#[derive(Debug, Serialize)]
struct OpenAiToolCallOut {
    id: String,
    #[serde(rename = "type")]
    call_type: String,
    function: OpenAiToolCallFunctionOut,
}

#[derive(Debug, Serialize)]
struct OpenAiToolCallFunctionOut {
    name: String,
    arguments: String,
}

#[derive(Debug, Serialize)]
struct OpenAiTool {
    #[serde(rename = "type")]
    tool_type: String,
    function: OpenAiFunction,
}

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

#[derive(Debug, Deserialize)]
struct OpenAiResponse {
    model: String,
    choices: Vec<OpenAiChoice>,
    usage: Option<OpenAiUsage>,
}

#[derive(Debug, Deserialize)]
struct OpenAiChoice {
    message: OpenAiResponseMessage,
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct OpenAiResponseMessage {
    content: Option<String>,
    reasoning_content: Option<String>,
    tool_calls: Option<Vec<OpenAiToolCall>>,
}

#[derive(Debug, Deserialize)]
struct OpenAiToolCall {
    id: String,
    function: OpenAiToolCallFunction,
}

#[derive(Debug, Deserialize)]
struct OpenAiToolCallFunction {
    name: String,
    arguments: String,
}

#[derive(Debug, Deserialize)]
struct OpenAiUsage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

#[cfg(test)]
mod arg_parse_tests {
    use super::*;

    #[test]
    fn well_formed_json_is_parsed_directly() {
        let v = parse_tool_arguments(r#"{"path":"/tmp/x","content":"hi"}"#);
        assert_eq!(v["path"], "/tmp/x");
        assert_eq!(v["content"], "hi");
    }

    #[test]
    fn raw_newlines_in_string_value_are_repaired() {
        // This is what DeepSeek emits: a literal LF inside the content
        // value, which serde_json::from_str rejects per the JSON spec.
        let raw = "{\"path\":\"/tmp/x.html\",\"content\":\"<html>\nhi\n</html>\"}";
        let v = parse_tool_arguments(raw);
        assert_eq!(v["path"], "/tmp/x.html");
        assert_eq!(v["content"], "<html>\nhi\n</html>");
    }

    #[test]
    fn raw_tabs_and_returns_in_string_value_are_repaired() {
        let raw = "{\"path\":\"/x\",\"content\":\"a\tb\rc\"}";
        let v = parse_tool_arguments(raw);
        assert_eq!(v["content"], "a\tb\rc");
    }

    #[test]
    fn already_escaped_sequences_are_preserved() {
        let raw = r#"{"path":"/x","content":"line1\nline2"}"#;
        let v = parse_tool_arguments(raw);
        assert_eq!(v["content"], "line1\nline2");
    }

    #[test]
    fn unrepairable_input_falls_back_to_string_wrapper() {
        let v = parse_tool_arguments("not even close to json {");
        assert!(v.is_string());
    }

    #[test]
    fn control_chars_outside_strings_are_left_alone() {
        // Newlines between fields are valid JSON whitespace. Make sure
        // we don't munge them.
        let raw = "{\n  \"path\": \"/x\",\n  \"content\": \"hi\"\n}";
        let v = parse_tool_arguments(raw);
        assert_eq!(v["path"], "/x");
    }

    #[tokio::test]
    async fn complete_rejects_multimodal_with_clear_error() {
        let provider = DeepSeekProvider::new("test-key".to_string());
        let req = LlmRequest {
            profile: ModelProfile::Powerful,
            messages: vec![ChatMessage {
                role: Role::User,
                content: MessageContent::Multimodal {
                    text: "describe this".to_string(),
                    images: vec![ImageInput {
                        mime_type: "image/png".to_string(),
                        data: ImageData::Base64 {
                            data: "AAAA".to_string(),
                        },
                    }],
                },
            }],
            max_tokens: None,
            temperature: None,
            tools: None,
            system_prompt: None,
            reasoning_effort: ReasoningEffort::default(),
        };
        let err = provider.complete(&req).await.unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.to_lowercase().contains("does not support image"),
            "expected vision rejection error, got: {msg}"
        );
    }

    fn deepseek_body_with(effort: ReasoningEffort) -> serde_json::Value {
        let provider = DeepSeekProvider::new("k".into());
        let req = LlmRequest {
            profile: ModelProfile::Powerful,
            messages: vec![ChatMessage {
                role: Role::User,
                content: MessageContent::Text("hi".into()),
            }],
            max_tokens: None,
            temperature: None,
            tools: None,
            system_prompt: None,
            reasoning_effort: effort,
        };
        serde_json::to_value(provider.build_request_body(&req)).expect("serializes")
    }

    #[test]
    fn deepseek_default_omits_reasoning_effort_field() {
        let body = deepseek_body_with(ReasoningEffort::Default);
        assert!(body.get("reasoning_effort").is_none(), "{body}");
    }

    #[test]
    fn deepseek_off_also_omits_field() {
        // DeepSeek interprets omission as "off"; sending an explicit
        // value would be wrong since `"off"` isn't a valid enum value.
        let body = deepseek_body_with(ReasoningEffort::Off);
        assert!(body.get("reasoning_effort").is_none());
    }

    #[test]
    fn deepseek_low_medium_high_all_map_to_high() {
        for eff in [
            ReasoningEffort::Minimal,
            ReasoningEffort::Low,
            ReasoningEffort::Medium,
            ReasoningEffort::High,
        ] {
            let body = deepseek_body_with(eff);
            assert_eq!(body["reasoning_effort"], "high", "for {eff:?}");
        }
    }

    #[test]
    fn deepseek_max_maps_to_max() {
        let body = deepseek_body_with(ReasoningEffort::Max);
        assert_eq!(body["reasoning_effort"], "max");
    }
}
