use clap::{Parser, Subcommand};

#[derive(Parser)]
#[command(name = "iperf", version, about = "AI inference backend performance benchmark")]
pub struct Cli {
    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Run a benchmark
    Run(RunArgs),
    /// Watch GPU and inference metrics
    Watch(WatchArgs),
    /// Generate default config.yaml
    Config(ConfigArgs),
    /// Model download and serving
    Hub(HubArgs),
}

/// All fields are Option so we can distinguish "user set" from "default"
/// when merging with config.yaml
#[derive(clap::Args, Clone)]
pub struct RunArgs {
    /// Path to config.yaml
    #[arg(long)]
    pub conf: Option<String>,
    /// Backend type (vllm, sglang)
    #[arg(long, short = 'b')]
    pub backend: Option<String>,
    /// Model name to benchmark
    #[arg(long, short = 'm')]
    pub model: Option<String>,
    /// Number of concurrent workers
    #[arg(long, short = 'c')]
    pub concurrency: Option<usize>,
    /// Benchmark duration (e.g. "60s", "5m", "1h")
    #[arg(long, short = 'd')]
    pub duration: Option<String>,
    /// Max requests (0 = unlimited)
    #[arg(long)]
    pub request_count: Option<usize>,
    /// Request mode: single, stream, continuous
    #[arg(long, short = 'M')]
    pub mode: Option<String>,
    /// Input prompt tokens
    #[arg(long)]
    pub prompt_tokens: Option<usize>,
    /// Max output tokens
    #[arg(long)]
    pub output_tokens: Option<usize>,
    /// Prepend UUID to each request (0% cache hit, overrides cache-rate)
    #[arg(long)]
    pub no_cache: bool,
    /// Prompt pool size
    #[arg(long)]
    pub num_prefix_prompts: Option<usize>,
    /// Cache hit percentage (0=fully unique, 100=fully cacheable)
    #[arg(long)]
    pub cache_rate: Option<usize>,
    /// Random seed
    #[arg(long)]
    pub seed: Option<i64>,
    /// Prompt length standard deviation
    #[arg(long)]
    pub prompt_tokens_stddev: Option<usize>,
    /// System prompt length in tokens (0 = disabled)
    #[arg(long)]
    pub system_prompt_tokens: Option<usize>,
    /// Number of unique system prompts (pool size)
    #[arg(long)]
    pub num_system_prompts: Option<usize>,
    /// Output format: table, json
    #[arg(long, short = 'f')]
    pub format: Option<String>,
    /// JSONL output directory
    #[arg(long)]
    pub output_dir: Option<String>,
    /// HTTP proxy URL
    #[arg(long)]
    pub http_proxy: Option<String>,
    /// Print a copy-pasteable curl command and exit (no benchmark).
    /// Optional value: show the Nth request (default: 1)
    #[arg(long, num_args = 0..=1, default_missing_value = "1")]
    pub trace: Option<usize>,
    /// Mark this run as warmup (output tagged with warmup: true)
    #[arg(long)]
    pub warmup: bool,
    /// Tag for results
    #[arg(long)]
    pub tag: Option<String>,
    /// Target base URL (positional)
    #[arg(index = 1)]
    pub target: Option<String>,
}

#[derive(clap::Args)]
pub struct WatchArgs {
    /// Target base URL (e.g., http://localhost:8000/v1)
    #[arg(index = 1)]
    pub target: Option<String>,
    /// Sampling interval in seconds
    #[arg(long, short = 'i', default_value_t = 2)]
    pub interval: u64,
    /// Enable nsys GPU profiling
    #[arg(long)]
    pub nsys: bool,
    /// Stop after duration (e.g., "60s", "5m", "1h")
    #[arg(long, short = 'd')]
    pub duration: Option<String>,
    /// Backend type (vllm, sglang) — affects metric name prefixes
    #[arg(long, short = 'b')]
    pub backend: Option<String>,
    /// Model name (used for JSONL filename)
    #[arg(long, short = 'm')]
    pub model: Option<String>,
    /// Tag for JSONL filename
    #[arg(long)]
    pub tag: Option<String>,
    /// JSONL output directory
    #[arg(long)]
    pub output_dir: Option<String>,
}

#[derive(clap::Args)]
pub struct ConfigArgs {
    /// Output file path
    #[arg(long, short = 'o', default_value = "config.yaml")]
    pub output: String,
}

#[derive(clap::Args)]
pub struct HubArgs {
    #[command(subcommand)]
    pub command: HubCommands,
}

#[derive(Subcommand)]
pub enum HubCommands {
    /// Download model from HuggingFace or custom hub
    Download(HubDownloadArgs),
    /// Serve models via HTTP
    Serve(HubServeArgs),
}

#[derive(clap::Args)]
pub struct HubDownloadArgs {
    /// HuggingFace model ID (e.g. "meta-llama/Llama-3-8B")
    pub model_id: String,
    /// Local directory
    #[arg(long)]
    pub local_dir: Option<String>,
    /// Branch/revision (HuggingFace only)
    #[arg(long, short = 'r', default_value = "main")]
    pub revision: String,
    /// Custom hub URL
    #[arg(long)]
    pub source: Option<String>,
    /// HTTP proxy
    #[arg(long)]
    pub http_proxy: Option<String>,
    /// Skip first N files
    #[arg(long, default_value_t = 0)]
    pub offset: usize,
    /// Download N files (0 = all)
    #[arg(long, default_value_t = 0)]
    pub count: usize,
    /// Role for distributed download (e.g. "1/4", "2/4")
    #[arg(long)]
    pub role: Option<String>,
}

#[derive(clap::Args)]
pub struct HubServeArgs {
    /// Models directory
    #[arg(long)]
    pub local_dir: String,
    /// Listen address
    #[arg(long, short = 'a', default_value = "0.0.0.0:8080")]
    pub addr: String,
}
