orign 0.2.3

use clap::{arg, ArgAction, Args, Parser, Subcommand};

/// Orign CLI.
#[derive(Parser)]
#[command(author, version, about, long_about = None)]
pub struct Cli {
    #[command(subcommand)]
    pub command: Commands,
}

/// The subcommands supported by the CLI.
#[derive(Subcommand)]
pub enum Commands {
    /// Chat with a model.
    Chat {
        /// Model name.
        #[arg(short, long)]
        model: String,

        /// Message to send.
        #[arg(long)]
        msg: String,

        /// Image to send. Can provide multiple images.
        #[arg(short, long, action = ArgAction::Append)]
        image: Vec<String>,

        /// Framework of the model.
        #[arg(short, long)]
        framework: Option<String>,

        /// Adapter to use.
        #[arg(short, long)]
        adapter: Option<String>,
    },

    /// Create resources.
    Create {
        #[command(subcommand)]
        command: CreateCommands,
    },

    /// Get resources.
    Get {
        #[command(subcommand)]
        command: GetCommands,
    },

    /// Stop resources.
    Stop {
        #[command(subcommand)]
        command: StopCommands,
    },

    /// Delete resources.
    Delete {
        #[command(subcommand)]
        command: DeleteCommands,
    },

    /// Train a model.
    Train {
        #[command(subcommand)]
        command: TrainCommands,
    },

    /// Prepare a dataset.
    Prepare {
        /// Type of dataset.
        #[arg(short, long)]
        dataset_type: String,

        /// URL of the dataset.
        #[arg(short, long)]
        url: String,

        /// Split ratio.
        #[arg(short, long)]
        split_ratio: f64,

        /// Base path.
        #[arg(short, long)]
        base_path: Option<String>,

        /// Image path.
        #[arg(short, long)]
        image_path: Option<String>,

        /// Number of workers.
        #[arg(long, default_value = "8")]
        num_workers: usize,
    },

    /// Serve orign.
    Serve {
        /// The address to bind to.
        #[arg(long, default_value = "127.0.0.1")]
        host: String,

        /// The port to bind to.
        #[arg(short, long, default_value_t = 8080)]
        port: u16,
    },

    /// Report training progress.
    Report {
        /// Model name
        #[clap(long)]
        watch_dir: String,

        /// Save directory
        #[clap(long)]
        save_dir: String,

        /// Debounce duration in seconds
        #[clap(long, default_value_t = 2)]
        debounce_secs: u64,
    },

    /// Send examples to a buffer.
    Send {
        /// Name of the buffer.
        buffer_name: String,

        /// Path to the file with data to send.
        #[arg(short, long)]
        file: String,

        /// Whether to train the model.
        #[arg(long)]
        train: Option<bool>,
    },

    /// Trigger a buffer.
    Trigger {
        /// Buffer name.
        buffer_name: String,
    },

    /// Orign workers.
    Work {
        #[command(subcommand)]
        command: WorkCommands,
    },

    /// Login to orign.
    Login,
}

/// Create resources.
#[derive(Subcommand)]
pub enum CreateCommands {
    /// Create an replay buffer.
    Buffer {
        #[command(flatten)]
        command: ReplayBufferCommands,
    },
}

/// vLLM specific options.
#[derive(Args)]
pub struct VllmOptions {
    /// Name of the vLLM model to deploy.
    #[arg(long = "model")]
    pub model: String,

    /// Type of the vLLM model to deploy. Options are "qwen2_vl" and "molmo"
    #[arg(long = "model-type")]
    pub model_type: Option<String>,

    /// Trust remote code.
    #[arg(long, default_value_t = true)]
    pub trust_remote_code: bool,

    /// Tensor parallel size.
    #[arg(long, default_value_t = 1)]
    pub tensor_parallel_size: i32,

    /// Maximum images per prompt.
    #[arg(long, default_value_t = 1)]
    pub max_images_per_prompt: i32,

    /// Device to use (e.g., "cuda" or "cpu").
    #[arg(long, default_value = "cuda")]
    pub device: String,

    /// Maximum model length.
    #[arg(long, default_value_t = 8192)]
    pub max_model_len: i32,

    /// Maximum number of sequences.
    #[arg(long, default_value_t = 5)]
    pub max_num_seqs: i32,

    /// GPU memory utilization.
    #[arg(long, default_value_t = 0.8)]
    pub gpu_memory_utilization: f32,

    /// Enforce eager execution.
    #[arg(long, default_value_t = true)]
    pub enforce_eager: bool,

    /// Enable adapter.
    #[arg(long, default_value_t = false)]
    pub enable_adapter: bool,
}

/// EasyOCR specific options.
#[derive(Args)]
pub struct EasyOcrOptions {
    /// Device to use (e.g., "cuda" or "cpu").
    #[arg(long, default_value = "cuda")]
    pub device: String,

    /// Use GPU for inference.
    #[arg(long, default_value_t = true)]
    pub gpu: bool,

    /// List of languages (comma-separated).
    #[arg(long = "lang-list", value_delimiter = ',', default_value = "en")]
    pub lang_list: Vec<String>,

    /// Enable quantization.
    #[arg(long, default_value_t = false)]
    pub quantize: bool,
}

/// Doctr specific options.
#[derive(Args)]
pub struct DoctrOptions {
    /// Detection architecture.
    #[arg(long = "det-arch", default_value = "fast_base")]
    pub det_arch: String,

    /// Recognition architecture.
    #[arg(long = "reco-arch", default_value = "crnn_vgg16_bn")]
    pub reco_arch: String,

    /// Use pretrained weights.
    #[arg(long, default_value_t = true)]
    pub pretrained: bool,
}

/// Sentence-TF specific options.
#[derive(Args)]
pub struct SentenceTfOptions {
    /// Model name to deploy.
    #[arg(long = "model", default_value = "clip-ViT-B-32")]
    pub model: String,

    /// Device to use (e.g., "cuda" or "cpu").
    #[arg(long, default_value = "cuda")]
    pub device: String,
}

/// LiteLLM specific options.
#[derive(Args, Debug)]
pub struct LiteLLMOptions {
    /// API keys in format "env_var=key" e.g. "OPENAI_API_KEY=sk-...". Can be specified multiple times.
    #[arg(long = "api-key", value_parser = parse_key_val::<String, String>)]
    pub api_keys: Vec<(String, String)>,
}

// Helper function to parse key-value pairs
fn parse_key_val<K, V>(s: &str) -> Result<(K, V), String>
where
    K: std::str::FromStr,
    V: std::str::FromStr,
    <K as std::str::FromStr>::Err: std::fmt::Display,
    <V as std::str::FromStr>::Err: std::fmt::Display,
{
    let pos = s
        .find('=')
        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{}`", s))?;

    Ok((
        s[..pos]
            .parse()
            .map_err(|e| format!("invalid key: {}", e))?,
        s[pos + 1..]
            .parse()
            .map_err(|e| format!("invalid value: {}", e))?,
    ))
}

/// Common arguments for model deployment commands.
#[derive(Args)]
pub struct CommonModelArgs {
    /// Amount of VRAM to request.
    #[arg(long)]
    pub vram: Option<String>,

    /// Data type (e.g., "fp16", "bf16").
    #[arg(long)]
    pub dtype: Option<String>,

    /// Maximum pixels for images.
    #[arg(long)]
    pub max_pixels: Option<i32>,
}

#[derive(Subcommand)]
pub enum DeploymentCommands {
    /// Create a vLLM chat model.
    Vllm {
        #[command(flatten)]
        args: CommonModelArgs,

        /// vLLM specific options.
        #[command(flatten)]
        options: VllmOptions,
    },

    /// Create an EasyOCR model.
    Easyocr {
        #[command(flatten)]
        args: CommonModelArgs,

        /// EasyOCR specific options.
        #[command(flatten)]
        options: EasyOcrOptions,
    },

    /// Create a Doctr OCR model.
    Doctr {
        #[command(flatten)]
        args: CommonModelArgs,

        /// Doctr specific options.
        #[command(flatten)]
        options: DoctrOptions,
    },

    /// Create a Sentence Transformers embedding model.
    STF {
        #[command(flatten)]
        args: CommonModelArgs,

        /// Sentence-TF specific options.
        #[command(flatten)]
        options: SentenceTfOptions,
    },

    /// Create a LiteLLM model.
    Litellm {
        /// LiteLLM specific options.
        #[command(flatten)]
        options: LiteLLMOptions,
    },
}

/// Get resources.
#[derive(Subcommand)]
pub enum GetCommands {
    /// Get deployments.
    Deployments {
        /// Optional deployment ID to get detailed information
        id: Option<String>,
    },
    /// Get a model training.
    Trainings {
        /// Training ID.
        id: Option<String>,
    },
    /// Get a buffer.
    Buffers {
        /// Buffer namespace/name.
        name: Option<String>,
    },
    /// Get a model.
    Models {},
    /// Get a dataset.
    Datasets {},
    /// Get an adapter.
    Adapters {},
}

/// Delete resources.
#[derive(Subcommand)]
pub enum DeleteCommands {
    /// Delete a deployment.
    Deployment {
        /// ID.
        id: String,
    },

    /// Delete a buffer.
    Buffer {
        /// Buffer ID.
        name: String,
    },

    /// Delete an adapter.
    Adapter {
        /// Adapter ID.
        name: String,
    },
}

/// Stop a model training.
#[derive(Subcommand)]
pub enum StopCommands {
    /// Stop a model training.
    Training {
        /// Training ID.
        #[arg(short, long)]
        id: String,
    },
}

#[derive(Args)]
pub struct MSSwiftArgs {
    /// Model name.
    #[arg(long, default_value = "Qwen/Qwen2-VL-7B-Instruct")]
    pub model: String,

    /// Model type (e.g., "qwen2-vl-7b-instruct").
    #[arg(long, default_value = "qwen2-vl-7b-instruct")]
    pub model_type: String,

    /// Training type (e.g., "lora").
    #[arg(long, default_value = "lora")]
    pub train_type: String,

    /// Deepspeed configuration (e.g., "zero3").
    #[arg(long, default_value = "zero3")]
    pub deepspeed: String,

    /// Torch data type (e.g., "bfloat16").
    #[arg(long, default_value = "bfloat16")]
    pub torch_dtype: String,

    /// Maximum sequence length.
    #[arg(long, default_value_t = 8192)]
    pub max_length: i32,

    /// Dataset path or identifier. e.g. https://mydata.com/train.jsonl or orign://pbarker/foo/train.jsonl
    #[arg(long)]
    pub dataset: String,

    /// Validation split ratio.
    #[arg(long, default_value_t = 0.90)]
    pub val_split_ratio: f32,

    /// Number of training epochs.
    #[arg(long, default_value_t = 3)]
    pub num_train_epochs: i32,

    /// Evaluation strategy (e.g., "epoch").
    #[arg(long, default_value = "epoch")]
    pub eval_strategy: String,

    /// Save strategy (e.g., "epoch").
    #[arg(long, default_value = "epoch")]
    pub save_strategy: String,

    /// Save total limit.
    #[arg(long, default_value_t = 3)]
    pub save_total_limit: i32,

    /// LoRA rank.
    #[arg(long)]
    pub lora_rank: Option<i32>,

    /// LoRA alpha.
    #[arg(long)]
    pub lora_alpha: Option<i32>,

    /// Size factor for images.
    #[arg(long, default_value_t = 28)]
    pub size_factor: i32,

    /// Maximum pixels for images.
    #[arg(long, default_value_t = 802816)]
    pub max_pixels: i32,

    /// Freeze ViT.
    #[arg(long, default_value_t = false)]
    pub freeze_vit: bool,

    /// RLHF type.
    #[arg(long)]
    pub rlhf_type: Option<String>,

    /// Gradient accumulation steps total.
    #[arg(long, default_value_t = 16)]
    pub gradient_accumulation_steps_total: i32,

    /// Learning rate.
    #[arg(long)]
    pub learning_rate: Option<f32>,

    /// Save steps.
    #[arg(long)]
    pub save_steps: Option<i32>,
}

/// Arguments for the Swift training command.
#[derive(Args)]
pub struct MSSwiftCommands {
    /// train args
    #[command(flatten)]
    pub train_args: TrainArgs,

    /// ms-swift args
    #[command(flatten)]
    pub ms_swift_args: MSSwiftArgs,
}

#[derive(Args)]
pub struct TrainArgs {
    /// Name of the training job.
    #[arg(long)]
    pub name: Option<String>,

    /// Namespace of the training job.
    #[arg(long)]
    pub namespace: Option<String>,

    /// Amount of VRAM to request (e.g., "16Gi").
    #[arg(long)]
    pub vram: Option<String>,

    /// Accelerators.
    #[arg(long)]
    pub accelerators: Option<Vec<String>>,

    /// Amount of CPU to request (e.g., "4").
    #[arg(long)]
    pub cpu_request: Option<String>,

    /// Trust remote code.
    #[arg(long)]
    pub trust_remote_code: Option<bool>,

    /// Path to the adapter.
    #[arg(long)]
    pub adapter: Option<String>,

    /// Optional buffer.
    #[arg(long)]
    pub buffer: Option<String>,

    /// Resume training.
    #[arg(long)]
    pub resume: Option<bool>,

    /// Queue to place job in.
    #[arg(long)]
    pub queue: Option<String>,

    /// Platform to use.
    #[arg(long)]
    pub platform: Option<String>,
}

/// Training subcommands.
#[derive(Subcommand)]
pub enum TrainCommands {
    /// Train a model using Swift.
    Swift(MSSwiftCommands),
}

#[derive(Args)]
pub struct MSSwiftBufferArgs {
    /// Model name.
    #[arg(long, default_value = "Qwen/Qwen2-VL-7B-Instruct")]
    pub model: String,

    /// Model type (e.g., "qwen2-vl-7b-instruct").
    #[arg(long, default_value = "qwen2-vl-7b-instruct")]
    pub model_type: String,

    /// Training type (e.g., "lora").
    #[arg(long, default_value = "lora")]
    pub train_type: String,

    /// Deepspeed configuration (e.g., "zero3").
    #[arg(long, default_value = "zero3")]
    pub deepspeed: String,

    /// Torch data type (e.g., "bfloat16").
    #[arg(long, default_value = "bfloat16")]
    pub torch_dtype: String,

    /// Maximum sequence length.
    #[arg(long, default_value_t = 8192)]
    pub max_length: i32,

    /// Validation split ratio.
    #[arg(long, default_value_t = 0.90)]
    pub val_split_ratio: f32,

    /// Number of training epochs.
    #[arg(long, default_value_t = 3)]
    pub num_train_epochs: i32,

    /// Evaluation strategy (e.g., "epoch").
    #[arg(long, default_value = "epoch")]
    pub eval_strategy: String,

    /// Save strategy (e.g., "epoch").
    #[arg(long, default_value = "epoch")]
    pub save_strategy: String,

    /// Save total limit.
    #[arg(long, default_value_t = 3)]
    pub save_total_limit: i32,

    /// LoRA rank.
    #[arg(long)]
    pub lora_rank: Option<i32>,

    /// LoRA alpha.
    #[arg(long)]
    pub lora_alpha: Option<i32>,

    /// Size factor for images.
    #[arg(long, default_value_t = 28)]
    pub size_factor: i32,

    /// Maximum pixels for images.
    #[arg(long, default_value_t = 802816)]
    pub max_pixels: i32,

    /// Freeze ViT.
    #[arg(long, default_value_t = false)]
    pub freeze_vit: bool,

    /// RLHF type.
    #[arg(long)]
    pub rlhf_type: Option<String>,

    /// Gradient accumulation steps total.
    #[arg(long, default_value_t = 16)]
    pub gradient_accumulation_steps_total: i32,

    /// Learning rate.
    #[arg(long)]
    pub learning_rate: Option<f32>,

    /// Save steps.
    #[arg(long)]
    pub save_steps: Option<i32>,
}

/// Arguments for the Replay Buffer create command.
#[derive(Args)]
pub struct ReplayBufferCommands {
    /// Name of the buffer.
    #[arg(long)]
    pub name: String,

    /// Namespace of the buffer.
    #[arg(long)]
    pub namespace: Option<String>,

    /// Train every N seconds
    #[arg(long)]
    pub train_every: Option<i32>,

    /// Sample N examples for training.
    #[arg(long, default_value_t = 100)]
    pub sample_n: i32,

    /// Sample strategy (e.g., "Random").
    #[arg(long, default_value = "Random")]
    pub sample_strategy: String,

    /// Image to use for training.
    #[arg(long)]
    pub image: String,

    /// Command to use for training.
    #[arg(long)]
    pub command: Option<String>,

    /// Accelerators.
    #[arg(long)]
    pub accelerators: Option<Vec<String>>,

    /// Number of epochs to train.
    #[arg(long, default_value_t = 1)]
    pub num_epochs: i32,
}

/// Subcommands for the "work" command
#[derive(Subcommand)]
pub enum WorkCommands {
    /// Report trainings by fetching K8s pods labeled "type=trainingjob".
    ReportTrainings {
        /// Kubernetes namespace in which to look for training job pods.
        #[arg(long, default_value = "default")]
        k8s_namespace: String,
    },
}