use serde::{Deserialize, Deserializer, Serialize};
use serde_json::Value;

#[derive(Debug, Clone, Deserialize)]
pub struct ChatCompletionRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,
    #[serde(default = "default_max_tokens", alias = "max_completion_tokens")]
    pub max_tokens: u32,
    #[serde(default = "default_temperature")]
    pub temperature: f32,
    #[serde(default = "default_top_p")]
    pub top_p: f32,
    #[serde(default)]
    pub stream: bool,
    #[serde(default, deserialize_with = "deserialize_stop_sequences")]
    pub stop: Vec<String>,
    #[serde(default = "default_n")]
    pub n: u32,
    #[serde(default)]
    pub seed: Option<u32>,
    #[serde(default)]
    pub user: Option<String>,
    #[serde(default)]
    pub response_format: Option<ResponseFormat>,
    #[serde(default)]
    pub tools: Vec<ToolDefinition>,
    #[serde(default)]
    pub tool_choice: Option<ToolChoice>,
    #[serde(default)]
    pub parallel_tool_calls: Option<bool>,
    #[serde(default)]
    pub stream_options: Option<StreamOptions>,
}

fn default_max_tokens() -> u32 {
    1024
}

fn default_temperature() -> f32 {
    0.15
}

fn default_top_p() -> f32 {
    0.95
}

fn default_n() -> u32 {
    1
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: Role,
    #[serde(
        default,
        deserialize_with = "deserialize_message_content",
        skip_serializing_if = "Option::is_none"
    )]
    pub content: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl ChatMessage {
    pub fn text(&self) -> &str {
        self.content.as_deref().unwrap_or("")
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    Developer,
    User,
    Assistant,
    Tool,
}

#[derive(Debug, Clone, Serialize)]
pub struct ChatCompletionResponse {
    pub id: String,
    pub object: &'static str,
    pub created: i64,
    pub model: String,
    pub choices: Vec<Choice>,
    pub usage: Usage,
}

#[derive(Debug, Clone, Serialize)]
pub struct Choice {
    pub index: u32,
    pub message: ChatMessage,
    pub finish_reason: FinishReason,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
    Stop,
    Length,
    ToolCalls,
}

#[derive(Debug, Clone, Serialize)]
pub struct ChatCompletionChunk {
    pub id: String,
    pub object: &'static str,
    pub created: i64,
    pub model: String,
    pub choices: Vec<ChunkChoice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ChunkChoice {
    pub index: u32,
    pub delta: Delta,
    pub finish_reason: Option<FinishReason>,
}

#[derive(Debug, Clone, Serialize)]
pub struct Delta {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub role: Option<Role>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<DeltaToolCall>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    #[serde(rename = "type")]
    pub tool_type: ToolType,
    pub function: FunctionDefinition,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ToolType {
    Function,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDefinition {
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(default)]
    pub parameters: Value,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    #[serde(rename = "type")]
    pub tool_type: ToolType,
    pub function: ToolFunctionCall,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunctionCall {
    pub name: String,
    pub arguments: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct DeltaToolCall {
    pub index: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
    pub tool_type: Option<ToolType>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function: Option<DeltaToolFunctionCall>,
}

#[derive(Debug, Clone, Serialize)]
pub struct DeltaToolFunctionCall {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub arguments: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum ToolChoice {
    Mode(String),
    Specific(ToolChoiceObject),
}

#[derive(Debug, Clone, Deserialize)]
pub struct ToolChoiceObject {
    #[serde(rename = "type")]
    pub choice_type: ToolType,
    pub function: ToolChoiceFunction,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ToolChoiceFunction {
    pub name: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ResponseFormat {
    #[serde(rename = "type")]
    pub format_type: String,
    #[serde(default)]
    pub json_schema: Option<Value>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct StreamOptions {
    #[serde(default)]
    pub include_usage: bool,
}

#[derive(Debug, Clone, Serialize, Default)]
pub struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
    pub tokens_per_second: f64,
}

fn deserialize_stop_sequences<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
where
    D: Deserializer<'de>,
{
    let value = Option::<Value>::deserialize(deserializer)?;
    let Some(value) = value else {
        return Ok(Vec::new());
    };

    match value {
        Value::Null => Ok(Vec::new()),
        Value::String(s) => Ok(vec![s]),
        Value::Array(items) => items
            .into_iter()
            .map(|item| match item {
                Value::String(s) => Ok(s),
                _ => Err(serde::de::Error::custom("stop items must be strings")),
            })
            .collect(),
        _ => Err(serde::de::Error::custom(
            "stop must be a string, array of strings, or null",
        )),
    }
}

fn deserialize_message_content<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
    D: Deserializer<'de>,
{
    let value = Option::<Value>::deserialize(deserializer)?;
    let Some(value) = value else {
        return Ok(None);
    };

    match value {
        Value::Null => Ok(None),
        Value::String(s) => Ok(Some(s)),
        Value::Array(items) => {
            let mut text_parts = Vec::new();

            for item in items {
                match item {
                    Value::String(s) => text_parts.push(s),
                    Value::Object(map) => {
                        let item_type = map.get("type").and_then(Value::as_str).unwrap_or("text");
                        if matches!(item_type, "text" | "input_text" | "output_text") {
                            if let Some(text) = map.get("text").and_then(Value::as_str) {
                                text_parts.push(text.to_owned());
                            } else if let Some(text) = map.get("content").and_then(Value::as_str) {
                                text_parts.push(text.to_owned());
                            }
                        }
                    }
                    _ => {
                        return Err(serde::de::Error::custom(
                            "message content array items must be strings or text parts",
                        ));
                    }
                }
            }

            if text_parts.is_empty() {
                Ok(None)
            } else {
                Ok(Some(text_parts.join("\n")))
            }
        }
        _ => Err(serde::de::Error::custom(
            "message content must be a string, array, or null",
        )),
    }
}

#[cfg(test)]
mod tests {
    use super::{ChatCompletionRequest, Role, ToolChoice};

    #[test]
    fn deserializes_openai_style_message_parts_and_aliases() {
        let request: ChatCompletionRequest = serde_json::from_str(
            r#"{
                "model": "alias-model",
                "messages": [
                    {
                        "role": "developer",
                        "content": [
                            { "type": "text", "text": "Write code" },
                            { "type": "input_text", "text": "Use Rust" }
                        ]
                    },
                    {
                        "role": "user",
                        "content": "Build a parser"
                    }
                ],
                "max_completion_tokens": 77,
                "stop": "</tool_call>",
                "tool_choice": "auto"
            }"#,
        )
        .unwrap();

        assert_eq!(request.max_tokens, 77);
        assert_eq!(request.stop, vec!["</tool_call>"]);
        assert!(matches!(request.messages[0].role, Role::Developer));
        assert_eq!(request.messages[0].text(), "Write code\nUse Rust");
        assert!(matches!(request.tool_choice, Some(ToolChoice::Mode(ref mode)) if mode == "auto"));
    }

    #[test]
    fn deserializes_stop_arrays_and_null_content() {
        let request: ChatCompletionRequest = serde_json::from_str(
            r#"{
                "model": "alias-model",
                "messages": [
                    { "role": "assistant", "content": null },
                    { "role": "tool", "tool_call_id": "call_1", "content": "done" }
                ],
                "stop": ["END", "STOP"]
            }"#,
        )
        .unwrap();

        assert_eq!(request.stop, vec!["END", "STOP"]);
        assert_eq!(request.messages[0].content, None);
        assert_eq!(request.messages[1].tool_call_id.as_deref(), Some("call_1"));
    }
}