openmodex 0.1.1

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// ─── Chat Completion ─────────────────────────────────────────────────────────

/// Request body for `POST /v1/chat/completions`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u32>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<i64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,

    /// OpenModex routing configuration extension.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub routing: Option<RoutingConfig>,

    /// OpenModex cache configuration extension.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cache: Option<CacheConfig>,
}

impl ChatCompletionRequest {
    /// Create a new chat completion request for the given model.
    pub fn new(model: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            messages: Vec::new(),
            stream: None,
            temperature: None,
            top_p: None,
            n: None,
            max_tokens: None,
            stop: None,
            frequency_penalty: None,
            presence_penalty: None,
            logit_bias: None,
            tools: None,
            tool_choice: None,
            response_format: None,
            seed: None,
            user: None,
            routing: None,
            cache: None,
        }
    }

    /// Add a message to the request.
    pub fn message(mut self, msg: ChatMessage) -> Self {
        self.messages.push(msg);
        self
    }

    /// Add multiple messages to the request.
    pub fn messages(mut self, msgs: impl IntoIterator<Item = ChatMessage>) -> Self {
        self.messages.extend(msgs);
        self
    }

    /// Set the sampling temperature.
    pub fn temperature(mut self, t: f64) -> Self {
        self.temperature = Some(t);
        self
    }

    /// Set the top-p (nucleus) sampling parameter.
    pub fn top_p(mut self, p: f64) -> Self {
        self.top_p = Some(p);
        self
    }

    /// Set the number of completions to generate.
    pub fn n(mut self, n: u32) -> Self {
        self.n = Some(n);
        self
    }

    /// Set the maximum number of tokens to generate.
    pub fn max_tokens(mut self, max: u32) -> Self {
        self.max_tokens = Some(max);
        self
    }

    /// Set stop sequences.
    pub fn stop(mut self, stop: impl Into<serde_json::Value>) -> Self {
        self.stop = Some(stop.into());
        self
    }

    /// Set the frequency penalty.
    pub fn frequency_penalty(mut self, p: f64) -> Self {
        self.frequency_penalty = Some(p);
        self
    }

    /// Set the presence penalty.
    pub fn presence_penalty(mut self, p: f64) -> Self {
        self.presence_penalty = Some(p);
        self
    }

    /// Set the seed for deterministic generation.
    pub fn seed(mut self, seed: i64) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Set the user identifier.
    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    /// Set the tools available for the model.
    pub fn tools(mut self, tools: serde_json::Value) -> Self {
        self.tools = Some(tools);
        self
    }

    /// Set the tool choice.
    pub fn tool_choice(mut self, choice: serde_json::Value) -> Self {
        self.tool_choice = Some(choice);
        self
    }

    /// Set the response format.
    pub fn response_format(mut self, format: serde_json::Value) -> Self {
        self.response_format = Some(format);
        self
    }

    /// Set the OpenModex routing configuration.
    pub fn routing(mut self, routing: RoutingConfig) -> Self {
        self.routing = Some(routing);
        self
    }

    /// Set the OpenModex cache configuration.
    pub fn cache(mut self, cache: CacheConfig) -> Self {
        self.cache = Some(cache);
        self
    }
}

/// A single message in a chat conversation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<serde_json::Value>,
}

impl ChatMessage {
    /// Create a user message.
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: Some(content.into()),
            name: None,
            tool_call_id: None,
            tool_calls: None,
        }
    }

    /// Create a system message.
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: Some(content.into()),
            name: None,
            tool_call_id: None,
            tool_calls: None,
        }
    }

    /// Create an assistant message.
    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: Some(content.into()),
            name: None,
            tool_call_id: None,
            tool_calls: None,
        }
    }

    /// Create a tool result message.
    pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: "tool".to_string(),
            content: Some(content.into()),
            name: None,
            tool_call_id: Some(tool_call_id.into()),
            tool_calls: None,
        }
    }
}

/// Response from `POST /v1/chat/completions`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionResponse {
    pub id: String,
    pub object: String,
    pub created: i64,
    pub model: String,
    pub choices: Vec<ChatChoice>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_fingerprint: Option<String>,

    /// OpenModex-specific metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub openmodex: Option<OpenModexMetadata>,
}

/// A single choice in a chat completion response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatChoice {
    pub index: u32,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<ChatMessage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub delta: Option<ChatMessage>,

    pub finish_reason: Option<String>,
}

// ─── Chat Completion Streaming ───────────────────────────────────────────────

/// A single chunk in a streaming chat completion response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
    pub id: String,
    pub object: String,
    pub created: i64,
    pub model: String,
    pub choices: Vec<StreamChoice>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_fingerprint: Option<String>,

    /// OpenModex-specific metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub openmodex: Option<OpenModexMetadata>,
}

/// A single choice in a streaming chunk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamChoice {
    pub index: u32,
    pub delta: StreamDelta,
    pub finish_reason: Option<String>,
}

/// The incremental content delta in a streaming chunk.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct StreamDelta {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub role: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<serde_json::Value>,
}

// ─── Completion (Legacy) ─────────────────────────────────────────────────────

/// Request body for `POST /v1/completions`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionRequest {
    pub model: String,
    pub prompt: serde_json::Value,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u32>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<serde_json::Value>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub suffix: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,

    /// OpenModex routing configuration extension.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub routing: Option<RoutingConfig>,
}

impl CompletionRequest {
    /// Create a new completion request with a string prompt.
    pub fn new(model: impl Into<String>, prompt: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            prompt: serde_json::Value::String(prompt.into()),
            max_tokens: None,
            temperature: None,
            top_p: None,
            n: None,
            stream: None,
            stop: None,
            frequency_penalty: None,
            presence_penalty: None,
            logit_bias: None,
            suffix: None,
            user: None,
            routing: None,
        }
    }

    /// Set the maximum number of tokens.
    pub fn max_tokens(mut self, max: u32) -> Self {
        self.max_tokens = Some(max);
        self
    }

    /// Set the sampling temperature.
    pub fn temperature(mut self, t: f64) -> Self {
        self.temperature = Some(t);
        self
    }

    /// Set the top-p sampling parameter.
    pub fn top_p(mut self, p: f64) -> Self {
        self.top_p = Some(p);
        self
    }

    /// Set stop sequences.
    pub fn stop(mut self, stop: impl Into<serde_json::Value>) -> Self {
        self.stop = Some(stop.into());
        self
    }

    /// Set the suffix to append after the completion.
    pub fn suffix(mut self, suffix: impl Into<String>) -> Self {
        self.suffix = Some(suffix.into());
        self
    }

    /// Set the user identifier.
    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    /// Set the OpenModex routing configuration.
    pub fn routing(mut self, routing: RoutingConfig) -> Self {
        self.routing = Some(routing);
        self
    }
}

/// Response from `POST /v1/completions`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionResponse {
    pub id: String,
    pub object: String,
    pub created: i64,
    pub model: String,
    pub choices: Vec<CompletionChoice>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_fingerprint: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub openmodex: Option<OpenModexMetadata>,
}

/// A single choice in a completion response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionChoice {
    pub index: u32,
    pub text: String,
    pub finish_reason: Option<String>,
}

// ─── Embeddings ──────────────────────────────────────────────────────────────

/// Request body for `POST /v1/embeddings`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingRequest {
    pub model: String,
    pub input: serde_json::Value,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub encoding_format: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub dimensions: Option<u32>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

impl EmbeddingRequest {
    /// Create an embedding request with a single string input.
    pub fn new(model: impl Into<String>, input: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            input: serde_json::Value::String(input.into()),
            encoding_format: None,
            dimensions: None,
            user: None,
        }
    }

    /// Create an embedding request with multiple string inputs.
    pub fn new_batch(model: impl Into<String>, inputs: Vec<String>) -> Self {
        Self {
            model: model.into(),
            input: serde_json::Value::Array(inputs.into_iter().map(serde_json::Value::String).collect()),
            encoding_format: None,
            dimensions: None,
            user: None,
        }
    }

    /// Set the encoding format (e.g. "float", "base64").
    pub fn encoding_format(mut self, format: impl Into<String>) -> Self {
        self.encoding_format = Some(format.into());
        self
    }

    /// Set the number of dimensions for the embedding.
    pub fn dimensions(mut self, dims: u32) -> Self {
        self.dimensions = Some(dims);
        self
    }

    /// Set the user identifier.
    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }
}

/// Response from `POST /v1/embeddings`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingResponse {
    pub object: String,
    pub data: Vec<EmbeddingData>,
    pub model: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub openmodex: Option<OpenModexMetadata>,
}

/// A single embedding vector.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingData {
    pub object: String,
    pub embedding: Vec<f64>,
    pub index: u32,
}

// ─── Models ──────────────────────────────────────────────────────────────────

/// A model available through the OpenModex API.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Model {
    pub id: String,
    pub object: String,
    pub name: String,
    pub provider: String,

    #[serde(default)]
    pub description: String,

    #[serde(default)]
    pub category: Vec<String>,

    #[serde(default)]
    pub context_length: u64,

    #[serde(default)]
    pub max_output_tokens: u64,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub modalities: Option<Modalities>,

    #[serde(default)]
    pub features: Vec<String>,

    #[serde(default)]
    pub license: String,

    #[serde(default)]
    pub language_optimization: Vec<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub pricing: Option<ModelPricing>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub performance: Option<ModelPerformance>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub quality_scores: Option<QualityScores>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage_stats: Option<ModelUsageStats>,

    #[serde(default)]
    pub recommended_for: Vec<String>,

    #[serde(default)]
    pub created_at: String,

    #[serde(default)]
    pub updated_at: String,
}

/// Input/output modalities of a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Modalities {
    pub input: Vec<String>,
    pub output: Vec<String>,
}

/// Pricing information for a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPricing {
    #[serde(default)]
    pub official_input_per_m: f64,
    #[serde(default)]
    pub official_output_per_m: f64,
    #[serde(default)]
    pub openmodex_input_per_m: f64,
    #[serde(default)]
    pub openmodex_output_per_m: f64,
    #[serde(default)]
    pub byok_fee_per_m: f64,
    #[serde(default)]
    pub cached_input_per_m: f64,
    #[serde(default)]
    pub currency: String,
}

/// Performance metrics for a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPerformance {
    #[serde(default)]
    pub availability: f64,
    #[serde(default)]
    pub latency_p50_ms: u64,
    #[serde(default)]
    pub latency_p99_ms: u64,
    #[serde(default)]
    pub throughput_rpm: u64,
}

/// Benchmark quality scores for a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QualityScores {
    #[serde(default)]
    pub overall: u32,
    #[serde(default)]
    pub code: u32,
    #[serde(default)]
    pub math: u32,
    #[serde(default)]
    pub chinese: u32,
    #[serde(default)]
    pub creative_writing: u32,
    #[serde(default)]
    pub multi_turn: u32,
    #[serde(default)]
    pub instruction_following: u32,
    #[serde(default)]
    pub reasoning: u32,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub source: Option<String>,
}

/// Usage statistics for a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelUsageStats {
    #[serde(default)]
    pub requests_7d: i64,
    #[serde(default)]
    pub trend: String,
    #[serde(default)]
    pub rank: u32,
}

/// Response from `GET /v1/models`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelListResponse {
    pub object: String,
    pub data: Vec<Model>,

    #[serde(default)]
    pub total: i64,

    #[serde(default)]
    pub has_more: bool,
}

/// Response from `GET /v1/models/compare`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCompareResponse {
    pub object: String,
    pub models: Vec<String>,
    pub comparison: HashMap<String, ComparisonItem>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub highlights: Option<ComparisonHighlights>,
}

/// A single model in a side-by-side comparison.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonItem {
    pub name: String,
    pub provider: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub pricing: Option<ComparisonPricing>,

    #[serde(default)]
    pub context_length: u64,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub performance: Option<ComparisonPerformance>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub quality_scores: Option<ComparisonQuality>,
}

/// Pricing in a model comparison.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonPricing {
    #[serde(default)]
    pub openmodex_input_per_m: f64,
    #[serde(default)]
    pub openmodex_output_per_m: f64,
}

/// Performance metrics in a model comparison.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonPerformance {
    #[serde(default)]
    pub latency_p50_ms: u64,
    #[serde(default)]
    pub availability: f64,
}

/// Quality scores in a model comparison.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonQuality {
    #[serde(default)]
    pub overall: u32,
    #[serde(default)]
    pub code: u32,
    #[serde(default)]
    pub math: u32,
    #[serde(default)]
    pub chinese: u32,
    #[serde(default)]
    pub reasoning: u32,
}

/// Winners in each comparison category.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComparisonHighlights {
    #[serde(default)]
    pub cheapest: String,
    #[serde(default)]
    pub fastest: String,
    #[serde(default)]
    pub best_quality: String,
    #[serde(default)]
    pub best_code: String,
    #[serde(default)]
    pub best_chinese: String,
    #[serde(default)]
    pub best_value: String,
    #[serde(default)]
    pub longest_context: String,
}

// ─── Shared ──────────────────────────────────────────────────────────────────

/// Token usage statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Usage {
    #[serde(default)]
    pub prompt_tokens: u32,
    #[serde(default)]
    pub completion_tokens: u32,
    #[serde(default)]
    pub total_tokens: u32,

    #[serde(default, skip_serializing_if = "is_zero_f64")]
    pub cost_usd: f64,
}

fn is_zero_f64(v: &f64) -> bool {
    *v == 0.0
}

/// OpenModex-specific response metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OpenModexMetadata {
    #[serde(default)]
    pub request_id: String,
    #[serde(default)]
    pub model_used: String,
    #[serde(default)]
    pub provider: String,
    #[serde(default)]
    pub cache_hit: bool,
    #[serde(default)]
    pub routing_strategy: String,
    #[serde(default)]
    pub fallback_used: bool,
    #[serde(default)]
    pub latency_ms: i64,
}

/// OpenModex routing configuration extension.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub strategy: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub fallback: Option<Vec<String>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub allow_upgrade: Option<bool>,
}

impl RoutingConfig {
    /// Create a new routing config with the given strategy.
    pub fn new(strategy: impl Into<String>) -> Self {
        Self {
            strategy: Some(strategy.into()),
            fallback: None,
            allow_upgrade: None,
        }
    }

    /// Set the fallback model chain.
    pub fn fallback(mut self, models: Vec<String>) -> Self {
        self.fallback = Some(models);
        self
    }

    /// Set whether to allow model upgrades.
    pub fn allow_upgrade(mut self, allow: bool) -> Self {
        self.allow_upgrade = Some(allow);
        self
    }
}

/// OpenModex cache configuration extension.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub enabled: Option<bool>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub ttl: Option<u32>,
}

impl CacheConfig {
    /// Create a cache config with caching enabled.
    pub fn enabled() -> Self {
        Self {
            enabled: Some(true),
            ttl: None,
        }
    }

    /// Create a cache config with caching disabled.
    pub fn disabled() -> Self {
        Self {
            enabled: Some(false),
            ttl: None,
        }
    }

    /// Set the cache TTL in seconds.
    pub fn ttl(mut self, ttl: u32) -> Self {
        self.ttl = Some(ttl);
        self
    }
}