omniference 0.3.2

use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, time::Duration};

// Provider-specific types are organized in the providers module
pub mod providers;

// Re-export provider types for convenience
pub use providers::*;

#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub enum ProviderKind {
	OpenAI,
	OpenAICompat,
	OpenRouter,
	Anthropic,
	Google,
	Custom(String),
}

#[derive(Clone, Serialize, Deserialize)]
pub struct ProviderEndpoint {
	pub kind: ProviderKind,
	pub base_url: String,
	#[serde(default, skip_serializing)]
	pub api_key: Option<String>,
	#[serde(default, skip_serializing)]
	pub extra_headers: BTreeMap<String, String>,
	pub timeout: Option<u64>,
}

impl std::fmt::Debug for ProviderEndpoint {
	fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		formatter
			.debug_struct("ProviderEndpoint")
			.field("kind", &self.kind)
			.field("base_url", &self.base_url)
			.field("api_key", &self.api_key.as_ref().map(|_| "[REDACTED]"))
			.field("extra_headers", &self.extra_headers.keys().collect::<Vec<_>>())
			.field("timeout", &self.timeout)
			.finish()
	}
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProviderConfig {
	pub name: String,
	pub endpoint: ProviderEndpoint,
	pub enabled: bool,
	#[serde(default)]
	pub catalog_provider_slug: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct DiscoveredModel {
	pub id: String,
	pub name: String,
	pub provider_name: String,
	pub provider_kind: ProviderKind,
	pub input_modalities: Vec<Modality>,
	pub output_modalities: Vec<Modality>,
	pub context_length: Option<u32>,
	pub max_tokens: Option<u32>,
	pub capabilities: Vec<ModelCapabilities>,
	#[serde(default)]
	pub pricing: Option<crate::catalog::ModelPricing>,
	#[serde(default)]
	pub reasoning_budget: Option<ReasoningBudget>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ReasoningBudget {
	pub min_tokens: Option<u32>,
	pub max_tokens: Option<u32>,
}

impl ReasoningBudget {
	pub fn from_legacy_capability(capability: &ModelCapabilities) -> Option<Self> {
		let (min_tokens, max_tokens) = match capability {
			ModelCapabilities::ReasoningBudgetTokens_1024_32000 => (1024, 32_000),
			ModelCapabilities::ReasoningBudgetTokens_1024_64000 => (1024, 64_000),
			ModelCapabilities::ReasoningBudgetTokens_128_32768 => (128, 32_768),
			ModelCapabilities::ReasoningBudgetTokens_128_24576 => (128, 24_576),
			_ => return None,
		};
		Some(Self {
			min_tokens: Some(min_tokens),
			max_tokens: Some(max_tokens),
		})
	}

	pub fn legacy_capability(&self) -> Option<ModelCapabilities> {
		match (self.min_tokens, self.max_tokens) {
			(Some(1024), Some(32_000)) => Some(ModelCapabilities::ReasoningBudgetTokens_1024_32000),
			(Some(1024), Some(64_000)) => Some(ModelCapabilities::ReasoningBudgetTokens_1024_64000),
			(Some(128), Some(32_768)) => Some(ModelCapabilities::ReasoningBudgetTokens_128_32768),
			(Some(128), Some(24_576)) => Some(ModelCapabilities::ReasoningBudgetTokens_128_24576),
			_ => None,
		}
	}
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
#[allow(non_camel_case_types)]
pub enum ModelCapabilities {
	ImageGeneration,
	ImageEditing,
	Reasoning,
	ReasoningEffortNone,
	ReasoningEffortMinimal,
	ReasoningEffortLow,
	ReasoningEffortMedium,
	ReasoningEffortHigh,
	#[serde(rename = "REASONING_EFFORT_XHIGH")]
	ReasoningEffortXHigh,
	ReasoningBudgetTokens_1024_32000,
	ReasoningBudgetTokens_1024_64000,
	ReasoningBudgetTokens_128_32768,
	ReasoningBudgetTokens_128_24576,
	Tools,
}

impl ModelCapabilities {
	pub fn as_str(&self) -> &'static str {
		match self {
			Self::ImageGeneration => "IMAGE_GENERATION",
			Self::ImageEditing => "IMAGE_EDITING",
			Self::Reasoning => "REASONING",
			Self::ReasoningEffortNone => "REASONING_EFFORT_NONE",
			Self::ReasoningEffortMinimal => "REASONING_EFFORT_MINIMAL",
			Self::ReasoningEffortLow => "REASONING_EFFORT_LOW",
			Self::ReasoningEffortMedium => "REASONING_EFFORT_MEDIUM",
			Self::ReasoningEffortHigh => "REASONING_EFFORT_HIGH",
			Self::ReasoningEffortXHigh => "REASONING_EFFORT_XHIGH",
			Self::ReasoningBudgetTokens_1024_32000 => "REASONING_BUDGET_TOKENS_1024_32000",
			Self::ReasoningBudgetTokens_1024_64000 => "REASONING_BUDGET_TOKENS_1024_64000",
			Self::ReasoningBudgetTokens_128_32768 => "REASONING_BUDGET_TOKENS_128_32768",
			Self::ReasoningBudgetTokens_128_24576 => "REASONING_BUDGET_TOKENS_128_24576",
			Self::Tools => "TOOLS",
		}
	}
	pub fn from_code(s: &str) -> Option<Self> {
		match s {
			"IMAGE_GENERATION" => Some(Self::ImageGeneration),
			"IMAGE_EDITING" => Some(Self::ImageEditing),
			"REASONING" => Some(Self::Reasoning),
			"REASONING_EFFORT_NONE" => Some(Self::ReasoningEffortNone),
			"REASONING_EFFORT_MINIMAL" => Some(Self::ReasoningEffortMinimal),
			"REASONING_EFFORT_LOW" => Some(Self::ReasoningEffortLow),
			"REASONING_EFFORT_MEDIUM" => Some(Self::ReasoningEffortMedium),
			"REASONING_EFFORT_HIGH" => Some(Self::ReasoningEffortHigh),
			"REASONING_EFFORT_XHIGH" => Some(Self::ReasoningEffortXHigh),
			"REASONING_BUDGET_TOKENS_1024_32000" => Some(Self::ReasoningBudgetTokens_1024_32000),
			"REASONING_BUDGET_TOKENS_1024_64000" => Some(Self::ReasoningBudgetTokens_1024_64000),
			"REASONING_BUDGET_TOKENS_128_32768" => Some(Self::ReasoningBudgetTokens_128_32768),
			"REASONING_BUDGET_TOKENS_128_24576" => Some(Self::ReasoningBudgetTokens_128_24576),
			"TOOLS" => Some(Self::Tools),
			_ => None,
		}
	}

	#[deprecated(since = "0.3.1", note = "use ModelCapabilities::from_code instead")]
	#[allow(clippy::should_implement_trait)]
	pub fn from_str(s: &str) -> Option<Self> {
		Self::from_code(s)
	}
}

#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ImageOperation {
	Generate,
	Edit,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ImageInput {
	pub bytes: Vec<u8>,
	pub media_type: String,
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ImageOptions {
	pub size: Option<String>,
	pub aspect_ratio: Option<String>,
	pub quality: Option<String>,
	pub output_format: Option<String>,
	pub background: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ImageRequestIR {
	pub model: ModelRef,
	pub operation: ImageOperation,
	pub prompt: String,
	#[serde(default)]
	pub request_id: Option<String>,
	#[serde(default)]
	pub input_images: Vec<ImageInput>,
	#[serde(default)]
	pub options: ImageOptions,
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ImageUsage {
	pub input_tokens: u64,
	pub output_tokens: u64,
	pub input_images: u32,
	pub output_images: u32,
	pub provider_cost: Option<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ImageOutput {
	pub bytes: Vec<u8>,
	pub media_type: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ImageResponse {
	pub images: Vec<ImageOutput>,
	pub usage: ImageUsage,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelCapabilitiesWithModalities {
	pub context_length: Option<u32>,
	pub max_tokens: Option<u32>,
	pub capabilities: Vec<ModelCapabilities>,
	pub input_modalities: Vec<Modality>,
	pub output_modalities: Vec<Modality>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum Modality {
	Text,
	Image,
	File,
	Audio,
	Video,
	Embeddings,
}

impl Modality {
	pub fn as_str(&self) -> &'static str {
		match self {
			Self::Text => "TEXT",
			Self::Image => "IMAGE",
			Self::File => "FILE",
			Self::Audio => "AUDIO",
			Self::Video => "VIDEO",
			Self::Embeddings => "EMBEDDINGS",
		}
	}
	pub fn from_code(s: &str) -> Option<Self> {
		match s {
			"TEXT" => Some(Self::Text),
			"IMAGE" => Some(Self::Image),
			"FILE" => Some(Self::File),
			"AUDIO" => Some(Self::Audio),
			"VIDEO" => Some(Self::Video),
			"EMBEDDINGS" => Some(Self::Embeddings),
			_ => None,
		}
	}

	#[deprecated(since = "0.3.1", note = "use Modality::from_code instead")]
	#[allow(clippy::should_implement_trait)]
	pub fn from_str(s: &str) -> Option<Self> {
		Self::from_code(s)
	}
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelRef {
	pub alias: String,
	pub provider: ProviderConfig,
	pub model_id: String,
	pub input_modalities: Vec<Modality>,
	pub output_modalities: Vec<Modality>,
}

#[derive(Clone, Debug, Serialize, Deserialize, Default)]
pub struct Sampling {
	pub temperature: Option<f32>,
	pub top_p: Option<f32>,
	pub top_k: Option<u32>,
	pub max_tokens: Option<u32>,
	pub presence_penalty: Option<f32>,
	pub frequency_penalty: Option<f32>,
	pub stop: Vec<String>,
	pub parallel_tool_calls: Option<bool>,
	pub seed: Option<i64>,
	pub logit_bias: Option<std::collections::HashMap<String, f32>>,
	pub logprobs: Option<bool>,
	pub top_logprobs: Option<u32>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ToolSpec {
	JsonSchema {
		name: String,
		description: Option<String>,
		schema: serde_json::Value,
		strict: Option<bool>,
	},
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ToolChoice {
	Auto,
	None,
	Required,
	Named(String),
	Allowed { mode: String, tools: Vec<String> },
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ContentPart {
	Text(String),
	ImageUrl {
		url: String,
		mime: Option<String>,
	},
	BlobRef {
		id: String,
		mime: String,
	},
	Audio {
		data: String,
		format: String,
	},
	File {
		file_id: Option<String>,
		filename: Option<String>,
		file_data: Option<String>,
	},
	/// Tool call from an assistant message
	ToolCall {
		id: String,
		name: String,
		arguments: String,
	},
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ResponseFormat {
	Text,
	JsonObject,
	JsonSchema {
		name: String,
		description: Option<String>,
		schema: serde_json::Value,
		strict: Option<bool>,
	},
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AudioOutput {
	pub voice: Option<String>,
	pub format: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WebSearchOptions {
	pub user_location: Option<UserLocation>,
	pub search_context_size: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct UserLocation {
	pub country: Option<String>,
	pub region: Option<String>,
	pub city: Option<String>,
	pub timezone: Option<String>,
}

/// Configuration for reasoning/thinking capabilities.
/// Supports both OpenAI-style effort levels and Anthropic-style token budgets.
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
pub struct ReasoningConfig {
	/// OpenAI-style reasoning effort level.
	/// Supported values: "none", "minimal", "low", "medium", "high", "xhigh"
	pub effort: Option<String>,
	/// Anthropic-style thinking budget in tokens.
	/// Specifies the maximum number of tokens the model can use for internal reasoning.
	pub budget_tokens: Option<u32>,
	/// Whether to include a summary of the reasoning in the response.
	/// OpenAI: "auto", "concise", "detailed"
	/// Anthropic: uses separate thinking block
	pub summary: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PredictionConfig {
	pub content: Option<PredictionContent>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum PredictionContent {
	Text(String),
	Parts(Vec<ContentPart>),
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum Role {
	Developer,
	System,
	User,
	Assistant,
	Tool,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Message {
	pub role: Role,
	pub parts: Vec<ContentPart>,
	pub name: Option<String>,
}

/// Gateway-agnostic provider-routing preferences.
///
/// Carries an upstream-provider preference (e.g. a user picking which OpenRouter provider serves
/// a request) down to adapters that support it. Adapters that have no notion of provider routing
/// ignore this. Currently only the OpenRouter adapter maps it (onto `provider.order`/`only`).
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ProviderRouting {
	/// Ordered list of provider slugs to prefer (router still falls back unless
	/// `allow_fallbacks` is `Some(false)`).
	pub order: Option<Vec<String>>,
	/// Allowlist of provider slugs — the request is restricted to these.
	pub only: Option<Vec<String>>,
	/// Whether to allow fallback to other providers (default at the gateway is `true`).
	pub allow_fallbacks: Option<bool>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ChatRequestIR {
	pub model: ModelRef,
	pub messages: Vec<Message>,
	pub tools: Vec<ToolSpec>,
	pub tool_choice: ToolChoice,
	pub sampling: Sampling,
	pub stream: bool,
	pub response_format: Option<ResponseFormat>,
	pub audio_output: Option<AudioOutput>,
	pub web_search_options: Option<WebSearchOptions>,
	pub prediction: Option<PredictionConfig>,
	pub reasoning: Option<ReasoningConfig>,
	pub metadata: BTreeMap<String, String>,
	pub request_timeout: Option<Duration>,
	pub cache_key: Option<String>,
	pub safety_identifier: Option<String>,
	/// Original Chat Completions request retained for lossless OpenAI-compatible forwarding.
	pub openai_chat_request: Option<Box<providers::openai::OpenAIChatRequest>>,
	/// Original Responses request retained for lossless OpenAI-compatible forwarding.
	pub openai_responses_request: Option<Box<serde_json::Value>>,
	/// Optional upstream-provider routing preferences (gateway adapters only).
	pub provider_routing: Option<ProviderRouting>,
}

impl Default for ChatRequestIR {
	fn default() -> Self {
		Self {
			model: ModelRef {
				alias: String::new(),
				provider: ProviderConfig {
					name: String::new(),
					enabled: true,
					endpoint: ProviderEndpoint {
						kind: ProviderKind::OpenAI,
						base_url: String::new(),
						api_key: None,
						extra_headers: BTreeMap::new(),
						timeout: None,
					},
					catalog_provider_slug: None,
				},
				model_id: String::new(),
				input_modalities: vec![Modality::Text],
				output_modalities: vec![Modality::Text],
			},
			messages: Vec::new(),
			tools: Vec::new(),
			tool_choice: ToolChoice::Auto,
			sampling: Sampling::default(),
			stream: false,
			response_format: None,
			audio_output: None,
			web_search_options: None,
			prediction: None,
			reasoning: None,
			metadata: BTreeMap::new(),
			request_timeout: None,
			cache_key: None,
			safety_identifier: None,
			openai_chat_request: None,
			openai_responses_request: None,
			provider_routing: None,
		}
	}
}