dynamo-protocols 2.0.1

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Re-exports upstream async-openai chat types and defines inference-serving
// extensions on top. Types prefixed with `Dynamo` or entirely absent from the
// upstream spec are documented with the rationale for the extension.

use std::pin::Pin;

use derive_builder::Builder;
use futures::Stream;
use serde::{Deserialize, Serialize};
use url::Url;
use uuid::Uuid;

use crate::error::OpenAIError;

// ---------------------------------------------------------------------------
// Re-exports from upstream async-openai (unchanged types)
// ---------------------------------------------------------------------------
// These types are structurally identical to the upstream definitions.
// Consumers should use them via `dynamo_protocols::types::*` as before.

pub use async_openai::types::chat::{
    ChatChoiceLogprobs,
    ChatCompletionAudio,
    ChatCompletionAudioFormat,
    ChatCompletionAudioVoice,
    ChatCompletionFunctionCall,
    ChatCompletionFunctions,
    ChatCompletionFunctionsArgs,
    ChatCompletionRequestAssistantMessageAudio,
    ChatCompletionRequestAssistantMessageContent,
    ChatCompletionRequestAssistantMessageContentPart,
    ChatCompletionRequestDeveloperMessage,
    ChatCompletionRequestDeveloperMessageArgs,
    ChatCompletionRequestDeveloperMessageContent,
    ChatCompletionRequestFunctionMessage,
    ChatCompletionRequestFunctionMessageArgs,
    ChatCompletionRequestMessageContentPartAudio,
    ChatCompletionRequestMessageContentPartRefusal,
    ChatCompletionRequestMessageContentPartText,
    ChatCompletionRequestSystemMessage,
    // Builder types (generated by derive_builder)
    ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestSystemMessageContent,
    ChatCompletionRequestSystemMessageContentPart,
    ChatCompletionRequestToolMessage,
    ChatCompletionRequestToolMessageArgs,
    ChatCompletionRequestToolMessageContent,
    ChatCompletionRequestToolMessageContentPart,
    ChatCompletionResponseMessageAudio,
    ChatCompletionTokenLogprob,
    Choice,
    CompletionFinishReason,
    CompletionTokensDetails,
    CompletionUsage,
    FunctionObject,
    FunctionObjectArgs,
    InputAudio,
    InputAudioFormat,
    Logprobs,
    PredictionContent,
    PredictionContentContent,
    Prompt,
    PromptTokensDetails,
    ResponseFormat,
    ResponseFormatJsonSchema,
    Role,
    ServiceTier,
    TopLogprobs,
    WebSearchContextSize,
    WebSearchLocation,
    WebSearchOptions,
    WebSearchUserLocation,
    WebSearchUserLocationType,
};

/// OpenAI stop configuration, with Dynamo's token-id stop extension.
///
/// The standard OpenAI shape accepts a string or string array. Dynamo also
/// accepts an integer array, e.g. `"stop": [576]`, to express token-id stop
/// conditions for tokenized in/out workflows. Strings like `"token_id:576"`
/// remain ordinary string stops; the `token_id:<id>` format is only an output
/// display format for logprobs.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
#[serde(untagged)]
pub enum Stop {
    String(String),
    StringArray(Vec<String>),
    TokenIdArray(Vec<u32>),
}

impl Stop {
    pub fn strings(&self) -> Option<Vec<String>> {
        match self {
            Stop::String(s) => Some(vec![s.clone()]),
            Stop::StringArray(arr) => Some(arr.clone()),
            Stop::TokenIdArray(_) => None,
        }
    }

    pub fn token_ids(&self) -> Option<Vec<u32>> {
        match self {
            Stop::TokenIdArray(arr) => Some(arr.clone()),
            Stop::String(_) | Stop::StringArray(_) => None,
        }
    }
}

impl From<String> for Stop {
    fn from(value: String) -> Self {
        Stop::String(value)
    }
}

impl From<&str> for Stop {
    fn from(value: &str) -> Self {
        Stop::String(value.to_string())
    }
}

impl From<Vec<String>> for Stop {
    fn from(value: Vec<String>) -> Self {
        Stop::StringArray(value)
    }
}

impl From<Vec<u32>> for Stop {
    fn from(value: Vec<u32>) -> Self {
        Stop::TokenIdArray(value)
    }
}

impl From<async_openai::types::chat::StopConfiguration> for Stop {
    fn from(value: async_openai::types::chat::StopConfiguration) -> Self {
        match value {
            async_openai::types::chat::StopConfiguration::String(value) => Stop::String(value),
            async_openai::types::chat::StopConfiguration::StringArray(value) => {
                Stop::StringArray(value)
            }
        }
    }
}

// Upstream renamed FinishReason (streaming) -- re-export
pub use async_openai::types::chat::FinishReason;

// Upstream uses FunctionType where we used ChatCompletionToolType.
// Re-export both names for compatibility.
pub use async_openai::types::chat::FunctionType;

/// Reasoning effort values accepted by OpenAI-compatible clients.
///
/// async-openai versions used by some Dynamo builds do not include `max`, but
/// DeepSeek-V4 compatible clients may send it by default. Keep this local enum
/// wire-compatible with upstream values and include `max`.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
    None,
    Minimal,
    Low,
    Medium,
    High,
    Xhigh,
    Max,
}

impl From<async_openai::types::chat::ReasoningEffort> for ReasoningEffort {
    fn from(value: async_openai::types::chat::ReasoningEffort) -> Self {
        match value {
            async_openai::types::chat::ReasoningEffort::None => ReasoningEffort::None,
            async_openai::types::chat::ReasoningEffort::Minimal => ReasoningEffort::Minimal,
            async_openai::types::chat::ReasoningEffort::Low => ReasoningEffort::Low,
            async_openai::types::chat::ReasoningEffort::Medium => ReasoningEffort::Medium,
            async_openai::types::chat::ReasoningEffort::High => ReasoningEffort::High,
            async_openai::types::chat::ReasoningEffort::Xhigh => ReasoningEffort::Xhigh,
        }
    }
}

// ---------------------------------------------------------------------------
// Flexible `arguments` deserialisation helpers
// ---------------------------------------------------------------------------
// Some agent frameworks (e.g. LangChain, custom harnesses) send tool-call
// arguments as a pre-parsed JSON object instead of the canonical JSON
// string.  The helpers below normalise both representations to a `String` so
// downstream code never needs to branch on the wire format.

fn deserialize_arguments<'de, D>(deserializer: D) -> Result<String, D::Error>
where
    D: serde::Deserializer<'de>,
{
    use serde::de::Error;
    let value = serde_json::Value::deserialize(deserializer)?;
    match value {
        serde_json::Value::String(s) => Ok(s),
        v @ serde_json::Value::Object(_) => {
            // serde_json::to_string on a Value is infallible
            Ok(serde_json::to_string(&v).unwrap())
        }
        other => Err(D::Error::custom(format!(
            "expected string or object for `arguments`, got {other}"
        ))),
    }
}

fn deserialize_arguments_opt<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    use serde::de::Error;
    let value = Option::<serde_json::Value>::deserialize(deserializer)?;
    match value {
        None => Ok(None),
        Some(serde_json::Value::String(s)) => Ok(Some(s)),
        Some(v @ serde_json::Value::Object(_)) => serde_json::to_string(&v)
            .map(Some)
            .map_err(|e| D::Error::custom(e.to_string())),
        Some(other) => Err(D::Error::custom(format!(
            "expected string or object for `arguments`, got {other}"
        ))),
    }
}

// ---------------------------------------------------------------------------
// FunctionCall / FunctionCallStream — local definitions with flexible deser
// ---------------------------------------------------------------------------
// Upstream `async-openai` only accepts a JSON string for `arguments`.
// We define these locally so we can attach `#[serde(deserialize_with)]` and
// accept both string and object representations on the wire.

/// The name and arguments of a function that should be called.
///
/// Accepts `arguments` as either a JSON string (`"{\"key\":\"value\"}"`) or a
/// JSON object (`{"key": "value"}`); both are normalised to a JSON string
/// on deserialisation so callers always see the canonical form.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Default)]
pub struct FunctionCall {
    pub name: String,
    #[serde(deserialize_with = "deserialize_arguments")]
    pub arguments: String,
}

/// Streaming variant of [`FunctionCall`] where both fields are optional.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Default)]
pub struct FunctionCallStream {
    pub name: Option<String>,
    #[serde(default, deserialize_with = "deserialize_arguments_opt")]
    pub arguments: Option<String>,
}

/// Streaming tool-call chunk.
///
/// Defined locally (instead of re-exporting from upstream) because its
/// `function` field references our local [`FunctionCallStream`] with the
/// flexible `arguments` deserialiser.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Default)]
pub struct ChatCompletionMessageToolCallChunk {
    pub index: u32,
    pub id: Option<String>,
    pub r#type: Option<FunctionType>,
    pub function: Option<FunctionCallStream>,
}

// ---------------------------------------------------------------------------
// Types with structural differences from upstream (kept locally)
// ---------------------------------------------------------------------------

/// Image detail level. Kept locally because upstream uses different field types in ImageUrl.
#[derive(Debug, Serialize, Deserialize, Default, Clone, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ImageDetail {
    #[default]
    Auto,
    Low,
    High,
}

/// Image content part -- uses our extended `ImageUrl` with `url::Url` and `uuid`.
#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "ChatCompletionRequestMessageContentPartImageArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionRequestMessageContentPartImage {
    pub image_url: ImageUrl,
}

/// Image URL with `url::Url` type and optional UUID.
///
/// Differs from upstream: uses `url::Url` instead of `String`, adds `uuid` field
/// for tracking multimodal assets through the pipeline.
#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "ImageUrlArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ImageUrl {
    pub url: Url,
    pub detail: Option<ImageDetail>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uuid: Option<Uuid>,
}

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

#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct FunctionName {
    pub name: String,
}

#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionNamedToolChoice {
    pub r#type: ChatCompletionToolType,
    pub function: FunctionName,
}

fn default_function_type() -> FunctionType {
    FunctionType::Function
}

/// Tool call kept locally to preserve `type: "function"` in unary request/response payloads.
///
/// Differs from upstream: `type` is serialized by default and also defaults to
/// `function` when omitted during deserialization, preserving compatibility with
/// both Dynamo's historical wire format and upstream spec-compliant inputs.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionMessageToolCall {
    pub id: String,
    #[serde(default = "default_function_type")]
    pub r#type: FunctionType,
    pub function: FunctionCall,
}

/// Tool choice enum kept locally because upstream changed variant names.
#[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ChatCompletionToolChoiceOption {
    #[default]
    None,
    Auto,
    Required,
    #[serde(untagged)]
    Named(ChatCompletionNamedToolChoice),
}

#[derive(Clone, Serialize, Default, Debug, Builder, Deserialize, PartialEq)]
#[builder(name = "ChatCompletionToolArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionTool {
    #[builder(default = "ChatCompletionToolType::Function")]
    pub r#type: ChatCompletionToolType,
    pub function: FunctionObject,
}

// ---------------------------------------------------------------------------
// Inference-serving extensions (not in upstream)
// ---------------------------------------------------------------------------

/// Matched stop condition from the backend.
///
/// Inference backends (vLLM, SGLang) report which stop condition triggered:
/// - `String`: a matched user-provided stop sequence
/// - `Int`: a matched stop token ID
/// - `IntArray`: matched stop token IDs reported as a sequence
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
#[serde(untagged)]
pub enum StopReason {
    String(String),
    Int(i64),
    IntArray(Vec<i64>),
}

/// Reasoning content from a previous assistant turn.
///
/// Deserializes from either:
/// - A plain string: `"reasoning_content": "thinking..."` -> `Text("thinking...")`
/// - An array of strings: `"reasoning_content": ["seg1", "seg2"]` -> `Segments(["seg1", "seg2"])`
///
/// The `Segments` variant preserves interleaved reasoning order needed for KV cache-correct
/// context reconstruction. `segments[i]` is the reasoning that preceded `tool_calls[i]`;
/// `segments[tool_calls.len()]` is any trailing reasoning after the last tool call.
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
#[serde(untagged)]
pub enum ReasoningContent {
    /// Flat string -- single reasoning block or legacy backward-compat form.
    Text(String),
    /// Interleaved segments. segments[i] precedes tool_calls[i];
    /// segments[N] is trailing reasoning after the last tool call.
    Segments(Vec<String>),
}

impl ReasoningContent {
    /// Join all segments (or return text as-is) into a single flat string.
    pub fn to_flat_string(&self) -> String {
        match self {
            ReasoningContent::Text(s) => s.clone(),
            ReasoningContent::Segments(segs) => segs
                .iter()
                .filter(|s| !s.is_empty())
                .cloned()
                .collect::<Vec<_>>()
                .join("\n"),
        }
    }

    /// Returns the segments if this is the `Segments` variant, `None` for `Text`.
    pub fn segments(&self) -> Option<&[String]> {
        match self {
            ReasoningContent::Segments(segs) => Some(segs),
            ReasoningContent::Text(_) => None,
        }
    }
}

// -- Multimodal content types for responses (not in upstream) --

/// Response content part for text in assistant messages
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionResponseContentPartText {
    pub text: String,
}

/// Response content part for image URLs in assistant messages
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionResponseContentPartImageUrl {
    pub image_url: ImageUrlResponse,
}

/// Response content part for video URLs in assistant messages
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionResponseContentPartVideoUrl {
    pub video_url: VideoUrlResponse,
}

/// Response content part for audio URLs in assistant messages
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ChatCompletionResponseContentPartAudioUrl {
    pub audio_url: AudioUrlResponse,
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ImageUrlResponse {
    pub url: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub detail: Option<String>,
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct VideoUrlResponse {
    pub url: String,
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct AudioUrlResponse {
    pub url: String,
}

/// Content parts for assistant responses supporting multiple modalities
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatCompletionResponseContentPart {
    Text(ChatCompletionResponseContentPartText),
    ImageUrl(ChatCompletionResponseContentPartImageUrl),
    VideoUrl(ChatCompletionResponseContentPartVideoUrl),
    AudioUrl(ChatCompletionResponseContentPartAudioUrl),
}

/// Assistant message content -- can be a simple string or multimodal content parts.
///
/// Upstream uses `Option<String>` for the content field. We extend this to
/// support multimodal responses (text + images + video + audio) from backends
/// like vLLM that can return non-text content.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ChatCompletionMessageContent {
    /// Simple text content (backward compatible)
    Text(String),
    /// Array of content parts (for multimodal responses)
    Parts(Vec<ChatCompletionResponseContentPart>),
}

// -- Multimodal input types (video/audio URL support, not in upstream) --

#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "VideoUrlArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct VideoUrl {
    pub url: Url,
    pub detail: Option<ImageDetail>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uuid: Option<Uuid>,
}

#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "ChatCompletionRequestMessageContentPartVideoArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionRequestMessageContentPartVideo {
    pub video_url: VideoUrl,
}

#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "AudioUrlArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct AudioUrl {
    pub url: Url,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uuid: Option<Uuid>,
}

#[derive(Debug, Serialize, Deserialize, Clone, Builder, PartialEq)]
#[builder(name = "ChatCompletionRequestMessageContentPartAudioUrlArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option))]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionRequestMessageContentPartAudioUrl {
    pub audio_url: AudioUrl,
}

// -- Extended request/response types --

/// User message content -- references our extended content part enum.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(untagged)]
pub enum ChatCompletionRequestUserMessageContent {
    Text(String),
    Array(Vec<ChatCompletionRequestUserMessageContentPart>),
}

#[derive(Debug, Serialize, Deserialize, Default, Clone, Builder, PartialEq)]
#[builder(name = "ChatCompletionRequestUserMessageArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionRequestUserMessage {
    pub content: ChatCompletionRequestUserMessageContent,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

impl Default for ChatCompletionRequestUserMessageContent {
    fn default() -> Self {
        Self::Text(String::new())
    }
}

impl From<&str> for ChatCompletionRequestUserMessageContent {
    fn from(value: &str) -> Self {
        Self::Text(value.into())
    }
}

impl From<String> for ChatCompletionRequestUserMessageContent {
    fn from(value: String) -> Self {
        Self::Text(value)
    }
}

impl From<Vec<ChatCompletionRequestUserMessageContentPart>>
    for ChatCompletionRequestUserMessageContent
{
    fn from(value: Vec<ChatCompletionRequestUserMessageContentPart>) -> Self {
        Self::Array(value)
    }
}

/// User message content part with video and audio URL support.
///
/// Extends upstream `ChatCompletionRequestUserMessageContentPart` with:
/// - `VideoUrl`: video input for multimodal models
/// - `AudioUrl`: audio URL input (distinct from base64 InputAudio)
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
pub enum ChatCompletionRequestUserMessageContentPart {
    Text(ChatCompletionRequestMessageContentPartText),
    ImageUrl(ChatCompletionRequestMessageContentPartImage),
    VideoUrl(ChatCompletionRequestMessageContentPartVideo),
    AudioUrl(ChatCompletionRequestMessageContentPartAudioUrl),
    InputAudio(ChatCompletionRequestMessageContentPartAudio),
}

/// Assistant message with reasoning content support.
///
/// Extends upstream `ChatCompletionRequestAssistantMessage` with:
/// - `reasoning_content`: interleaved reasoning segments for KV cache correctness
///   (DeepSeek-R1, QwQ models)
#[derive(Debug, Serialize, Deserialize, Default, Clone, Builder, PartialEq)]
#[builder(name = "ChatCompletionRequestAssistantMessageArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct ChatCompletionRequestAssistantMessage {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<ChatCompletionRequestAssistantMessageContent>,
    /// Reasoning content from a previous assistant turn.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<ReasoningContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<ChatCompletionRequestAssistantMessageAudio>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ChatCompletionMessageToolCall>>,
    #[deprecated]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_call: Option<FunctionCall>,
}

/// Chat completion request message enum.
///
/// Redefined to use our extended `ChatCompletionRequestAssistantMessage`
/// (with reasoning_content) and `ChatCompletionRequestUserMessage`
/// (which references our extended content parts with video/audio).
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(tag = "role")]
#[serde(rename_all = "lowercase")]
pub enum ChatCompletionRequestMessage {
    Developer(ChatCompletionRequestDeveloperMessage),
    System(ChatCompletionRequestSystemMessage),
    User(ChatCompletionRequestUserMessage),
    Assistant(ChatCompletionRequestAssistantMessage),
    Tool(ChatCompletionRequestToolMessage),
    Function(ChatCompletionRequestFunctionMessage),
}

/// Response tier enum for responses (distinct from request `ServiceTier`).
///
/// Not in upstream -- backends report which tier actually served the request.
#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ServiceTierResponse {
    Scale,
    Default,
    Flex,
    Priority,
}

/// Chat completion response message with multimodal content and reasoning.
///
/// Extends upstream `ChatCompletionResponseMessage` with:
/// - `content`: `Option<ChatCompletionMessageContent>` (multimodal) instead of `Option<String>`
/// - `reasoning_content`: model reasoning output (DeepSeek-R1, QwQ)
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct ChatCompletionResponseMessage {
    /// Always serialized (as `null` when None) so clients can rely on the
    /// `content` key being present alongside `reasoning_content` or
    /// `tool_calls`. Matches the upstream OpenAI API shape (DGH-651).
    pub content: Option<ChatCompletionMessageContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ChatCompletionMessageToolCall>>,
    pub role: Role,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[deprecated]
    pub function_call: Option<FunctionCall>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<ChatCompletionResponseMessageAudio>,
    /// Reasoning content produced by the model (DeepSeek-R1, QwQ).
    pub reasoning_content: Option<String>,
}

/// Stream options with per-chunk usage reporting.
///
/// Extends upstream `ChatCompletionStreamOptions` with:
/// - `continuous_usage_stats`: emit usage in every chunk, not just the final one
#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq)]
pub struct ChatCompletionStreamOptions {
    pub include_usage: bool,
    /// When true, usage statistics are included in every streaming chunk.
    /// Backends like vLLM/SGLang support this for real-time token counting.
    #[serde(default)]
    pub continuous_usage_stats: bool,
}

/// Chat completion request with multimodal processor support.
///
/// Extends upstream `CreateChatCompletionRequest` with:
/// - `mm_processor_kwargs`: multimodal processor configuration (vLLM-specific)
/// - Uses our extended `ChatCompletionRequestMessage` (with reasoning, video/audio)
/// - Uses our extended `ChatCompletionStreamOptions` (with continuous_usage_stats)
#[derive(Clone, Serialize, Default, Debug, Builder, Deserialize, PartialEq)]
#[builder(name = "CreateChatCompletionRequestArgs")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "OpenAIError"))]
pub struct CreateChatCompletionRequest {
    pub messages: Vec<ChatCompletionRequestMessage>,
    pub model: String,
    /// Multimodal processor configuration (vLLM-specific)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mm_processor_kwargs: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub store: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<ReasoningEffort>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<std::collections::HashMap<String, serde_json::Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logprobs: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_logprobs: Option<u8>,
    #[deprecated]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_completion_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u8>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub modalities: Option<Vec<async_openai::types::chat::ResponseModalities>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prediction: Option<PredictionContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<ChatCompletionAudio>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub service_tier: Option<ServiceTier>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<Stop>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream_options: Option<ChatCompletionStreamOptions>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<ChatCompletionTool>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ChatCompletionToolChoiceOption>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parallel_tool_calls: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
    #[deprecated]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_call: Option<ChatCompletionFunctionCall>,
    #[deprecated]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub functions: Option<Vec<ChatCompletionFunctions>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub web_search_options: Option<WebSearchOptions>,
}

/// Chat choice with extended response message.
///
/// Uses our `ChatCompletionResponseMessage` (multimodal content + reasoning).
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct ChatChoice {
    pub index: u32,
    pub message: ChatCompletionResponseMessage,
    pub finish_reason: Option<FinishReason>,
    pub logprobs: Option<ChatChoiceLogprobs>,
}

/// Non-streaming chat completion response.
#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)]
pub struct CreateChatCompletionResponse {
    pub id: String,
    pub choices: Vec<ChatChoice>,
    pub created: u32,
    pub model: String,
    pub service_tier: Option<ServiceTierResponse>,
    pub system_fingerprint: Option<String>,
    pub object: String,
    pub usage: Option<CompletionUsage>,
}

pub type ChatCompletionResponseStream =
    Pin<Box<dyn Stream<Item = Result<CreateChatCompletionStreamResponse, OpenAIError>> + Send>>;

/// Streaming delta with reasoning content.
///
/// Extends upstream `ChatCompletionStreamResponseDelta` with:
/// - `content`: `Option<ChatCompletionMessageContent>` (multimodal) instead of `Option<String>`
/// - `reasoning_content`: streaming reasoning tokens (DeepSeek-R1, QwQ)
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct ChatCompletionStreamResponseDelta {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<ChatCompletionMessageContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_call: Option<ChatCompletionStreamResponseDeltaFunctionCall>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ChatCompletionMessageToolCallChunk>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub role: Option<Role>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<String>,
    /// Streaming reasoning content (DeepSeek-R1, QwQ models).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_content: Option<String>,
}

#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct ChatCompletionStreamResponseDeltaFunctionCall {
    pub name: Option<String>,
    #[serde(default, deserialize_with = "deserialize_arguments_opt")]
    pub arguments: Option<String>,
}

/// Streaming chat choice.
#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)]
pub struct ChatChoiceStream {
    pub index: u32,
    pub delta: ChatCompletionStreamResponseDelta,
    pub finish_reason: Option<FinishReason>,
    pub logprobs: Option<ChatChoiceLogprobs>,
}

/// Streaming chat completion response with extended choices.
#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)]
pub struct CreateChatCompletionStreamResponse {
    pub id: String,
    pub choices: Vec<ChatChoiceStream>,
    pub created: u32,
    pub model: String,
    pub service_tier: Option<ServiceTierResponse>,
    pub system_fingerprint: Option<String>,
    pub object: String,
    pub usage: Option<CompletionUsage>,
}

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

    #[test]
    fn stop_accepts_token_id_array() {
        let stop: Stop = serde_json::from_value(serde_json::json!([32, 34])).unwrap();

        assert_eq!(stop, Stop::TokenIdArray(vec![32, 34]));
    }

    #[test]
    fn stop_accepts_string_and_string_array() {
        let stop: Stop = serde_json::from_value(serde_json::json!(" The")).unwrap();

        assert_eq!(stop, Stop::String(" The".to_string()));

        let stop: Stop = serde_json::from_value(serde_json::json!(["A", "B"])).unwrap();

        assert_eq!(
            stop,
            Stop::StringArray(vec!["A".to_string(), "B".to_string()])
        );
    }

    #[test]
    fn stop_token_id_display_string_remains_string_stop() {
        let stop: Stop = serde_json::from_value(serde_json::json!("token_id:576")).unwrap();

        assert_eq!(stop, Stop::String("token_id:576".to_string()));

        let stop: Stop = serde_json::from_value(serde_json::json!(["token_id:576"])).unwrap();

        assert_eq!(stop, Stop::StringArray(vec!["token_id:576".to_string()]));
    }

    #[test]
    fn stop_rejects_single_token_id() {
        let result = serde_json::from_value::<Stop>(serde_json::json!(576));

        assert!(result.is_err());
    }

    #[test]
    fn stop_converts_from_upstream_stop_configuration() {
        let upstream =
            async_openai::types::chat::StopConfiguration::StringArray(vec!["END".to_string()]);

        assert_eq!(
            Stop::from(upstream),
            Stop::StringArray(vec!["END".to_string()])
        );
    }

    #[test]
    fn request_builder_accepts_upstream_reasoning_effort() {
        let request = CreateChatCompletionRequestArgs::default()
            .reasoning_effort(async_openai::types::chat::ReasoningEffort::High)
            .build()
            .unwrap();

        assert_eq!(request.reasoning_effort, Some(ReasoningEffort::High));
    }

    #[test]
    fn tool_call_defaults_type_on_deserialize() {
        let tool_call: ChatCompletionMessageToolCall = serde_json::from_value(serde_json::json!({
            "id": "call_123",
            "function": {
                "name": "get_weather",
                "arguments": "{\"location\":\"SF\"}"
            }
        }))
        .unwrap();

        assert_eq!(tool_call.r#type, FunctionType::Function);
    }

    #[test]
    fn tool_call_serializes_type_for_wire_compat() {
        let tool_call = ChatCompletionMessageToolCall {
            id: "call_123".into(),
            r#type: FunctionType::Function,
            function: FunctionCall {
                name: "get_weather".into(),
                arguments: "{\"location\":\"SF\"}".into(),
            },
        };

        let json = serde_json::to_value(tool_call).unwrap();
        assert_eq!(json["type"], "function");
    }

    // -- dict-format arguments tests --

    #[test]
    fn function_call_accepts_string_arguments() {
        let fc: FunctionCall = serde_json::from_value(serde_json::json!({
            "name": "get_weather",
            "arguments": "{\"location\":\"SF\"}"
        }))
        .unwrap();
        assert_eq!(fc.arguments, "{\"location\":\"SF\"}");
    }

    #[test]
    fn function_call_accepts_dict_arguments() {
        let fc: FunctionCall = serde_json::from_value(serde_json::json!({
            "name": "get_weather",
            "arguments": {"location": "SF"}
        }))
        .unwrap();
        assert_eq!(fc.arguments, "{\"location\":\"SF\"}");
    }

    #[test]
    fn function_call_rejects_integer_arguments() {
        let result = serde_json::from_value::<FunctionCall>(serde_json::json!({
            "name": "f",
            "arguments": 42
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_rejects_boolean_arguments() {
        let result = serde_json::from_value::<FunctionCall>(serde_json::json!({
            "name": "f",
            "arguments": true
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_rejects_null_arguments() {
        let result = serde_json::from_value::<FunctionCall>(serde_json::json!({
            "name": "f",
            "arguments": null
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_rejects_array_arguments() {
        let result = serde_json::from_value::<FunctionCall>(serde_json::json!({
            "name": "f",
            "arguments": [1, 2, 3]
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_stream_null_arguments_produces_none() {
        let fcs: FunctionCallStream = serde_json::from_value(serde_json::json!({
            "name": "f",
            "arguments": null
        }))
        .unwrap();
        assert_eq!(fcs.arguments, None);
    }

    #[test]
    fn function_call_stream_rejects_integer_arguments() {
        let result = serde_json::from_value::<FunctionCallStream>(serde_json::json!({
            "name": "f",
            "arguments": 42
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_stream_rejects_boolean_arguments() {
        let result = serde_json::from_value::<FunctionCallStream>(serde_json::json!({
            "name": "f",
            "arguments": true
        }));
        assert!(result.is_err());
    }

    #[test]
    fn function_call_stream_accepts_dict_arguments() {
        let fcs: FunctionCallStream = serde_json::from_value(serde_json::json!({
            "name": "get_weather",
            "arguments": {"location": "SF"}
        }))
        .unwrap();
        assert_eq!(fcs.arguments.as_deref(), Some("{\"location\":\"SF\"}"));
    }

    #[test]
    fn function_call_stream_accepts_null_arguments() {
        let fcs: FunctionCallStream = serde_json::from_value(serde_json::json!({
            "name": "get_weather"
        }))
        .unwrap();
        assert_eq!(fcs.arguments, None);
    }

    #[test]
    fn tool_call_with_dict_arguments_roundtrip() {
        let tc: ChatCompletionMessageToolCall = serde_json::from_value(serde_json::json!({
            "id": "call_abc",
            "type": "function",
            "function": {
                "name": "search",
                "arguments": {"query": "hello", "limit": 10}
            }
        }))
        .unwrap();
        // Compare as parsed JSON values since key order is non-deterministic
        let parsed: serde_json::Value = serde_json::from_str(&tc.function.arguments).unwrap();
        assert_eq!(parsed, serde_json::json!({"query": "hello", "limit": 10}));
        // Re-serialisation produces a string, not an object
        let json = serde_json::to_value(&tc).unwrap();
        assert!(json["function"]["arguments"].is_string());
    }

    #[test]
    fn stream_delta_function_call_accepts_dict_arguments() {
        let delta: ChatCompletionStreamResponseDeltaFunctionCall =
            serde_json::from_value(serde_json::json!({
                "name": "get_weather",
                "arguments": {"location": "SF"}
            }))
            .unwrap();
        assert_eq!(delta.arguments.as_deref(), Some("{\"location\":\"SF\"}"));
    }
}