//! Federation Benchmark Framework (§8.11)
//!
//! Provides a standardised benchmarking suite for measuring federation
//! performance across six categories:
//!
//! 1. **single_query** — single-store query latency & throughput
//! 2. **single_write** — single-store write latency & throughput
//! 3. **cross_query**  — cross-store (JOIN) query
//! 4. **batch_write**  — multi-operation batch write
//! 5. **transaction**  — 2PC transaction commit
//! 6. **concurrent**   — multi-threaded query throughput
//!
//! Each benchmark yields a [`BenchmarkResult`] with P50 / P99 / max latency,
//! QPS, and error count.  Results can be pretty-printed as an ASCII table via
//! [`BenchmarkReport`].

use std::sync::Arc;
use std::time::Instant;

use super::error::{FederationError, Result};
use super::traits::Store;
use super::types::{
    UniversalQuery, WriteOp, Record,
};

// ---------------------------------------------------------------------------
// Data structures
// ---------------------------------------------------------------------------

/// Result of a single benchmark run.
#[derive(Debug, Clone)]
pub struct BenchmarkResult {
    /// Benchmark name (e.g. "single_query")
    pub name: String,
    /// Total number of operations executed
    pub total_ops: u64,
    /// Wall-clock duration in seconds
    pub duration_secs: f64,
    /// Operations per second
    pub ops_per_sec: f64,
    /// 50th-percentile latency in milliseconds
    pub latency_p50_ms: f64,
    /// 99th-percentile latency in milliseconds
    pub latency_p99_ms: f64,
    /// Maximum observed latency in milliseconds
    pub latency_max_ms: f64,
    /// Number of errors encountered
    pub errors: u64,
}

/// Configuration for a benchmark run.
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    /// Iterations for single-store query benchmark
    pub single_query_iters: u64,
    /// Iterations for single-store write benchmark
    pub single_write_iters: u64,
    /// Iterations for cross-store query benchmark
    pub cross_query_iters: u64,
    /// Number of batch write batches
    pub batch_write_batches: u64,
    /// Operations per batch
    pub batch_size: u64,
    /// Iterations for transaction benchmark
    pub transaction_iters: u64,
    /// Number of concurrent threads
    pub concurrent_threads: u32,
    /// Operations per concurrent thread
    pub concurrent_per_thread: u64,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            single_query_iters: 1000,
            single_write_iters: 500,
            cross_query_iters: 200,
            batch_write_batches: 50,
            batch_size: 20,
            transaction_iters: 200,
            concurrent_threads: 4,
            concurrent_per_thread: 500,
        }
    }
}

/// ASCII report formatter.
pub struct BenchmarkReport;

impl BenchmarkReport {
    /// Format benchmark results as a pretty ASCII table.
    pub fn format(results: &[BenchmarkResult]) -> String {
        let mut out = String::new();

        let version = env!("CARGO_PKG_VERSION");

        // Header
        out.push_str(&format!(
            "╔══════════════════════════════════════════════════════════════════╗\n"
        ));
        out.push_str(&format!(
            "║         Federation Benchmark Report v{:<24}║\n", version
        ));
        out.push_str(
            "╠════════════════════╦═══════╦══════════╦═══════╦═══════╦════════╣\n",
        );
        out.push_str(
            "║ Test               ║  Ops  ║ QPS      ║ P50   ║ P99   ║ Errors ║\n",
        );
        out.push_str(
            "╠════════════════════╬═══════╬══════════╬═══════╬═══════╬════════╣\n",
        );

        // Rows
        for r in results {
            out.push_str(&format!(
                "║ {:<18} ║ {:>5} ║ {:>8.0} ║ {:>5.1}ms║ {:>5.1}ms║ {:>6} ║\n",
                r.name,
                r.total_ops,
                r.ops_per_sec,
                r.latency_p50_ms,
                r.latency_p99_ms,
                r.errors,
            ));
        }

        // Footer
        out.push_str(
            "╚════════════════════╩═══════╩══════════╩═══════╩═══════╩════════╝\n",
        );

        out
    }

    /// Format as a compact single-line summary per benchmark.
    pub fn format_compact(results: &[BenchmarkResult]) -> String {
        let mut out = String::new();
        for r in results {
            out.push_str(&format!(
                "[{}] ops={} qps={:.0} p50={:.2}ms p99={:.2}ms max={:.2}ms err={}\n",
                r.name, r.total_ops, r.ops_per_sec,
                r.latency_p50_ms, r.latency_p99_ms, r.latency_max_ms, r.errors,
            ));
        }
        out
    }
}

// ---------------------------------------------------------------------------
// FederationBenchmark
// ---------------------------------------------------------------------------

/// Federation performance benchmark runner.
///
/// Operates against registered stores through the `Store` trait, so it does
/// not require a full `Federation` instance — making it easy to test in
/// isolation.
pub struct FederationBenchmark {
    /// Registered stores: alias → Arc<dyn Store>
    stores: std::collections::HashMap<String, Arc<dyn Store>>,
    /// Benchmark configuration
    config: BenchmarkConfig,
    /// Collected results
    results: Vec<BenchmarkResult>,
}

impl FederationBenchmark {
    /// Create a new benchmark runner with the given configuration.
    pub fn new(config: BenchmarkConfig) -> Self {
        Self {
            stores: std::collections::HashMap::new(),
            config,
            results: Vec::new(),
        }
    }

    /// Create with default configuration.
    pub fn with_defaults() -> Self {
        Self::new(BenchmarkConfig::default())
    }

    /// Register a store for benchmarking.
    pub fn register_store(&mut self, alias: &str, store: Arc<dyn Store>) {
        self.stores.insert(alias.to_string(), store);
    }

    /// Run the full benchmark suite.
    ///
    /// Requires at least one store registered. If two stores are registered,
    /// cross_query and transaction benchmarks are also run.
    pub fn run_all(&mut self) -> Result<Vec<BenchmarkResult>> {
        let aliases: Vec<String> = self.stores.keys().cloned().collect();
        if aliases.is_empty() {
            return Err(FederationError::ConfigError(
                "No stores registered for benchmark".into(),
            ));
        }

        let primary = &aliases[0];
        let mut results = Vec::new();

        // 1. single_query
        results.push(self.bench_single_query(primary));

        // 2. single_write
        results.push(self.bench_single_write(primary));

        // 3. cross_query (only if ≥ 2 stores)
        if aliases.len() >= 2 {
            results.push(self.bench_cross_query(&aliases[0], &aliases[1]));
        }

        // 4. batch_write
        results.push(self.bench_batch_write(primary));

        // 5. transaction
        results.push(self.bench_transaction(primary));

        // 6. concurrent queries
        results.push(self.bench_concurrent_queries(primary));

        self.results = results.clone();
        Ok(results)
    }

    /// Get the results from the last run.
    pub fn results(&self) -> &[BenchmarkResult] {
        &self.results
    }

    /// Generate an ASCII report from the last run.
    pub fn report(&self) -> String {
        BenchmarkReport::format(&self.results)
    }

    // ===== Individual benchmarks =============================================

    /// Benchmark: single-store query.
    pub fn bench_single_query(&self, alias: &str) -> BenchmarkResult {
        let iters = self.config.single_query_iters;
        let store = self.stores.get(alias);
        let query = Self::sample_query();
        let mut latencies = Vec::with_capacity(iters as usize);
        let mut errors = 0u64;
        let start = Instant::now();

        for _ in 0..iters {
            let op_start = Instant::now();
            match store {
                Some(s) => {
                    match s.query(&query) {
                        Ok(_) => {
                            latencies.push(op_start.elapsed().as_secs_f64() * 1000.0);
                        }
                        Err(_) => { errors += 1; }
                    }
                }
                None => { errors += 1; }
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("single_query", iters, duration, &latencies, errors)
    }

    /// Benchmark: single-store write.
    pub fn bench_single_write(&self, alias: &str) -> BenchmarkResult {
        let iters = self.config.single_write_iters;
        let store = self.stores.get(alias);
        let mut latencies = Vec::with_capacity(iters as usize);
        let mut errors = 0u64;
        let start = Instant::now();

        for i in 0..iters {
            let op_start = Instant::now();
            let op = Self::sample_write_op(i);
            match store {
                Some(s) => {
                    match s.write(op) {
                        Ok(_) => {
                            latencies.push(op_start.elapsed().as_secs_f64() * 1000.0);
                        }
                        Err(_) => { errors += 1; }
                    }
                }
                None => { errors += 1; }
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("single_write", iters, duration, &latencies, errors)
    }

    /// Benchmark: cross-store query (simulated via two sequential single-store queries + merge).
    pub fn bench_cross_query(&self, alias_a: &str, alias_b: &str) -> BenchmarkResult {
        let iters = self.config.cross_query_iters;
        let store_a = self.stores.get(alias_a);
        let store_b = self.stores.get(alias_b);
        let query = Self::sample_query();
        let mut latencies = Vec::with_capacity(iters as usize);
        let mut errors = 0u64;
        let start = Instant::now();

        for _ in 0..iters {
            let op_start = Instant::now();
            let ok = (|| -> std::result::Result<(), ()> {
                let _a = store_a.ok_or(())?.query(&query).map_err(|_| ())?;
                let _b = store_b.ok_or(())?.query(&query).map_err(|_| ())?;
                // Simulate merge/join overhead
                Ok(())
            })();
            match ok {
                Ok(_) => {
                    latencies.push(op_start.elapsed().as_secs_f64() * 1000.0);
                }
                Err(_) => { errors += 1; }
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("cross_query", iters, duration, &latencies, errors)
    }

    /// Benchmark: batch write.
    pub fn bench_batch_write(&self, alias: &str) -> BenchmarkResult {
        let batches = self.config.batch_write_batches;
        let batch_size = self.config.batch_size;
        let total_ops = batches * batch_size;
        let store = self.stores.get(alias);
        let mut latencies = Vec::with_capacity(batches as usize);
        let mut errors = 0u64;
        let start = Instant::now();

        for batch_idx in 0..batches {
            let ops: Vec<WriteOp> = (0..batch_size)
                .map(|j| Self::sample_write_op(batch_idx * batch_size + j))
                .collect();

            let op_start = Instant::now();
            match store {
                Some(s) => {
                    match s.batch_write(ops) {
                        Ok(r) => {
                            errors += r.failed_count as u64;
                            latencies.push(op_start.elapsed().as_secs_f64() * 1000.0);
                        }
                        Err(_) => { errors += batch_size; }
                    }
                }
                None => { errors += batch_size; }
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("batch_write", total_ops, duration, &latencies, errors)
    }

    /// Benchmark: transaction begin + write + commit.
    pub fn bench_transaction(&self, alias: &str) -> BenchmarkResult {
        let iters = self.config.transaction_iters;
        let store = self.stores.get(alias);
        let mut latencies = Vec::with_capacity(iters as usize);
        let mut errors = 0u64;
        let start = Instant::now();

        for i in 0..iters {
            let op_start = Instant::now();
            let result = (|| -> Result<()> {
                let s = store.ok_or_else(|| FederationError::StoreNotFound(alias.into()))?;
                let mut tx = s.begin_local_tx()?;
                tx.write(Self::sample_write_op(i + 100_000))?;
                let prepared = tx.prepare()?;
                if prepared {
                    tx.commit()?;
                } else {
                    tx.rollback()?;
                }
                Ok(())
            })();
            match result {
                Ok(_) => {
                    latencies.push(op_start.elapsed().as_secs_f64() * 1000.0);
                }
                Err(_) => { errors += 1; }
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("transaction", iters, duration, &latencies, errors)
    }

    /// Benchmark: concurrent queries from multiple threads.
    pub fn bench_concurrent_queries(&self, alias: &str) -> BenchmarkResult {
        let threads = self.config.concurrent_threads;
        let per_thread = self.config.concurrent_per_thread;
        let total_ops = threads as u64 * per_thread;

        let store = match self.stores.get(alias) {
            Some(s) => s.clone(),
            None => {
                return Self::compute_result("concurrent", total_ops, 0.0, &[], total_ops);
            }
        };

        let query = Self::sample_query();
        let start = Instant::now();

        let handles: Vec<_> = (0..threads)
            .map(|_| {
                let s = store.clone();
                let q = query.clone();
                let iters = per_thread;
                std::thread::spawn(move || {
                    let mut lats = Vec::with_capacity(iters as usize);
                    let mut errs = 0u64;
                    for _ in 0..iters {
                        let op_start = Instant::now();
                        match s.query(&q) {
                            Ok(_) => {
                                lats.push(op_start.elapsed().as_secs_f64() * 1000.0);
                            }
                            Err(_) => { errs += 1; }
                        }
                    }
                    (lats, errs)
                })
            })
            .collect();

        let mut all_latencies = Vec::new();
        let mut total_errors = 0u64;

        for h in handles {
            if let Ok((lats, errs)) = h.join() {
                all_latencies.extend(lats);
                total_errors += errs;
            }
        }

        let duration = start.elapsed().as_secs_f64();
        Self::compute_result("concurrent", total_ops, duration, &all_latencies, total_errors)
    }

    // ===== Helpers ===========================================================

    /// Compute percentile statistics from a latency vector.
    fn compute_result(
        name: &str,
        total_ops: u64,
        duration_secs: f64,
        latencies: &[f64],
        errors: u64,
    ) -> BenchmarkResult {
        let mut sorted = latencies.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let p50 = if sorted.is_empty() {
            0.0
        } else {
            sorted[sorted.len() / 2]
        };

        let p99 = if sorted.is_empty() {
            0.0
        } else {
            let idx = ((sorted.len() as f64) * 0.99) as usize;
            sorted[idx.min(sorted.len() - 1)]
        };

        let max = sorted.last().copied().unwrap_or(0.0);

        let ops_per_sec = if duration_secs > 0.0 {
            total_ops as f64 / duration_secs
        } else {
            0.0
        };

        BenchmarkResult {
            name: name.to_string(),
            total_ops,
            duration_secs,
            ops_per_sec,
            latency_p50_ms: p50,
            latency_p99_ms: p99,
            latency_max_ms: max,
            errors,
        }
    }

    /// Generate a simple sample query for benchmarking.
    fn sample_query() -> UniversalQuery {
        UniversalQuery {
            target: "*".to_string(),
            filters: vec![],
            order_by: None,
            limit: Some(10),
            offset: None,
            traverse: None,
            vector_search: None,
        }
    }

    /// Generate a sample write operation for benchmarking.
    fn sample_write_op(seq: u64) -> WriteOp {
        WriteOp::CreateDocument {
            id: format!("bench-doc-{}", seq),
            doc_type: "text".to_string(),
            content: serde_json::json!({
                "text": format!("Benchmark document #{}", seq),
                "importance": 0.5,
            }),
            tags: vec!["benchmark".to_string()],
            owner_id: None,
            visibility: None,
        }
    }
}

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

    /// Minimal mock store for benchmark tests.
    struct MockBenchStore {
        alias: String,
    }

    impl super::super::traits::AsAny for MockBenchStore {
        fn as_any(&self) -> &dyn std::any::Any { self }
        fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> { self }
    }

    impl Store for MockBenchStore {
        fn store_type(&self) -> super::super::types::StoreType {
            super::super::types::StoreType::Document
        }
        fn alias(&self) -> &str { &self.alias }
        fn query(&self, _q: &UniversalQuery) -> Result<Vec<Record>> {
            Ok(vec![])
        }
        fn write(&self, _op: WriteOp) -> Result<super::super::types::WriteResult> {
            Ok(super::super::types::WriteResult::Created("ok".into()))
        }
        fn batch_write(&self, ops: Vec<WriteOp>) -> Result<super::super::types::BatchWriteResult> {
            Ok(super::super::types::BatchWriteResult {
                success_count: ops.len(),
                failed_count: 0,
                results: ops.iter().map(|_| Ok(super::super::types::WriteResult::Created("ok".into()))).collect(),
                execution_time_ms: 1,
            })
        }
        fn begin_local_tx(&self) -> Result<Box<dyn super::super::traits::LocalTransaction + Send>> {
            Ok(Box::new(MockTx { alias: self.alias.clone() }))
        }
        fn get(&self, _id: &str) -> Result<Option<Record>> { Ok(None) }
        fn exists(&self, _id: &str) -> Result<bool> { Ok(false) }
        fn query_iter(&self, _q: &UniversalQuery) -> Result<Box<dyn Iterator<Item = Result<Record>> + Send>> {
            Ok(Box::new(std::iter::empty()))
        }
        fn stats(&self) -> super::super::traits::StoreStats {
            super::super::traits::StoreStats::default()
        }
        fn update_ownership(&self, _id: &str, _ownership: super::super::types::OwnershipMeta) -> Result<()> { Ok(()) }
        fn sync(&self) -> Result<()> { Ok(()) }
        fn close(&self) -> Result<()> { Ok(()) }
    }

    struct MockTx {
        alias: String,
    }

    impl super::super::traits::LocalTransaction for MockTx {
        fn store_alias(&self) -> &str { &self.alias }
        fn write(&mut self, _op: WriteOp) -> Result<()> { Ok(()) }
        fn prepare(&mut self) -> Result<bool> { Ok(true) }
        fn commit(&mut self) -> Result<()> { Ok(()) }
        fn rollback(&mut self) -> Result<()> { Ok(()) }
    }

    #[test]
    fn test_compute_result_empty() {
        let r = FederationBenchmark::compute_result("test", 0, 1.0, &[], 0);
        assert_eq!(r.name, "test");
        assert_eq!(r.latency_p50_ms, 0.0);
        assert_eq!(r.latency_p99_ms, 0.0);
    }

    #[test]
    fn test_compute_result_percentiles() {
        // 100 values: 1.0, 2.0, ..., 100.0
        let latencies: Vec<f64> = (1..=100).map(|i| i as f64).collect();
        let r = FederationBenchmark::compute_result("pct", 100, 1.0, &latencies, 0);
        assert_eq!(r.total_ops, 100);
        assert_eq!(r.ops_per_sec, 100.0);
        // P50 = index 50 → 51.0
        assert!((r.latency_p50_ms - 51.0).abs() < 0.01);
        // P99 = index 99 → 100.0
        assert!((r.latency_p99_ms - 100.0).abs() < 0.01);
        assert!((r.latency_max_ms - 100.0).abs() < 0.01);
    }

    #[test]
    fn test_run_all_no_stores() {
        let mut bench = FederationBenchmark::with_defaults();
        let result = bench.run_all();
        assert!(result.is_err());
    }

    #[test]
    fn test_run_all_single_store() {
        let mock = Arc::new(MockBenchStore { alias: "s1".into() });
        let mut bench = FederationBenchmark::new(BenchmarkConfig {
            single_query_iters: 50,
            single_write_iters: 30,
            cross_query_iters: 10,
            batch_write_batches: 5,
            batch_size: 10,
            transaction_iters: 20,
            concurrent_threads: 2,
            concurrent_per_thread: 50,
        });
        bench.register_store("s1", mock);

        let results = bench.run_all().unwrap();
        // single_query, single_write, batch_write, transaction, concurrent (no cross_query with 1 store)
        assert_eq!(results.len(), 5);
        assert_eq!(results[0].name, "single_query");
        assert_eq!(results[0].total_ops, 50);
        assert_eq!(results[0].errors, 0);

        assert_eq!(results[1].name, "single_write");
        assert_eq!(results[1].total_ops, 30);

        assert_eq!(results[2].name, "batch_write");
        assert_eq!(results[2].total_ops, 50); // 5 * 10

        assert_eq!(results[3].name, "transaction");
        assert_eq!(results[3].total_ops, 20);

        assert_eq!(results[4].name, "concurrent");
        assert_eq!(results[4].total_ops, 100); // 2 * 50

        // All should have 0 errors
        for r in &results {
            assert_eq!(r.errors, 0, "benchmark '{}' had errors", r.name);
        }
    }

    #[test]
    fn test_run_all_two_stores() {
        let mock_a = Arc::new(MockBenchStore { alias: "a".into() });
        let mock_b = Arc::new(MockBenchStore { alias: "b".into() });
        let mut bench = FederationBenchmark::new(BenchmarkConfig {
            single_query_iters: 10,
            single_write_iters: 10,
            cross_query_iters: 10,
            batch_write_batches: 2,
            batch_size: 5,
            transaction_iters: 10,
            concurrent_threads: 2,
            concurrent_per_thread: 10,
        });
        bench.register_store("a", mock_a);
        bench.register_store("b", mock_b);

        let results = bench.run_all().unwrap();
        // With 2 stores: includes cross_query → 6 benchmarks
        assert_eq!(results.len(), 6);

        let names: Vec<&str> = results.iter().map(|r| r.name.as_str()).collect();
        assert!(names.contains(&"cross_query"));
    }

    #[test]
    fn test_report_format() {
        let results = vec![
            BenchmarkResult {
                name: "single_query".into(),
                total_ops: 1000,
                duration_secs: 1.0,
                ops_per_sec: 1000.0,
                latency_p50_ms: 0.8,
                latency_p99_ms: 3.2,
                latency_max_ms: 5.0,
                errors: 0,
            },
            BenchmarkResult {
                name: "single_write".into(),
                total_ops: 500,
                duration_secs: 0.5,
                ops_per_sec: 1000.0,
                latency_p50_ms: 1.2,
                latency_p99_ms: 5.8,
                latency_max_ms: 8.0,
                errors: 0,
            },
        ];

        let report = BenchmarkReport::format(&results);
        assert!(report.contains("single_query"));
        assert!(report.contains("single_write"));
        assert!(report.contains("Federation Benchmark Report"));
        assert!(report.contains("QPS"));
        assert!(report.contains("P50"));
    }

    #[test]
    fn test_report_compact() {
        let results = vec![
            BenchmarkResult {
                name: "test".into(),
                total_ops: 100,
                duration_secs: 1.0,
                ops_per_sec: 100.0,
                latency_p50_ms: 5.0,
                latency_p99_ms: 15.0,
                latency_max_ms: 20.0,
                errors: 2,
            },
        ];
        let compact = BenchmarkReport::format_compact(&results);
        assert!(compact.contains("[test]"));
        assert!(compact.contains("ops=100"));
        assert!(compact.contains("err=2"));
    }

    #[test]
    fn test_bench_concurrent_zero_threads() {
        let _mock = Arc::new(MockBenchStore { alias: "s".into() });
        let bench = FederationBenchmark::new(BenchmarkConfig {
            concurrent_threads: 0,
            concurrent_per_thread: 100,
            ..BenchmarkConfig::default()
        });
        // No store registered — alias mismatch
        let result = bench.bench_concurrent_queries("s");
        // 0 threads × 100 = 0 total ops
        assert_eq!(result.total_ops, 0);
    }

    #[test]
    fn test_individual_benchmarks() {
        let mock = Arc::new(MockBenchStore { alias: "store".into() });
        let mut bench = FederationBenchmark::new(BenchmarkConfig {
            single_query_iters: 20,
            single_write_iters: 15,
            transaction_iters: 10,
            ..BenchmarkConfig::default()
        });
        bench.register_store("store", mock);

        let r = bench.bench_single_query("store");
        assert_eq!(r.total_ops, 20);
        assert!(r.ops_per_sec > 0.0);
        assert_eq!(r.errors, 0);

        let r = bench.bench_single_write("store");
        assert_eq!(r.total_ops, 15);
        assert_eq!(r.errors, 0);

        let r = bench.bench_transaction("store");
        assert_eq!(r.total_ops, 10);
        assert_eq!(r.errors, 0);
    }
}
