Function bodies 96 total

fn default_log_level() -> LogLevel {
    LogLevel::Info
}

fn default_metrics_interval() -> String {
    "1m".to_string()
}

fn default_admin_listen() -> String {
    "127.0.0.1:9090".to_string()
}

fn default_admin_enabled() -> bool {
    true
}

    pub fn load(path: &PathBuf) -> anyhow::Result<Self> {
        let contents = std::fs::read_to_string(path)?;
        let config: Config = toml::from_str(&contents)?;
        config.validate()?;
        Ok(config)
    }

    pub fn validate(&self) -> anyhow::Result<()> {
        if self.load.total_requests.is_none() && self.load.duration_seconds.is_none() {
            anyhow::bail!("Either total_requests or duration_seconds must be specified");
        }

        if self.load.total_requests.is_some() && self.load.duration_seconds.is_some() {
            anyhow::bail!("Only one of total_requests or duration_seconds can be specified");
        }

        if self.load.concurrent_requests == 0 {
            anyhow::bail!("concurrent_requests must be greater than 0");
        }

        // If qps is specified, we're in fixed QPS mode
        if let Some(qps) = self.load.qps
            && qps <= 0.0
        {
            anyhow::bail!("qps must be greater than 0");
        }

        if self.runtime.worker_threads == 0 {
            anyhow::bail!("worker_threads must be greater than 0");
        }

        Ok(())
    }

    pub fn new(arrival_dist: &ArrivalDistribution, qps: f64) -> Self {
        let dist_type = match arrival_dist {
            ArrivalDistribution::Uniform => {
                let interval_ms = (1000.0 / qps) as u64;
                DistributionType::Uniform {
                    interval: Duration::from_millis(interval_ms),
                }
            }
            ArrivalDistribution::Poisson => {
                // For Poisson arrivals, inter-arrival times follow exponential distribution
                // λ (lambda) = rate = qps
                let exp_dist =
                    Exp::new(qps).expect("QPS must be positive for exponential distribution");
                DistributionType::Exponential { exp_dist }
            }
        };

        Self { dist_type }
    }

    pub fn next_delay(&self) -> Duration {
        match &self.dist_type {
            DistributionType::Uniform { interval } => *interval,
            DistributionType::Exponential { exp_dist } => {
                let mut rng = thread_rng();
                // Sample returns time in seconds
                let wait_secs = exp_dist.sample(&mut rng);
                Duration::from_secs_f64(wait_secs)
            }
        }
    }

    pub fn distribution_name(&self) -> &str {
        match &self.dist_type {
            DistributionType::Uniform { .. } => "Uniform",
            DistributionType::Exponential { .. } => "Poisson",
        }
    }

    fn test_uniform_distribution() {
        let dist = RequestDistribution::new(&ArrivalDistribution::Uniform, 10.0);

        // For 10 QPS, expect 100ms intervals
        let delay = dist.next_delay();
        assert_eq!(delay, Duration::from_millis(100));

        // Should be deterministic
        let delay2 = dist.next_delay();
        assert_eq!(delay, delay2);
    }

    fn test_poisson_distribution_variability() {
        let dist = RequestDistribution::new(&ArrivalDistribution::Poisson, 10.0);

        // Collect several samples
        let mut delays = Vec::new();
        for _ in 0..100 {
            delays.push(dist.next_delay());
        }

        // Poisson should produce variable delays
        let all_same = delays.iter().all(|d| *d == delays[0]);
        assert!(
            !all_same,
            "Poisson distribution should produce variable delays"
        );

        // Average should be roughly 1/rate = 1/10 = 0.1 seconds
        let avg_secs: f64 =
            delays.iter().map(|d| d.as_secs_f64()).sum::<f64>() / delays.len() as f64;
        assert!(
            (avg_secs - 0.1).abs() < 0.05,
            "Average delay should be close to 1/rate (0.1s), got {}",
            avg_secs
        );
    }

    fn test_distribution_name() {
        let uniform = RequestDistribution::new(&ArrivalDistribution::Uniform, 10.0);
        assert_eq!(uniform.distribution_name(), "Uniform");

        let poisson = RequestDistribution::new(&ArrivalDistribution::Poisson, 10.0);
        assert_eq!(poisson.distribution_name(), "Poisson");
    }

fn parse_log_filters(filters: &[String]) -> HashMap<String, LevelFilter> {
    let mut map = HashMap::new();
    for filter in filters {
        if let Some((module, level)) = filter.split_once('=') {
            let level_filter = match level.to_lowercase().as_str() {
                "error" => LevelFilter::Error,
                "warn" => LevelFilter::Warn,
                "info" => LevelFilter::Info,
                "debug" => LevelFilter::Debug,
                "trace" => LevelFilter::Trace,
                "off" => LevelFilter::Off,
                _ => continue,
            };
            map.insert(module.to_string(), level_filter);
        }
    }
    map
}

fn should_log(metadata: &Metadata, filters: &HashMap<String, LevelFilter>) -> bool {
    let target = metadata.target();

    // Check each filter to see if it matches this target
    for (module_prefix, level_filter) in filters {
        if target.starts_with(module_prefix) {
            return metadata.level() <= *level_filter;
        }
    }

    // If no filter matched, allow the log (will be caught by global level filter)
    true
}

    fn enabled(&self, metadata: &Metadata) -> bool {
        metadata.level() <= self.max_level && should_log(metadata, &self.filters)
    }

    fn log(&self, record: &Record) {
        if self.enabled(record.metadata())
            && let Ok(mut output) = self.output.lock()
        {
            let message = format!("{}\n", record.args());
            let _ = output.write_all(message.as_bytes());
        }
    }

    fn flush(&self) {
        if let Ok(mut output) = self.output.lock() {
            let _ = output.flush();
        }
    }

fn main() -> Result<()> {
    // Parse CLI arguments
    let cli = Cli::parse_args();

    // Load configuration first to check verbosity setting
    let config = Config::load(&cli.config)?;

    // Set up logging with ringlog and per-module filtering
    let log_level = config.log.level.to_level_filter();

    // Configure output destination
    let output: Box<dyn Output> = if let Some(ref log_file) = config.output.trace_log {
        // Log to file with rotation
        let backup_file = log_file.with_extension("old");
        Box::new(File::new(log_file.clone(), backup_file, LOG_FILE_MAX_SIZE)?)
    } else {
        // Log to stderr
        Box::new(Stderr::new())
    };

    // Parse per-module filters from config
    let filters = parse_log_filters(&config.log.filter);

    // Create logger with per-module filtering if configured
    if filters.is_empty() {
        // No filters - use ringlog directly
        let base_log = LogBuilder::new()
            .output(output)
         

async fn run_benchmark(config: Config) -> Result<()> {
    // Start admin server if configured
    if let Some(ref admin_config) = config.admin
        && admin_config.enabled
    {
        let addr: std::net::SocketAddr = admin_config
            .listen
            .parse()
            .expect("Invalid admin listen address");

        info!("Starting metrics server on {}", addr);
        tokio::spawn(async move {
            llm_bench::admin::start_server(addr).await;
        });
    }

    debug!("Initializing benchmark runner");
    let runner = llm_bench::BenchmarkRunner::new(config).await?;
    info!("Starting benchmark run");
    runner.run().await?;
    info!("Benchmark completed successfully");
    Ok(())
}

    pub fn init() {
        // Metriken metrics are automatically registered via the #[metric] attribute
        // No explicit initialization needed
    }

    pub fn record_request_sent() {
        REQUESTS_SENT.increment();
        REQUESTS_INFLIGHT.increment();
    }

    pub fn record_request_complete(status: RequestStatus) {
        REQUESTS_INFLIGHT.decrement();
        match status {
            RequestStatus::Success => {
                REQUESTS_SUCCESS.increment();
            }
            RequestStatus::Failed(error_type) => {
                REQUESTS_FAILED.increment();
                match error_type {
                    ErrorType::Connection => ERRORS_CONNECTION.increment(),
                    ErrorType::Http4xx(_) => ERRORS_HTTP_4XX.increment(),
                    ErrorType::Http5xx(_) => ERRORS_HTTP_5XX.increment(),
                    ErrorType::Parse => ERRORS_PARSE.increment(),
                    ErrorType::Timeout => REQUESTS_TIMEOUT.increment(),
                    ErrorType::Other => ERRORS_OTHER.increment(),
                };
            }
            RequestStatus::Timeout => {
                REQUESTS_TIMEOUT.increment();
                ERRORS_OTHER.increment();
            }
        }
    }

    pub fn record_tokens(input: u64, output: u64) {
        TOKENS_INPUT.add(input);
        TOKENS_OUTPUT.add(output);
    }

    pub fn record_ttft_with_context(duration: Duration, input_tokens: u64) {
        let nanos = duration.as_nanos() as u64;

        // Record in context-specific histogram based on production patterns
        match input_tokens {
            0..=200 => {
                let _ = TTFT_SMALL.increment(nanos); // ~50% of production traffic
            }
            201..=500 => {
                let _ = TTFT_MEDIUM.increment(nanos); // ~30% of production traffic
            }
            501..=2000 => {
                let _ = TTFT_LARGE.increment(nanos); // ~15% of production traffic
            }
            2001..=8000 => {
                let _ = TTFT_XLARGE.increment(nanos); // ~4% of production traffic
            }
            _ => {
                let _ = TTFT_XXLARGE.increment(nanos); // ~1% of production traffic
            }
        }
    }

    pub fn record_tpot(duration: Duration) {
        let _ = TPOT.increment(duration.as_nanos() as u64);
    }

    pub fn record_itl_with_context(duration: Duration, input_tokens: u64) {
        let nanos = duration.as_nanos() as u64;

        // Record in context-specific histogram
        match input_tokens {
            0..=200 => {
                let _ = ITL_SMALL.increment(nanos);
            }
            201..=500 => {
                let _ = ITL_MEDIUM.increment(nanos);
            }
            501..=1000 => {
                let _ = ITL_LARGE.increment(nanos);
            }
            1001..=2000 => {
                let _ = ITL_XLARGE.increment(nanos);
            }
            _ => {
                let _ = ITL_XXLARGE.increment(nanos);
            }
        }
    }

    pub fn record_latency(duration: Duration) {
        let _ = REQUEST_LATENCY.increment(duration.as_nanos() as u64);
    }

    pub fn record_retry() {
        REQUESTS_RETRIED.increment();
    }

    fn default() -> Self {
        Self::new()
    }

    pub fn new() -> Self {
        Self {
            start_time: SystemTime::now(),
            config: None,
            duration: None,
        }
    }

    pub fn with_config(mut self, config: crate::config::Config) -> Self {
        self.config = Some(config);
        self
    }

    pub fn with_duration(mut self, duration: Duration) -> Self {
        self.duration = Some(duration);
        self
    }

    pub fn build(&self) -> Result<BenchmarkReport> {
        // Use provided duration if available, otherwise calculate from elapsed time
        let duration = if let Some(d) = self.duration {
            d
        } else {
            let end_time = SystemTime::now();
            end_time.duration_since(self.start_time)?
        };

        // Access metrics directly for now
        // In production, we'd use metriken-exposition properly
        use crate::metrics::{
            ERRORS_CONNECTION, ERRORS_HTTP_4XX, ERRORS_HTTP_5XX, ERRORS_OTHER, REQUESTS_FAILED,
            REQUESTS_RETRIED, REQUESTS_SENT, REQUESTS_SUCCESS, REQUESTS_TIMEOUT, TOKENS_INPUT,
            TOKENS_OUTPUT,
        };

        let requests_sent = REQUESTS_SENT.value();
        let requests_success = REQUESTS_SUCCESS.value();
        let requests_failed = REQUESTS_FAILED.value();
        let requests_timeout = REQUESTS_TIMEOUT.value();

        let input_tokens = TOKENS_INPUT.value();
        let output_tokens

    fn calculate_histogram_mean(histogram: &metriken::AtomicHistogram) -> f64 {
        if let Some(loaded) = histogram.load() {
            let mut sum = 0u64;
            let mut count = 0u64;

            // Iterate through all buckets
            for bucket in loaded.iter() {
                let bucket_count = bucket.count();
                if bucket_count > 0 {
                    // Use upper edge for consistency with percentile calculation
                    sum += bucket.end() * bucket_count;
                    count += bucket_count;
                }
            }

            if count > 0 {
                return (sum as f64 / count as f64) / 1_000_000.0; // Convert to ms
            }
        }
        0.0
    }

    fn build_latency_stats(&self) -> Result<LatencyStats> {
        use crate::metrics::{REQUEST_LATENCY, TPOT};

        // Calculate means from histograms
        let tpot_mean = Self::calculate_histogram_mean(&TPOT);
        let request_mean = Self::calculate_histogram_mean(&REQUEST_LATENCY);

        // Extract TPOT percentiles
        let mut tpot_p50 = 0.0;
        let mut tpot_p90 = 0.0;
        let mut tpot_p95 = 0.0;
        let mut tpot_p99 = 0.0;

        if let Some(tpot_histogram) = TPOT.load()
            && let Ok(Some(percentiles)) = tpot_histogram.percentiles(&[50.0, 90.0, 95.0, 99.0])
        {
            for (percentile, bucket) in percentiles.iter() {
                let value_ms = bucket.end() as f64 / 1_000_000.0;
                match percentile.round() as u32 {
                    50 => tpot_p50 = value_ms,
                    90 => tpot_p90 = value_ms,
                    95 => tpot_p95 = value_ms,
                    99 => tpot_p99 = value_ms,
              

    fn build_context_latency_stats(&self) -> Option<ContextLatencyStats> {
        use crate::metrics::{TTFT_LARGE, TTFT_MEDIUM, TTFT_SMALL, TTFT_XLARGE, TTFT_XXLARGE};

        let extract_percentiles =
            |histogram: &metriken::AtomicHistogram| -> Option<LatencyPercentiles> {
                if let Some(loaded) = histogram.load()
                    && let Ok(Some(percentiles)) = loaded.percentiles(&[50.0, 90.0, 95.0, 99.0])
                    && !percentiles.is_empty()
                {
                    let mut p50 = 0.0;
                    let mut p90 = 0.0;
                    let mut p95 = 0.0;
                    let mut p99 = 0.0;

                    for (percentile, bucket) in percentiles.iter() {
                        let value_ms = bucket.end() as f64 / 1_000_000.0;
                        match percentile.round() as u32 {
                            50 => p50 = value_ms,
                            90 => p90 = value_ms,
                            95 => p9

    fn build_context_itl_stats(&self) -> Option<ContextITLStats> {
        use crate::metrics::{ITL_LARGE, ITL_MEDIUM, ITL_SMALL, ITL_XLARGE, ITL_XXLARGE};

        let extract_percentiles =
            |histogram: &metriken::AtomicHistogram| -> Option<ITLPercentiles> {
                if let Some(loaded) = histogram.load()
                    && let Ok(Some(percentiles)) = loaded.percentiles(&[50.0, 90.0, 95.0, 99.0])
                    && !percentiles.is_empty()
                {
                    let mut p50 = 0.0;
                    let mut p90 = 0.0;
                    let mut p95 = 0.0;
                    let mut p99 = 0.0;

                    for (percentile, bucket) in percentiles.iter() {
                        let value_ms = bucket.end() as f64 / 1_000_000.0;
                        match percentile.round() as u32 {
                            50 => p50 = value_ms,
                            90 => p90 = value_ms,
                            95 => p95 = value_ms,
   

    pub fn print_console_report(&self) -> Result<()> {
        let report = self.build()?;

        use chrono::Utc;
        let now = Utc::now();
        let timestamp = now.to_rfc3339_opts(chrono::SecondsFormat::Millis, false);

        println!();
        println!("{}", timestamp);
        println!("{} -----", timestamp);
        println!("{} Benchmark Complete", timestamp);
        println!(
            "{} Duration: {:.1}s",
            timestamp,
            report.duration.as_secs_f64()
        );
        println!(
            "{} Requests: Sent: {} Retries: {}",
            timestamp, report.summary.total_requests, report.summary.retries
        );
        println!(
            "{} Responses: Received: {} Ok: {} Err: {} Success: {:.2}%",
            timestamp,
            report.summary.successful_requests + report.summary.failed_requests,
            report.summary.successful_requests,
            report.summary.failed_requests,
            report.summary.success_rate * 100.0

pub async fn capture_snapshots(config: Config) -> Result<()> {
    let metrics_config = match config.metrics.as_ref() {
        Some(cfg) => cfg,
        None => return Ok(()), // No metrics configured
    };

    // Parse interval
    let interval_duration = humantime::parse_duration(&metrics_config.interval)?;

    // Create a temporary file for msgpack data (will be converted to parquet)
    let temp_file = NamedTempFile::new()?;
    let file = File::from_std(temp_file.reopen()?);
    let mut writer = BufWriter::new(file);

    // Get an aligned start time (aligned to the second)
    let start = Instant::now() - Duration::from_nanos(Utc::now().nanosecond() as u64)
        + Duration::from_secs(1);

    // Create interval timer
    let mut interval = interval_at(start, interval_duration);

    // Build snapshotter with metadata
    let snapshotter = SnapshotterBuilder::new()
        .metadata("source".to_string(), env!("CARGO_PKG_NAME").to_string())
        .metadata("version".to_str

    fn new() -> Self {
        Self {
            requests_sent: REQUESTS_SENT.value(),
            requests_success: REQUESTS_SUCCESS.value(),
            requests_failed: REQUESTS_FAILED.value(),
            requests_timeout: REQUESTS_TIMEOUT.value(),
            tokens_input: TOKENS_INPUT.value(),
            tokens_output: TOKENS_OUTPUT.value(),
            errors_connection: ERRORS_CONNECTION.value(),
            errors_4xx: ERRORS_HTTP_4XX.value(),
            errors_5xx: ERRORS_HTTP_5XX.value(),
            errors_parse: ERRORS_PARSE.value(),
            errors_other: ERRORS_OTHER.value(),
            tpot_histogram: TPOT.load(),
            request_histogram: REQUEST_LATENCY.load(),
        }
    }

    fn update(&mut self) {
        self.requests_sent = REQUESTS_SENT.value();
        self.requests_success = REQUESTS_SUCCESS.value();
        self.requests_failed = REQUESTS_FAILED.value();
        self.requests_timeout = REQUESTS_TIMEOUT.value();
        self.tokens_input = TOKENS_INPUT.value();
        self.tokens_output = TOKENS_OUTPUT.value();
        self.errors_connection = ERRORS_CONNECTION.value();
        self.errors_4xx = ERRORS_HTTP_4XX.value();
        self.errors_5xx = ERRORS_HTTP_5XX.value();
        self.errors_parse = ERRORS_PARSE.value();
        self.errors_other = ERRORS_OTHER.value();
        self.tpot_histogram = TPOT.load();
        self.request_histogram = REQUEST_LATENCY.load();
    }

pub async fn periodic_stats(config: Config, warmup_complete: Arc<Notify>) {
    // Default to 1 minute interval if not specified
    let interval_duration = if let Some(metrics_config) = config.metrics.as_ref() {
        humantime::parse_duration(&metrics_config.interval).unwrap_or(Duration::from_secs(60))
    } else {
        Duration::from_secs(60)
    };

    // Build snapshotter for reading metrics (not used but kept for compatibility)
    let snapshotter = SnapshotterBuilder::new().build();

    // Wait for warmup to complete before starting stats intervals
    warmup_complete.notified().await;

    // Get an aligned start time (aligned to the second) AFTER warmup
    let start = Instant::now() - Duration::from_nanos(Utc::now().nanosecond() as u64)
        + Duration::from_secs(1);

    // Create interval timer
    let mut interval = interval_at(start, interval_duration);
    let mut window_id = 0;

    // Wait a bit for initial metrics
    tokio::time::sleep(Duration::from_secs(

    pub fn new(model: &str) -> Result<Self> {
        let (encoder, model_type) = if model.contains("gpt-4o") {
            (Arc::new(o200k_base()?), ModelType::O200k)
        } else {
            // Default to cl100k for most models (GPT-3.5, GPT-4, etc.)
            (Arc::new(cl100k_base()?), ModelType::Cl100k)
        };

        Ok(Self {
            encoder,
            model_type,
        })
    }

    pub fn count_tokens(&self, text: &str) -> usize {
        // Note: This counts tokens in the raw text content only.
        // When using chat APIs, additional tokens are added for:
        // - Chat format markers (e.g., <|im_start|>, <|im_end|>)
        // - Role indicators (e.g., "user", "assistant")
        // - Other protocol overhead
        // So the actual tokens sent to the API will be higher than this count.
        self.encoder.encode_with_special_tokens(text).len()
    }

    pub fn model_type(&self) -> ModelType {
        self.model_type
    }

    fn test_token_counting() {
        let tokenizer = Tokenizer::new("gpt-3.5-turbo").unwrap();

        // Test basic text
        let count = tokenizer.count_tokens("Hello, world!");
        assert!(count > 0);

        // Test that whitespace counting differs from token counting
        let text = "This is a test";
        let word_count = text.split_whitespace().count();
        let token_count = tokenizer.count_tokens(text);
        // Tokens and words are usually different
        println!("Words: {}, Tokens: {}", word_count, token_count);
    }