// src/mp.rs
//
// SPDX-License-Identifier: Apache-2.0 OR MIT
// SPDX-FileCopyrightText: 2025 Russ Fellows <russ.fellows@gmail.com>

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::hash::{BuildHasher, BuildHasherDefault};
use std::collections::hash_map::DefaultHasher;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::time::Instant;

/// Configuration for multi-process GET operations
#[derive(Debug, Clone)]
pub struct MpGetConfig {
    /// Number of worker processes to spawn
    pub procs: usize,
    /// Concurrent operations per worker process
    pub concurrent_per_proc: usize,
    /// Path to file containing object keys (one per line)
    pub keylist: PathBuf,
    /// Optional timeout in seconds
    pub duration_secs: Option<u64>,
    /// Directory for per-worker operation logs
    pub op_log_dir: Option<PathBuf>,
    /// Worker command (default: "s3-cli")
    pub worker_cmd: String,
    /// Worker subcommand (default: "get")
    pub worker_subcmd: String,
    /// Extra environment variables for workers
    pub extra_env: Vec<(String, String)>,
    /// Tokio runtime threads per worker
    pub rt_threads: Option<usize>,
    /// Max HTTP connections per worker
    pub max_http_connections: Option<usize>,
    /// Enable optimized HTTP settings
    pub optimized_http: bool,
    /// Arguments to forward to each worker
    pub forward_args: Vec<String>,
    /// Whether to pass through worker stdout/stderr (for debugging)
    pub passthrough_io: bool,
}

impl Default for MpGetConfig {
    fn default() -> Self {
        Self {
            procs: 4,
            concurrent_per_proc: 64,
            keylist: PathBuf::from("keys.txt"),
            duration_secs: None,
            op_log_dir: None,
            worker_cmd: "s3-cli".to_string(),
            worker_subcmd: "get".to_string(),
            extra_env: Vec::new(),
            rt_threads: None,
            max_http_connections: None,
            optimized_http: true,
            forward_args: Vec::new(),
            passthrough_io: false,
        }
    }
}

/// Summary of a single worker's performance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerSummary {
    pub worker_id: usize,
    pub attempted: u64,
    pub succeeded: u64,
    pub failed: u64,
    pub bytes: Option<u64>,
}

/// Indicates whether byte totals are based on observed metrics or unavailable.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum BytesSource {
    Observed,
    Unknown,
}

/// Overall run summary with aggregated statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunSummary {
    pub started_at: String,
    pub finished_at: String,
    pub elapsed_seconds: f64,
    pub procs: usize,
    pub concurrent_per_proc: usize,
    pub attempted: u64,
    pub succeeded: u64,
    pub failed: u64,
    pub total_bytes: Option<u64>,
    pub bytes_source: BytesSource,
    pub per_worker: Vec<WorkerSummary>,
}

impl RunSummary {
    /// Calculate total throughput in MB/s
    pub fn throughput_mbps(&self) -> Option<f64> {
        self.total_bytes.and_then(|bytes| {
            if self.elapsed_seconds > 0.0 {
                Some((bytes as f64 / (1024.0 * 1024.0)) / self.elapsed_seconds)
            } else {
                None
            }
        })
    }

    /// Calculate operations per second
    pub fn ops_per_second(&self) -> f64 {
        if self.elapsed_seconds > 0.0 {
            self.succeeded as f64 / self.elapsed_seconds
        } else {
            0.0
        }
    }

    /// Calculate average object size in bytes
    pub fn avg_object_size(&self) -> Option<f64> {
        match (self.total_bytes, self.succeeded) {
            (Some(bytes), succeeded) if succeeded > 0 => Some(bytes as f64 / succeeded as f64),
            _ => None,
        }
    }
}

/// Run a multi-process GET operation with the given configuration
pub fn run_get_shards(cfg: &MpGetConfig) -> Result<RunSummary> {
    let start = Instant::now();
    let started_at = chrono::Utc::now().to_rfc3339();

    // Read and shard the key list
    let keys = read_lines(&cfg.keylist)
        .with_context(|| format!("Failed to read keylist: {:?}", cfg.keylist))?;
    
    if keys.is_empty() {
        anyhow::bail!("No keys found in keylist: {:?}", cfg.keylist);
    }

    let shards = shard_keys(&keys, cfg.procs);
    
    // Create temporary directory for this run
    let run_dir = tempfile::tempdir().context("Failed to create temp directory")?;
    let run_dir_path = run_dir.path();

    // Write shard files
    let shard_files = write_shards(run_dir_path, &shards)
        .context("Failed to write shard files")?;

    // Setup oplog files if requested
    let oplog_files = if let Some(ref oplog_dir) = cfg.op_log_dir {
        std::fs::create_dir_all(oplog_dir)
            .with_context(|| format!("Failed to create oplog dir: {:?}", oplog_dir))?;
        
        (0..cfg.procs)
            .map(|i| Some(oplog_dir.join(format!("worker_{}.jsonl", i))))
            .collect()
    } else {
        vec![None; cfg.procs]
    };

    // Spawn worker processes
    let mut stderr_files: Vec<Option<PathBuf>> = vec![None; cfg.procs];
    let mut children = Vec::with_capacity(cfg.procs);
    for (i, shard_file) in shard_files.iter().enumerate() {
        let mut cmd = Command::new(&cfg.worker_cmd);
        
        // Add global arguments first (like --op-log)
        if let Some(ref oplog_file) = oplog_files[i] {
            cmd.arg("--op-log").arg(oplog_file);
        }
        
        // Add subcommand and its arguments
        cmd.arg(&cfg.worker_subcmd)
           .arg("--concurrent")
           .arg(cfg.concurrent_per_proc.to_string())
           .arg("--keylist")
           .arg(shard_file);

        // Add duration if specified
        if let Some(duration) = cfg.duration_secs {
            cmd.arg("--duration").arg(duration.to_string());
        }

        // Add forward args
        cmd.args(&cfg.forward_args);

        // Setup environment
        let mut env_vars = cfg.extra_env.clone();
        
        // Pass through AWS environment variables
        let aws_vars = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION", 
                       "AWS_ENDPOINT_URL", "AWS_PROFILE", "AWS_DEFAULT_REGION"];
        for var in &aws_vars {
            if let Ok(value) = std::env::var(var) {
                env_vars.push((var.to_string(), value));
            }
        }

        // Set runtime threads
        if let Some(rt_threads) = cfg.rt_threads {
            env_vars.push(("S3DLIO_RT_THREADS".to_string(), rt_threads.to_string()));
        }

        // Set max HTTP connections
        if let Some(max_http) = cfg.max_http_connections {
            env_vars.push(("S3DLIO_MAX_HTTP_CONNECTIONS".to_string(), max_http.to_string()));
        }

        // Set optimized HTTP
        if cfg.optimized_http {
            env_vars.push(("S3DLIO_USE_OPTIMIZED_HTTP".to_string(), "true".to_string()));
        }

        // Apply environment variables
        for (key, value) in env_vars {
            cmd.env(key, value);
        }

        // Setup IO redirection
        if cfg.passthrough_io {
            cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
        } else {
            let stderr_path = run_dir_path.join(format!("worker_{}.stderr.log", i));
            let stderr_file = File::create(&stderr_path)
                .with_context(|| format!("Failed to create stderr file: {:?}", stderr_path))?;
            cmd.stdout(Stdio::null()).stderr(Stdio::from(stderr_file));
            stderr_files[i] = Some(stderr_path);
        }

        // Spawn the worker
        let child = cmd.spawn()
            .with_context(|| format!("Failed to spawn worker {} with command: {:?}", i, cfg.worker_cmd))?;
        
        children.push((i, child));
    }

    // Wait for all workers to complete
    let mut summaries = Vec::with_capacity(cfg.procs);
    let mut worker_failures: Vec<String> = Vec::new();
    for (worker_id, mut child) in children {
        let status = child.wait()
            .with_context(|| format!("Failed to wait for worker {}", worker_id))?;

        // Parse worker results
        let shard_size = shards.get(worker_id).map(|s| s.len()).unwrap_or(0) as u64;
        let stderr_path = stderr_files.get(worker_id).and_then(|p| p.as_ref());
        if !status.success() {
            let first_stderr = stderr_path
                .map(|p| read_first_non_empty_line(p.as_path()))
                .flatten()
                .unwrap_or_else(|| "(no stderr output captured)".to_string());
            let status_code = status
                .code()
                .map(|c| c.to_string())
                .unwrap_or_else(|| "terminated by signal".to_string());

            worker_failures.push(format!(
                "worker={} status={} stderr={} stderr_path={}",
                worker_id,
                status_code,
                first_stderr,
                stderr_path
                    .map(|p| p.display().to_string())
                    .unwrap_or_else(|| "(passthrough)".to_string())
            ));
        }

        let summary = summarize_worker(
            worker_id,
            oplog_files.get(worker_id).and_then(|o| o.as_ref()),
            shard_size,
            status.success(),
        )
            .with_context(|| format!("Failed to summarize worker {}", worker_id))?;
        
        summaries.push(summary);
    }

    let elapsed = start.elapsed().as_secs_f64();
    let finished_at = chrono::Utc::now().to_rfc3339();

    let attempted: u64 = summaries.iter().map(|s| s.attempted).sum();
    let succeeded: u64 = summaries.iter().map(|s| s.succeeded).sum();
    let failed: u64 = summaries.iter().map(|s| s.failed).sum();

    let all_bytes_observed = summaries.iter().all(|s| s.bytes.is_some());
    let (total_bytes, bytes_source) = if all_bytes_observed {
        (
            Some(summaries.iter().filter_map(|s| s.bytes).sum::<u64>()),
            BytesSource::Observed,
        )
    } else {
        (None, BytesSource::Unknown)
    };

    if !worker_failures.is_empty() {
        eprintln!("\n=== Multi-Process GET Result (invalid) ===");
        eprintln!(
            "attempted={}, succeeded={}, failed={}, bytes_source={:?}",
            attempted, succeeded, failed, bytes_source
        );
        eprintln!("Worker failures:");
        for detail in &worker_failures {
            eprintln!("  - {}", detail);
        }
        anyhow::bail!(
            "mp-get failed: one or more workers exited non-zero; metrics are invalid due to worker failures, throughput omitted"
        );
    }

    Ok(RunSummary {
        started_at,
        finished_at,
        elapsed_seconds: elapsed,
        procs: cfg.procs,
        concurrent_per_proc: cfg.concurrent_per_proc,
        attempted,
        succeeded,
        failed,
        total_bytes,
        bytes_source,
        per_worker: summaries,
    })
}

/// Builder for MpGetConfig
pub struct MpGetConfigBuilder {
    config: MpGetConfig,
}

impl MpGetConfigBuilder {
    pub fn new() -> Self {
        Self {
            config: MpGetConfig::default(),
        }
    }

    pub fn procs(mut self, procs: usize) -> Self {
        self.config.procs = procs;
        self
    }

    pub fn concurrent_per_proc(mut self, concurrent: usize) -> Self {
        self.config.concurrent_per_proc = concurrent;
        self
    }

    pub fn keylist<P: Into<PathBuf>>(mut self, path: P) -> Self {
        self.config.keylist = path.into();
        self
    }

    pub fn duration_secs(mut self, secs: u64) -> Self {
        self.config.duration_secs = Some(secs);
        self
    }

    pub fn op_log_dir<P: Into<PathBuf>>(mut self, dir: P) -> Self {
        self.config.op_log_dir = Some(dir.into());
        self
    }

    pub fn worker_cmd<S: Into<String>>(mut self, cmd: S) -> Self {
        self.config.worker_cmd = cmd.into();
        self
    }

    pub fn optimized_http(mut self, enabled: bool) -> Self {
        self.config.optimized_http = enabled;
        self
    }

    pub fn rt_threads(mut self, threads: usize) -> Self {
        self.config.rt_threads = Some(threads);
        self
    }

    pub fn max_http_connections(mut self, conns: usize) -> Self {
        self.config.max_http_connections = Some(conns);
        self
    }

    pub fn forward_arg<S: Into<String>>(mut self, arg: S) -> Self {
        self.config.forward_args.push(arg.into());
        self
    }

    pub fn forward_args<I, S>(mut self, args: I) -> Self 
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.config.forward_args.extend(args.into_iter().map(|s| s.into()));
        self
    }

    pub fn env<K, V>(mut self, key: K, value: V) -> Self 
    where
        K: Into<String>,
        V: Into<String>,
    {
        self.config.extra_env.push((key.into(), value.into()));
        self
    }

    pub fn passthrough_io(mut self, enabled: bool) -> Self {
        self.config.passthrough_io = enabled;
        self
    }

    pub fn build(self) -> MpGetConfig {
        self.config
    }
}

impl Default for MpGetConfigBuilder {
    fn default() -> Self {
        Self::new()
    }
}

// ----------------------------- Helper Functions -----------------------------

fn read_lines(path: &Path) -> Result<Vec<String>> {
    let file = File::open(path)
        .with_context(|| format!("Failed to open file: {:?}", path))?;
    let reader = BufReader::new(file);
    
    let mut lines = Vec::new();
    for line in reader.lines() {
        let line = line.context("Failed to read line")?;
        let trimmed = line.trim();
        if !trimmed.is_empty() {
            lines.push(trimmed.to_string());
        }
    }
    Ok(lines)
}

fn write_shards(run_dir: &Path, shards: &[Vec<String>]) -> Result<Vec<PathBuf>> {
    let mut shard_files = Vec::with_capacity(shards.len());
    
    for (i, shard) in shards.iter().enumerate() {
        let shard_file = run_dir.join(format!("shard_{}.txt", i));
        let mut file = File::create(&shard_file)
            .with_context(|| format!("Failed to create shard file: {:?}", shard_file))?;
        
        for key in shard {
            writeln!(file, "{}", key)
                .context("Failed to write key to shard file")?;
        }
        
        shard_files.push(shard_file);
    }
    
    Ok(shard_files)
}

fn shard_keys(keys: &[String], procs: usize) -> Vec<Vec<String>> {
    let mut shards: Vec<Vec<String>> = vec![Vec::new(); procs];
    
    for key in keys {
        let hash = hash64(key.as_bytes());
        let shard_index = (hash as usize) % procs;
        shards[shard_index].push(key.clone());
    }
    
    shards
}

fn summarize_worker(
    worker_id: usize,
    oplog_path: Option<&PathBuf>,
    shard_size: u64,
    worker_succeeded: bool,
) -> Result<WorkerSummary> {
    if let Some(path) = oplog_path {
        if path.exists() {
            let (observed_objects, observed_bytes) = parse_oplog(path)?;
            let succeeded = observed_objects.min(shard_size);
            let failed = shard_size.saturating_sub(succeeded);
            return Ok(WorkerSummary {
                worker_id,
                attempted: shard_size,
                succeeded,
                failed,
                bytes: observed_bytes,
            });
        }
    }

    let succeeded = if worker_succeeded { shard_size } else { 0 };
    let failed = shard_size.saturating_sub(succeeded);
    Ok(WorkerSummary {
        worker_id,
        attempted: shard_size,
        succeeded,
        failed,
        bytes: None,
    })
}

fn parse_oplog(path: &Path) -> Result<(u64, Option<u64>)> {
    let file = File::open(path)
        .with_context(|| format!("Failed to open oplog: {:?}", path))?;
    let reader = BufReader::new(file);
    
    let mut objects = 0u64;
    let mut bytes = 0u64;
    let mut saw_bytes = false;
    
    for line in reader.lines() {
        let line = line.context("Failed to read oplog line")?;
        if let Ok(value) = serde_json::from_str::<serde_json::Value>(&line) {
            // Extract bytes if present
            if let Some(size) = value.get("bytes").and_then(|v| v.as_u64()) {
                bytes += size;
                saw_bytes = true;
            }
            objects += 1;
        }
    }

    Ok((objects, if saw_bytes { Some(bytes) } else { None }))
}

fn read_first_non_empty_line(path: &Path) -> Option<String> {
    let file = File::open(path).ok()?;
    let reader = BufReader::new(file);
    for line in reader.lines() {
        let line = line.ok()?;
        let trimmed = line.trim();
        if !trimmed.is_empty() {
            return Some(trimmed.to_string());
        }
    }
    None
}

// Simple hash function for sharding keys
fn hash64(data: &[u8]) -> u64 {
    
    
    BuildHasherDefault::<DefaultHasher>::default().hash_one(data)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_shard_keys() {
        let keys = vec![
            "key1".to_string(),
            "key2".to_string(),
            "key3".to_string(),
            "key4".to_string(),
        ];
        
        let shards = shard_keys(&keys, 2);
        
        // Should have 2 shards
        assert_eq!(shards.len(), 2);
        
        // All keys should be distributed
        let total_keys: usize = shards.iter().map(|s| s.len()).sum();
        assert_eq!(total_keys, keys.len());
        
        // Sharding should be deterministic
        let shards2 = shard_keys(&keys, 2);
        assert_eq!(shards, shards2);
    }

    #[test]
    fn test_config_builder() {
        let config = MpGetConfigBuilder::new()
            .procs(8)
            .concurrent_per_proc(32)
            .keylist("/tmp/keys.txt")
            .optimized_http(true)
            .rt_threads(4)
            .forward_arg("s3://bucket/prefix")
            .env("AWS_REGION", "us-west-2")
            .build();

        assert_eq!(config.procs, 8);
        assert_eq!(config.concurrent_per_proc, 32);
        assert_eq!(config.keylist, PathBuf::from("/tmp/keys.txt"));
        assert!(config.optimized_http);
        assert_eq!(config.rt_threads, Some(4));
        assert_eq!(config.forward_args, vec!["s3://bucket/prefix"]);
        assert_eq!(config.extra_env, vec![("AWS_REGION".to_string(), "us-west-2".to_string())]);
    }

    #[test]
    fn test_run_summary_calculations() {
        let summary = RunSummary {
            started_at: "2025-01-01T00:00:00Z".to_string(),
            finished_at: "2025-01-01T00:01:00Z".to_string(),
            elapsed_seconds: 60.0,
            procs: 4,
            concurrent_per_proc: 64,
            attempted: 1000,
            succeeded: 1000,
            failed: 0,
            total_bytes: Some(1024 * 1024 * 100), // 100 MB
            bytes_source: BytesSource::Observed,
            per_worker: vec![],
        };

        // Test throughput calculation
        let throughput = summary.throughput_mbps().unwrap();
        assert!((throughput - (100.0 / 60.0)).abs() < 0.01);

        // Test ops per second
        let ops_per_sec = summary.ops_per_second();
        assert!((ops_per_sec - (1000.0 / 60.0)).abs() < 0.01);

        // Test average object size
        let avg_size = summary.avg_object_size().unwrap();
        // 100 MiB / 1000 objects = 104857.6 bytes per object
        assert_eq!(avg_size, 104857.6);
    }

    #[test]
    fn test_hash_distribution() {
        let keys: Vec<String> = (0..1000).map(|i| format!("key_{}", i)).collect();
        let shards = shard_keys(&keys, 4);
        
        // Check that we have 4 shards
        assert_eq!(shards.len(), 4);
        
        // Check that all keys are distributed
        let total: usize = shards.iter().map(|s| s.len()).sum();
        assert_eq!(total, 1000);
        
        // Check that distribution is reasonably balanced
        // (no shard should be completely empty for 1000 keys)
        for shard in &shards {
            assert!(!shard.is_empty());
        }
    }
}