orign 0.2.3

// src/models.rs

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

pub trait Request: Serialize + for<'de> Deserialize<'de> {
    fn set_user_id(&mut self, user_id: String);
    fn set_output_topic(&mut self, output_topic: String);
    fn set_request_id(&mut self, request_id: String);
    fn set_organizations(&mut self, organizations: HashMap<String, HashMap<String, String>>);
    fn set_handle(&mut self, handle: String);
}

// Implement Request trait for ChatRequest
impl Request for ChatRequest {
    fn set_user_id(&mut self, user_id: String) {
        self.user_id = Some(user_id);
    }
    fn set_output_topic(&mut self, output_topic: String) {
        self.output_topic = Some(output_topic);
    }
    fn set_request_id(&mut self, request_id: String) {
        self.request_id = Some(request_id);
    }
    fn set_organizations(&mut self, organizations: HashMap<String, HashMap<String, String>>) {
        self.organizations = Some(organizations);
    }
    fn set_handle(&mut self, handle: String) {
        self.handle = Some(handle);
    }
}

// Implement Request trait for EmbeddingRequest
impl Request for EmbeddingRequest {
    fn set_user_id(&mut self, user_id: String) {
        self.user_id = Some(user_id);
    }
    fn set_output_topic(&mut self, output_topic: String) {
        self.output_topic = Some(output_topic);
    }
    fn set_request_id(&mut self, request_id: String) {
        self.request_id = Some(request_id);
    }
    fn set_organizations(&mut self, organizations: HashMap<String, HashMap<String, String>>) {
        self.organizations = Some(organizations);
    }
    fn set_handle(&mut self, handle: String) {
        self.handle = Some(handle);
    }
}

// Implement Request trait for OCRRequest
impl Request for OCRRequest {
    fn set_user_id(&mut self, user_id: String) {
        self.user_id = Some(user_id);
    }
    fn set_output_topic(&mut self, output_topic: String) {
        self.output_topic = Some(output_topic);
    }
    fn set_request_id(&mut self, request_id: String) {
        self.request_id = Some(request_id);
    }
    fn set_organizations(&mut self, organizations: HashMap<String, HashMap<String, String>>) {
        self.organizations = Some(organizations);
    }
    fn set_handle(&mut self, handle: String) {
        self.handle = Some(handle);
    }
}

// === Streaming Chat Request ===

#[derive(Deserialize, Serialize, Default)]
/// Image URL content for chat requests
pub struct ImageUrlContent {
    pub url: String,
}

#[derive(Deserialize, Serialize, Default)]
/// Content item for chat requests
pub struct ContentItem {
    #[serde(rename = "type")]
    pub content_type: String,
    pub text: Option<String>,
    pub image_url: Option<ImageUrlContent>,
}

#[derive(Deserialize, Serialize, Default)]
/// Batch item with content for chat requests
pub struct BatchItemWithContent {
    pub role: String,
    pub content: Vec<ContentItem>,
}

#[derive(Deserialize, Serialize)]
#[serde(untagged)]
/// Message content can be either a string or a ContentItem
pub enum MessageContent {
    Text(String),
    Item(Vec<ContentItem>),
}

#[derive(Deserialize, Serialize)]
/// Message item for chat requests
pub struct MessageItem {
    pub role: String,
    pub content: MessageContent,
}

#[derive(Deserialize, Serialize)]
/// Prompt for chat requests
pub struct Prompt {
    pub messages: Vec<MessageItem>,
}

#[derive(Deserialize, Serialize, Default)]
/// Sampling parameters for chat requests
pub struct SamplingParams {
    pub n: i32,
    pub best_of: Option<i32>,
    pub presence_penalty: f32,
    pub frequency_penalty: f32,
    pub repetition_penalty: f32,
    pub temperature: f32,
    pub top_p: f32,
    pub top_k: i32,
    pub min_p: f32,
    pub seed: Option<i32>,
    pub stop: Option<Vec<String>>,
    pub stop_token_ids: Option<Vec<i32>>,
    pub min_tokens: i32,
    pub logprobs: Option<i32>,
    pub prompt_logprobs: Option<i32>,
    pub detokenize: bool,
    pub skip_special_tokens: bool,
    pub spaces_between_special_tokens: bool,
    pub truncate_prompt_tokens: Option<i32>,
}

#[derive(Deserialize, Serialize, Default)]
pub struct Usage {
    pub prompt_tokens: usize,
    pub completion_tokens: usize,
    pub total_tokens: usize,
}

#[derive(Deserialize, Serialize, Default)]
/// Chat request
pub struct ChatRequest {
    #[serde(rename = "type")]
    #[serde(default = "default_chat_request_type")]
    pub request_type: String,
    pub request_id: Option<String>,
    pub model: Option<String>,
    pub framework: Option<String>,
    pub kind: Option<String>,
    pub namespace: Option<String>,
    pub prompt: Option<Prompt>,
    pub batch: Option<Vec<Prompt>>,
    pub adapter: Option<String>,
    pub max_tokens: Option<i32>,
    pub sampling_params: Option<SamplingParams>,
    #[serde(default)]
    pub stream: bool,
    pub user_id: Option<String>,
    pub organizations: Option<HashMap<String, HashMap<String, String>>>,
    pub handle: Option<String>,
    pub output_topic: Option<String>,
    pub output_partition: Option<i32>,
}

fn default_chat_request_type() -> String {
    "ChatRequest".to_string()
}

#[derive(Deserialize, Serialize, Default)]
/// Individual choice in the token response
pub struct Choice {
    pub index: i32,
    pub text: String,
    pub tokens: Option<Vec<String>>,
    pub token_ids: Option<Vec<i32>>,
    pub logprobs: Option<Vec<HashMap<String, serde_json::Value>>>,
    pub finish_reason: Option<String>,
}

#[derive(Deserialize, Serialize, Default)]
/// Chat response
pub struct ChatResponse {
    #[serde(rename = "type", default = "default_chat_response_type")]
    pub response_type: String,
    pub request_id: String,
    pub choices: Vec<Choice>,
    pub usage: Option<Usage>,
    pub trip_time: Option<f64>,
}

fn default_chat_response_type() -> String {
    "ChatResponse".to_string()
}

#[derive(Deserialize, Serialize, Default)]
/// Token response
pub struct TokenResponse {
    #[serde(rename = "type", default = "default_token_response_type")]
    pub response_type: String,
    pub request_id: String,
    pub tokens: Vec<String>,
    pub token_ids: Option<Vec<i32>>,
    pub usage: Option<Usage>,
    pub logprobs: Option<Vec<HashMap<String, serde_json::Value>>>,
}

fn default_token_response_type() -> String {
    "TokenResponse".to_string()
}

fn default_chat_event_type() -> String {
    "ChatEvent".to_string()
}

#[derive(Deserialize, Serialize, Default)]
pub struct ChatEvent {
    #[serde(rename = "type", default = "default_chat_event_type")]
    pub event_type: String,
    pub id: Option<String>,
    pub request: ChatRequest,
    pub response: ChatResponse,
    pub token_count: Option<i32>,
    pub trip_time: Option<f64>,
    pub approved: Option<bool>,
    pub metadata: Option<HashMap<String, String>>,
    pub owner_id: Option<String>,
    pub organization_id: Option<String>,
    pub handle: Option<String>,
    pub created: Option<i64>,
}

// === Model Instances ===

#[derive(Deserialize, Serialize)]
pub struct ModelReadyResponse {
    #[serde(rename = "type")]
    pub response_type: String,
    pub request_id: String,
    pub ready: bool,
    pub error: Option<String>,
}

impl Default for ModelReadyResponse {
    fn default() -> Self {
        Self {
            response_type: "ModelReadyResponse".to_string(),
            request_id: "".to_string(),
            ready: false,
            error: None,
        }
    }
}

#[derive(Serialize, Deserialize, Default)]
pub struct ReplicasInfo {
    pub desired: i32,
    pub ready: i32,
    pub updated: i32,
    pub available: i32,
    pub unavailable: i32,
}

// class SentenceTFConfig(BaseSettings):
//     model: str = "clip-ViT-B-32"
//     device: str = "cuda"
#[derive(Deserialize, Serialize)]
pub struct SentenceTFParams {
    #[serde(default = "default_sentence_tf_model")]
    pub model: String,

    #[serde(default = "default_cuda_device")]
    pub device: String,
}

impl Default for SentenceTFParams {
    fn default() -> Self {
        Self {
            model: default_sentence_tf_model(),
            device: default_cuda_device(),
        }
    }
}
// class vLLMConfig(BaseSettings):
//     model: str
//     trust_remote_code: bool = True
//     tensor_parallel_size: int = 1
//     dtype: str = "auto"
//     max_images_per_prompt: int = 1
//     device: str = "cuda"
#[derive(Deserialize, Serialize)]
pub struct VLLMParams {
    pub model: String,
    pub model_type: Option<String>,

    #[serde(default = "default_option_true")]
    pub trust_remote_code: Option<bool>,

    #[serde(default = "default_tensor_parallel_size")]
    pub tensor_parallel_size: Option<i32>,

    #[serde(default = "default_auto_dtype")]
    pub dtype: Option<String>,

    #[serde(default = "default_max_images")]
    pub max_images_per_prompt: Option<i32>,

    #[serde(default = "default_cuda_device_option")]
    pub device: Option<String>,

    #[serde(default = "default_max_model_len")]
    pub max_model_len: Option<i32>,

    #[serde(default = "default_max_num_seqs")]
    pub max_num_seqs: Option<i32>,

    #[serde(default = "default_gpu_memory_utilization")]
    pub gpu_memory_utilization: Option<f32>,

    #[serde(default = "default_enforce_eager")]
    pub enforce_eager: Option<bool>,

    #[serde(default = "default_enable_adapter")]
    pub enable_adapter: Option<bool>,
}

// class DoctrConfig(BaseSettings):
//     det_arch: str = "fast_base"
//     reco_arch: str = "crnn_vgg16_bn"
//     pretrained: bool = True
#[derive(Deserialize, Serialize)]
pub struct DoctrParams {
    #[serde(default = "default_det_arch")]
    pub det_arch: Option<String>,

    #[serde(default = "default_reco_arch")]
    pub reco_arch: Option<String>,

    #[serde(default = "default_option_true")]
    pub pretrained: Option<bool>,
}

// class EasyOCRConfig(BaseSettings):
//     device: str = "cuda"
//     gpu: bool = True
//     lang_list: list[str] = ["en"]
//     quantize: bool = False
#[derive(Deserialize, Serialize)]
pub struct EasyOCRParams {
    #[serde(default = "default_cuda_device_option")]
    pub device: Option<String>,

    #[serde(default = "default_option_true")]
    pub gpu: Option<bool>,

    #[serde(default = "default_lang_list")]
    pub lang_list: Option<Vec<String>>,

    #[serde(default = "default_option_false")]
    pub quantize: Option<bool>,
}

// class LiteLLMConfig(BaseSettings):
//     api_keys: dict[str, str] = {}
//     lang_list: list[str] = ["en"]
//     quantize: bool = False
#[derive(Deserialize, Serialize)]
pub struct LiteLLMParams {
    #[serde()]
    pub api_keys: HashMap<String, String>,
}

#[derive(Deserialize, Serialize, Default)]
pub struct ModelDeployment {
    pub id: String,
    pub namespace: String,
    pub framework: String,
    pub kind: String,
    pub replicas: Option<ReplicasInfo>,
    pub status: String,
    pub image: String,
    pub vram_request: Option<String>,
    pub accelerators: Option<Vec<String>>,
    pub gpu_type: Option<String>,
    pub cpu_request: Option<String>,
    pub vllm_params: Option<VLLMParams>,
    pub sentence_tf_params: Option<SentenceTFParams>,
    pub doctr_params: Option<DoctrParams>,
    pub easyocr_params: Option<EasyOCRParams>,
    pub litellm_params: Option<LiteLLMParams>,
    pub platform: Option<String>,
}

#[derive(Deserialize, Serialize, Default)]
pub struct ModelDeploymentRequest {
    pub namespace: Option<String>,
    pub framework: String,
    pub vram_request: Option<String>,
    pub accelerators: Option<Vec<String>>,
    pub memory_request: Option<String>,
    pub cpu_request: Option<String>,
    pub vllm_params: Option<VLLMParams>,
    pub sentence_tf_params: Option<SentenceTFParams>,
    pub doctr_params: Option<DoctrParams>,
    pub easyocr_params: Option<EasyOCRParams>,
    pub litellm_params: Option<LiteLLMParams>,
    pub max_pixels: Option<i32>,
    pub platform: Option<String>,
}

// === OCR Models ===

#[derive(Deserialize, Serialize, Default)]
pub struct OCRRequest {
    #[serde(rename = "type")]
    #[serde(default = "default_ocr_request_type")]
    pub request_type: String,

    pub request_id: Option<String>,
    pub model: Option<String>,
    pub framework: Option<String>,
    pub image: String,
    pub languages: Vec<String>,

    #[serde(default = "default_true")]
    pub gpu: bool,

    #[serde(default = "default_true")]
    pub detail: bool,

    #[serde(default = "default_false")]
    pub paragraph: bool,

    #[serde(default = "default_min_confidence")]
    pub min_confidence: Option<f32>,

    pub user_id: Option<String>,
    pub organizations: Option<HashMap<String, HashMap<String, String>>>,
    pub handle: Option<String>,
    pub output_topic: Option<String>,
    pub output_partition: Option<i32>,
}

fn default_ocr_request_type() -> String {
    "OCRRequest".to_string()
}

fn default_min_confidence() -> Option<f32> {
    Some(0.0)
}

#[derive(Deserialize, Serialize, Default)]
pub struct BoundingBox {
    pub points: Vec<Vec<i32>>,
    pub text: String,
    pub confidence: f32,
}

#[derive(Deserialize, Serialize)]
#[serde(untagged)]
pub enum OCRResult {
    Detailed(Vec<BoundingBox>),
    Simple(Vec<String>),
}

impl Default for OCRResult {
    fn default() -> Self {
        OCRResult::Detailed(Vec::new())
    }
}

#[derive(Deserialize, Serialize, Default)]
pub struct OCRResponse {
    #[serde(rename = "type")]
    #[serde(default = "default_ocr_response_type")]
    pub response_type: String,
    pub request_id: String,
    pub results: OCRResult,
    pub processing_time: Option<f64>,
    pub usage: Option<Usage>,
}

fn default_ocr_response_type() -> String {
    "OCRResponse".to_string()
}

// === Embedding Models ===

#[derive(Deserialize, Serialize, Default)]
pub struct EmbeddingRequest {
    #[serde(rename = "type")]
    #[serde(default = "default_embedding_request_type")]
    pub request_type: String,
    pub request_id: Option<String>,
    pub model: Option<String>,
    pub framework: Option<String>,
    pub text: Option<String>,
    pub image: Option<String>,
    pub user_id: Option<String>,
    pub organizations: Option<HashMap<String, HashMap<String, String>>>,
    pub handle: Option<String>,
    pub output_topic: Option<String>,
    pub output_partition: Option<i32>,
}

fn default_embedding_request_type() -> String {
    "EmbeddingRequest".to_string()
}

#[derive(Deserialize, Serialize, Default)]
pub struct Embedding {
    pub object: String,
    pub index: i32,
    pub embedding: Vec<f32>,
}

#[derive(Deserialize, Serialize, Default)]
pub struct EmbeddingResponse {
    #[serde(rename = "type")]
    #[serde(default = "default_embedding_response_type")]
    pub response_type: String,
    pub request_id: String,
    pub object: String,
    pub data: Vec<Embedding>,
    pub model: String,
    pub usage: Option<Usage>,
}

fn default_embedding_response_type() -> String {
    "EmbeddingResponse".to_string()
}

// === Auth ===

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
pub struct V1UserProfile {
    pub email: String,
    pub display_name: Option<String>,
    pub handle: Option<String>,
    pub picture: Option<String>,
    pub organization: Option<String>,
    pub role: Option<String>,
    pub external_id: Option<String>,
    pub actor: Option<String>,
    pub organizations: Option<HashMap<String, HashMap<String, String>>>,
    pub created: Option<i64>,
    pub updated: Option<i64>,
    pub token: Option<String>,
}

// === Training ===
#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct LlamaFactoryParams {
    pub model: String,
}

/// Enumeration of valid torch dtypes recognized by TRL.
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(rename_all = "lowercase")]
pub enum TorchDType {
    Auto,
    Bfloat16,
    Float16,
    Float32,
}

/// Utility enum for handling either a single string or a list of strings (per the "Union[str, list[str]]" spec).
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(untagged)]
pub enum OneOrManyStrings {
    Single(String),
    Multiple(Vec<String>),
}

#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct TRLParams {
    /// Model checkpoint for weights initialization.
    #[serde(default)]
    pub model_name_or_path: Option<String>,

    /// Specific model version to use.
    #[serde(default = "default_model_revision")]
    pub model_revision: String,

    /// Override the default torch.dtype and load the model under this dtype.
    /// Valid values: "auto", "bfloat16", "float16", "float32".
    #[serde(default)]
    pub torch_dtype: Option<TorchDType>,

    /// Whether to allow for custom models defined on the Hub in their own modeling files.
    #[serde(default)]
    pub trust_remote_code: bool,

    /// Which attention implementation to use (e.g. "flash_attention_2").
    #[serde(default)]
    pub attn_implementation: Option<String>,

    /// Whether to use PEFT for training.
    #[serde(default)]
    pub use_peft: bool,

    /// LoRA R value.
    #[serde(default = "default_lora_r")]
    pub lora_r: i32,

    /// LoRA alpha.
    #[serde(default = "default_lora_alpha")]
    pub lora_alpha: i32,

    /// LoRA dropout.
    #[serde(default = "default_lora_dropout")]
    pub lora_dropout: f32,

    /// LoRA target modules (Union[str, list[str]]).
    #[serde(default)]
    pub lora_target_modules: Option<OneOrManyStrings>,

    /// Model layers to unfreeze & train.
    #[serde(default)]
    pub lora_modules_to_save: Option<Vec<String>>,

    /// Task type to pass for LoRA (use "SEQ_CLS" for reward modeling).
    #[serde(default = "default_lora_task_type")]
    pub lora_task_type: String,

    /// Whether to use Rank-Stabilized LoRA.
    #[serde(default)]
    pub use_rslora: bool,

    /// Enable Weight-Decomposed Low-Rank Adaptation (DoRA).
    #[serde(default)]
    pub use_dora: bool,

    /// Whether to use 8-bit precision for the base model. Works only with LoRA.
    #[serde(default)]
    pub load_in_8bit: bool,

    /// Whether to use 4-bit precision for the base model. Works only with LoRA.
    #[serde(default)]
    pub load_in_4bit: bool,

    /// 4-bit quantization type ("fp4" or "nf4").
    #[serde(default = "default_bnb_4bit_quant_type")]
    pub bnb_4bit_quant_type: String,

    /// Whether to use nested quantization for 4-bit.
    #[serde(default)]
    pub use_bnb_nested_quant: bool,
}

impl Default for TRLParams {
    fn default() -> Self {
        Self {
            model_name_or_path: None,
            model_revision: default_model_revision(),
            torch_dtype: None,
            trust_remote_code: false,
            attn_implementation: None,
            use_peft: false,
            lora_r: default_lora_r(),
            lora_alpha: default_lora_alpha(),
            lora_dropout: default_lora_dropout(),
            lora_target_modules: None,
            lora_modules_to_save: None,
            lora_task_type: default_lora_task_type(),
            use_rslora: false,
            use_dora: false,
            load_in_8bit: false,
            load_in_4bit: false,
            bnb_4bit_quant_type: default_bnb_4bit_quant_type(),
            use_bnb_nested_quant: false,
        }
    }
}

fn default_model_revision() -> String {
    "main".to_string()
}

fn default_lora_r() -> i32 {
    16
}

fn default_lora_alpha() -> i32 {
    32
}

fn default_lora_dropout() -> f32 {
    0.05
}

fn default_lora_task_type() -> String {
    "CAUSAL_LM".to_string()
}

fn default_bnb_4bit_quant_type() -> String {
    "nf4".to_string()
}

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct MSSwiftParams {
    #[serde(default = "default_ms_swift_model")]
    pub model: String,

    #[serde(default = "default_ms_swift_model_type")]
    pub model_type: String,

    #[serde(default = "default_ms_swift_train_type")]
    pub train_type: String,

    #[serde(default = "default_ms_swift_deepspeed")]
    pub deepspeed: String,

    #[serde(default = "default_ms_swift_dtype")]
    pub torch_dtype: String,

    #[serde(default = "default_ms_swift_max_length")]
    pub max_length: i32,

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

    #[serde(default = "default_ms_swift_val_split_ratio")]
    pub val_split_ratio: f32,

    #[serde(default = "default_ms_swift_num_train_epochs")]
    pub num_train_epochs: i32,

    #[serde(default = "default_ms_swift_eval_strategy")]
    pub eval_strategy: String,

    #[serde(default = "default_ms_swift_save_strategy")]
    pub save_strategy: String,

    #[serde(default)]
    pub save_steps: Option<i32>,

    #[serde(default = "default_ms_swift_save_total_limit")]
    pub save_total_limit: i32,

    #[serde(default)]
    pub lora_rank: Option<i32>,

    #[serde(default)]
    pub lora_alpha: Option<i32>,

    #[serde(default = "default_ms_swift_size_factor")]
    pub size_factor: i32,

    #[serde(default = "default_ms_swift_max_pixels")]
    pub max_pixels: i32,

    #[serde(default)]
    pub resume_from_checkpoint: Option<String>,

    #[serde(default)]
    pub freeze_vit: Option<bool>,

    #[serde(default)]
    pub rlhf_type: Option<String>,

    #[serde(default = "default_ms_swift_gradient_accumulation_steps_total")]
    pub gradient_accumulation_steps_total: i32,

    #[serde(default)]
    pub learning_rate: Option<f32>,
}

// Add these helper functions at the bottom of the file:
fn default_ms_swift_model() -> String {
    "Qwen/Qwen2-VL-7B-Instruct".to_string()
}

fn default_ms_swift_model_type() -> String {
    "qwen2-vl-7b-instruct".to_string()
}

fn default_ms_swift_train_type() -> String {
    "lora".to_string()
}

fn default_ms_swift_deepspeed() -> String {
    "zero3".to_string()
}

fn default_ms_swift_dtype() -> String {
    "bfloat16".to_string()
}

fn default_ms_swift_max_length() -> i32 {
    8192
}

fn default_ms_swift_num_train_epochs() -> i32 {
    3
}

fn default_ms_swift_eval_strategy() -> String {
    "epoch".to_string()
}

fn default_ms_swift_save_strategy() -> String {
    "epoch".to_string()
}

fn default_ms_swift_save_total_limit() -> i32 {
    3
}

fn default_ms_swift_size_factor() -> i32 {
    28
}

fn default_ms_swift_max_pixels() -> i32 {
    1025000
}

// fn default_ms_swift_max_pixels() -> i32 {
//     802816
// }

fn default_ms_swift_val_split_ratio() -> f32 {
    0.90
}

fn default_ms_swift_gradient_accumulation_steps_total() -> i32 {
    16
}

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct TrainingRequest {
    pub name: Option<String>,
    pub namespace: Option<String>,
    pub framework: String,
    pub vram_request: Option<String>,
    pub accelerators: Option<Vec<String>>,
    pub cpu_request: Option<String>,
    pub trust_remote_code: Option<bool>,
    pub adapter: Option<String>,
    pub buffer: Option<String>,
    pub queue: Option<String>,
    pub ms_swift_params: Option<MSSwiftParams>,
    pub llama_factory_params: Option<LlamaFactoryParams>,
    pub trl_params: Option<TRLParams>,
    pub resume: Option<bool>,
    pub labels: Option<HashMap<String, String>>,
    pub platform: Option<String>,
}

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct Checkpoint {
    pub step: i32,
    pub trainer_state: Option<String>,
    pub args: Option<String>,
    pub adapter_config: Option<String>,
}

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct TrainingJob {
    pub id: String,
    pub name: String,
    pub namespace: String,
    pub framework: String,
    pub vram_request: Option<String>,
    pub accelerators: Option<Vec<String>>,
    pub cpu_request: Option<String>,
    pub trust_remote_code: Option<bool>,
    pub adapter: Option<String>,
    pub buffer: Option<String>,
    pub ms_swift_params: Option<MSSwiftParams>,
    pub llama_factory_params: Option<LlamaFactoryParams>,
    pub trl_params: Option<TRLParams>,
    pub resume: Option<bool>,
    pub status: String,
    pub checkpoints: Option<Vec<Checkpoint>>,
    pub queue: Option<String>,
    pub platform: Option<String>,
    pub labels: Option<HashMap<String, String>>,
    pub created: i64,
    pub updated: i64,
}

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct TRLBufferParams {
    #[serde(default)]
    pub model_name_or_path: Option<String>,

    #[serde(default = "default_model_revision")]
    pub model_revision: String,

    #[serde(default)]
    pub torch_dtype: Option<TorchDType>,

    #[serde(default)]
    pub trust_remote_code: bool,

    #[serde(default)]
    pub attn_implementation: Option<String>,

    #[serde(default)]
    pub use_peft: bool,

    #[serde(default = "default_lora_r")]
    pub lora_r: i32,

    #[serde(default = "default_lora_alpha")]
    pub lora_alpha: i32,

    #[serde(default = "default_lora_dropout")]
    pub lora_dropout: f32,

    #[serde(default)]
    pub lora_target_modules: Option<OneOrManyStrings>,

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

    #[serde(default = "default_lora_task_type")]
    pub lora_task_type: String,

    #[serde(default)]
    pub use_rslora: bool,

    #[serde(default)]
    pub use_dora: bool,

    #[serde(default)]
    pub load_in_8bit: bool,

    #[serde(default)]
    pub load_in_4bit: bool,

    #[serde(default = "default_bnb_4bit_quant_type")]
    pub bnb_4bit_quant_type: String,

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

#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct MSSwiftBufferParams {
    #[serde(default = "default_ms_swift_model")]
    pub model: String,

    #[serde(default = "default_ms_swift_model_type")]
    pub model_type: String,

    #[serde(default = "default_ms_swift_train_type")]
    pub train_type: String,

    #[serde(default = "default_ms_swift_deepspeed")]
    pub deepspeed: String,

    #[serde(default = "default_ms_swift_dtype")]
    pub torch_dtype: String,

    #[serde(default = "default_ms_swift_max_length")]
    pub max_length: i32,

    #[serde(default = "default_ms_swift_val_split_ratio")]
    pub val_split_ratio: f32,

    #[serde(default = "default_ms_swift_num_train_epochs")]
    pub num_train_epochs: i32,

    #[serde(default = "default_ms_swift_eval_strategy")]
    pub eval_strategy: String,

    #[serde(default = "default_ms_swift_save_strategy")]
    pub save_strategy: String,

    #[serde(default = "default_ms_swift_save_total_limit")]
    pub save_total_limit: i32,

    #[serde(default)]
    pub save_steps: Option<i32>,

    #[serde(default)]
    pub lora_rank: Option<i32>,

    #[serde(default)]
    pub lora_alpha: Option<i32>,

    #[serde(default = "default_ms_swift_size_factor")]
    pub size_factor: i32,

    #[serde(default = "default_ms_swift_max_pixels")]
    pub max_pixels: i32,

    #[serde(default)]
    pub resume_from_checkpoint: Option<String>,

    #[serde(default)]
    pub freeze_vit: Option<bool>,

    #[serde(default)]
    pub rlhf_type: Option<String>,

    #[serde(default = "default_ms_swift_gradient_accumulation_steps_total")]
    pub gradient_accumulation_steps_total: i32,

    #[serde(default)]
    pub learning_rate: Option<f32>,
}

#[derive(Deserialize, Serialize, Clone, Default)]
pub struct TrainingJobsResponse {
    pub jobs: Vec<TrainingJob>,
}

// #[derive(Deserialize, Serialize, Clone, Default, Debug)]
// pub struct ReplayBufferRequest {
//     pub name: String,
//     pub namespace: Option<String>,
//     pub framework: String,
//     pub vram_request: Option<String>,
//     pub accelerators: Option<Vec<String>>,
//     pub cpu_request: Option<String>,
//     pub trust_remote_code: Option<bool>,
//     pub adapter: Option<String>,
//     pub train_every: Option<i32>,

//     #[serde(default = "default_sample_n")]
//     pub sample_n: i32,

//     #[serde(default = "default_sample_strategy")]
//     pub sample_strategy: String,

//     pub queue: Option<String>,
//     pub ms_swift_params: Option<MSSwiftBufferParams>,
//     pub llama_factory_params: Option<LlamaFactoryParams>,
//     pub trl_params: Option<TRLBufferParams>,
//     pub platform: Option<String>,
//     pub labels: Option<HashMap<String, String>>,
// }

// #[derive(Deserialize, Serialize, Clone, Default, Debug)]
// pub struct UpdateReplayBufferRequest {
//     pub name: Option<String>,
//     pub namespace: Option<String>,
//     pub framework: Option<String>,
//     pub accelerators: Option<Vec<String>>,
//     pub vram_request: Option<String>,
//     pub cpu_request: Option<String>,
//     pub trust_remote_code: Option<bool>,
//     pub adapter: Option<String>,
//     pub train_every: Option<i32>,
//     pub sample_n: Option<i32>,
//     pub sample_strategy: Option<String>,
//     pub queue: Option<String>,
//     pub ms_swift_params: Option<MSSwiftBufferParams>,
//     pub llama_factory_params: Option<LlamaFactoryParams>,
//     pub trl_params: Option<TRLBufferParams>,
//     pub platform: Option<String>,
//     pub labels: Option<HashMap<String, String>>,
// }

// #[derive(Deserialize, Serialize, Clone, Default)]
// pub struct ReplayBuffer {
//     pub id: String,
//     pub name: String,
//     pub namespace: String,
//     pub framework: String,
//     pub vram_request: Option<String>,
//     pub accelerators: Option<Vec<String>>,
//     pub cpu_request: Option<String>,
//     pub trust_remote_code: Option<bool>,
//     pub adapter: Option<String>,
//     pub train_every: Option<i32>,
//     pub sample_n: i32,
//     pub sample_strategy: String,
//     pub ms_swift_params: Option<MSSwiftBufferParams>,
//     pub llama_factory_params: Option<LlamaFactoryParams>,
//     pub trl_params: Option<TRLBufferParams>,
//     pub labels: Option<HashMap<String, String>>,
//     pub queue: Option<String>,
//     pub platform: Option<String>,
//     pub num_records: Option<i32>,
//     pub train_idx: Option<i32>,
// }

// fn default_sample_n() -> i32 {
//     100
// }

// fn default_sample_strategy() -> String {
//     "Random".to_string()
// }

// #[derive(Deserialize, Serialize, Clone, Default)]
// pub struct ReplayBuffersResponse {
//     pub buffers: Vec<ReplayBuffer>,
// }

// #[derive(Deserialize, Serialize, Clone, Default)]
// pub struct ReplayBufferData {
//     pub examples: Vec<serde_json::Value>,
//     pub train: Option<bool>,
// }

#[derive(Deserialize, Serialize, Clone, Default)]
pub struct Adapters {
    pub adapters: Vec<String>,
}

#[derive(Deserialize, Serialize, Clone, Default)]
pub struct Models {
    pub models: Vec<String>,
}

#[derive(Deserialize, Serialize, Clone, Default)]
pub struct Datasets {
    pub datasets: Vec<String>,
}

#[derive(Deserialize, Serialize, Clone, Default)]
pub struct ModelFile {
    pub latest_checkpoint: Option<String>,
}

// === Common ===

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
pub struct MessageEnvelope<T> {
    pub user_profile: V1UserProfile,
    pub request_id: String,
    pub output_topic: String,
    pub payload: T,
}

#[derive(Debug, Serialize, Deserialize, Clone, Default)]
pub struct ErrorResponse {
    #[serde(rename = "type", default = "default_error_response_type")]
    pub response_type: String,
    pub request_id: String,
    pub error: String,
    pub traceback: Option<String>,
}

fn default_error_response_type() -> String {
    "ErrorResponse".to_string()
}

// Add these helper functions at the bottom of the file:
fn default_sentence_tf_model() -> String {
    "clip-ViT-B-32".to_string()
}

fn default_cuda_device() -> String {
    "cuda".to_string()
}

fn default_cuda_device_option() -> Option<String> {
    // New name
    Some("cuda".to_string())
}

fn default_tensor_parallel_size() -> Option<i32> {
    Some(1)
}

fn default_auto_dtype() -> Option<String> {
    Some("auto".to_string())
}

fn default_max_images() -> Option<i32> {
    Some(1)
}

fn default_det_arch() -> Option<String> {
    Some("fast_base".to_string())
}

fn default_reco_arch() -> Option<String> {
    Some("crnn_vgg16_bn".to_string())
}

fn default_option_true() -> Option<bool> {
    Some(true)
}

fn default_true() -> bool {
    true
}

fn default_option_false() -> Option<bool> {
    Some(false)
}

fn default_false() -> bool {
    false
}

fn default_lang_list() -> Option<Vec<String>> {
    Some(vec!["en".to_string()])
}

fn default_max_model_len() -> Option<i32> {
    Some(8192)
}

fn default_max_num_seqs() -> Option<i32> {
    Some(5)
}

fn default_gpu_memory_utilization() -> Option<f32> {
    Some(0.8)
}

fn default_enforce_eager() -> Option<bool> {
    Some(true)
}

fn default_enable_adapter() -> Option<bool> {
    Some(false)
}

// === LocalQueue ===

use schemars::JsonSchema;

#[derive(CustomResource, Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
#[kube(
    group = "kueue.x-k8s.io",
    version = "v1beta1",
    kind = "LocalQueue",
    namespaced
)]
#[serde(rename_all = "camelCase")]
pub struct LocalQueueSpec {
    pub cluster_queue: String,
    pub stop_policy: Option<String>,
    pub resource_groups: Option<Vec<ResourceGroup>>,
}

// Example auxiliary structs
#[derive(Serialize, Deserialize, Debug, Default, Clone, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ResourceGroup {
    pub covered_resources: Option<Vec<String>>,
    pub flavors: Option<Vec<Flavor>>,
}

#[derive(Serialize, Deserialize, Debug, Default, Clone, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct Flavor {
    pub name: String,
    pub resources: Option<Vec<FlavorResource>>,
}

#[derive(Serialize, Deserialize, Debug, Default, Clone, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FlavorResource {
    pub name: String,
    pub nominal_quota: String,
}

//
// === ClusterQueue ===
//

pub type ResourceName = String;

#[derive(CustomResource, Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
#[kube(
    group = "kueue.x-k8s.io",
    version = "v1beta1",
    kind = "ClusterQueue",
    namespaced = false
)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueueSpec {
    /// resourceGroups describes groups of resources.
    pub resource_groups: Vec<ClusterQueueResourceGroup>,

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

    /// cohort is now optional (skipped if None).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cohort: Option<String>,

    /// queueingStrategy now optional.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub queueing_strategy: Option<String>,

    /// flavorFungibility now optional.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub flavor_fungibility: Option<FlavorFungibility>,

    /// preemption now optional.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub preemption: Option<ClusterQueuePreemption>,

    /// admissionChecks lists the AdmissionChecks required by this ClusterQueue (cannot be used along with admissionChecksStrategy).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub admission_checks: Option<Vec<String>>,

    /// admissionChecksStrategy defines a list of strategies to determine which ResourceFlavors require AdmissionChecks.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub admission_checks_strategy: Option<AdmissionChecksStrategy>,

    /// stopPolicy now optional.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_policy: Option<String>,

    /// fairSharing now optional.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fair_sharing: Option<FairSharing>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct LabelSelectorRequirement {
    pub key: String,
    pub operator: String,

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

#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct CustomLabelSelector {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub match_labels: Option<std::collections::BTreeMap<String, String>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub match_expressions: Option<Vec<LabelSelectorRequirement>>,
}

/// ClusterQueueStatus defines the observed state of ClusterQueue.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueueStatus {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub flavors_reservation: Option<Vec<FlavorUsage>>,

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

    pub pending_workloads: i32,
    pub reserving_workloads: i32,
    pub admitted_workloads: i32,

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

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

/// A single resource group within a ClusterQueue, referencing flavors and their quotas.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueueResourceGroup {
    pub covered_resources: Vec<ResourceName>,
    pub flavors: Vec<FlavorQuotas>,
}

/// FlavorQuotas associates a flavor with resource quota definitions.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FlavorQuotas {
    pub name: String,
    pub resources: Vec<ResourceQuota>,
}

/// ResourceQuota defines nominal, borrowing, and lending quotas for a particular resource.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ResourceQuota {
    pub name: ResourceName,
    pub nominal_quota: String,

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

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

/// FlavorUsage tracks usage of a particular flavor and its resources.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FlavorUsage {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resources: Option<Vec<ResourceUsage>>,
}

/// ResourceUsage captures the total usage and how much was borrowed.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ResourceUsage {
    pub name: ResourceName,
    pub total: String,
    pub borrowed: String,
}

/// ClusterQueuePendingWorkloadsStatus is a listing of pending workloads in the cluster queue.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueuePendingWorkloadsStatus {
    #[serde(rename = "clusterQueuePendingWorkload")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cluster_queue_pending_workload: Option<Vec<ClusterQueuePendingWorkload>>,

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

/// ClusterQueuePendingWorkload identifies a single pending workload by name/namespace.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueuePendingWorkload {
    pub name: String,
    pub namespace: String,
}

/// Makes the entire block optional so it won't appear if None.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FairSharing {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub weight: Option<String>,
}

/// FairSharingStatus contains the current fair sharing usage details.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FairSharingStatus {
    pub weighted_share: i64,
}

/// Make the entire block optional so it won't appear if None.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct FlavorFungibility {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub when_can_borrow: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub when_can_preempt: Option<String>,
}

/// Make this optional as well if desired.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct ClusterQueuePreemption {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reclaim_within_cohort: Option<String>,

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

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

/// BorrowWithinCohort defines whether and how to preempt other CQs to reclaim resources while borrowing.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct BorrowWithinCohort {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub policy: Option<String>,

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

/// AdmissionCheckStrategyRule defines rules for exactly one AdmissionCheck
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct AdmissionCheckStrategyRule {
    pub name: String,

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

/// AdmissionChecksStrategy is a list of AdmissionCheck rules.
#[derive(Deserialize, Serialize, Debug, Clone, Default, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct AdmissionChecksStrategy {
    pub admission_checks: Vec<AdmissionCheckStrategyRule>,
}