#![expect(
    clippy::redundant_pub_crate,
    reason = "criterion bench entry points are re-exported only to the private bench root module"
)]

use super::support::{
    BenchmarkId, Criterion, DownloadEngine, DownloadStatus, PieceId, PieceState, RuntimeConfig,
    Throughput,
};

#[derive(Clone, Copy, Debug)]
struct BackpressureScenario {
    task_count: usize,
    rounds: usize,
    disk_cache_bytes: u64,
    split: usize,
    max_connections_per_server: usize,
}

fn seed_instrumented_engine(task_count: usize) -> DownloadEngine {
    let runtime = RuntimeConfig {
        split: 4,
        max_connections_per_server: 4,
        max_connection_per_server: 4,
        min_split_size: 1024,
        piece_length: 1024,
        ..RuntimeConfig::default()
    };
    let mut engine = DownloadEngine::with_runtime(runtime);
    for i in 0..task_count {
        let gid = engine
            .add_uri(format!("magnet:?xt=urn:btih:{:040x}", 90_001 + i))
            .gid();
        let group = engine
            .handle_mut(gid)
            .expect("newly inserted benchmark group should exist");
        group.set_status(if i % 3 == 0 {
            DownloadStatus::Waiting
        } else {
            DownloadStatus::Active
        });
        group.set_total_length(8 * 1024);
        group.set_completed_length((i % 4) as u64 * 1024);
        group.set_piece_length(1024);
        group.set_piece_state(PieceId(0), PieceState::Verified);
        group.set_piece_state(PieceId(1), PieceState::Pending);
        group.set_piece_state(PieceId(2), PieceState::Downloading);
        group.set_piece_state(PieceId(3), PieceState::Queued);
        group.set_piece_state(PieceId(4), PieceState::Missing);
    }

    for _ in 0..3 {
        let _ = engine.schedule_once();
    }
    engine
}

fn seed_backpressure_engine(scenario: BackpressureScenario) -> DownloadEngine {
    let runtime = RuntimeConfig {
        split: scenario.split,
        max_connections_per_server: scenario.max_connections_per_server,
        max_connection_per_server: scenario.max_connections_per_server,
        min_split_size: 1024,
        piece_length: 1024,
        disk_cache_bytes: scenario.disk_cache_bytes,
        ..RuntimeConfig::default()
    };
    let mut engine = DownloadEngine::with_runtime(runtime);
    for i in 0..scenario.task_count {
        let gid = engine
            .add_uri(format!("https://example.org/backpressure-{i}.bin"))
            .gid();
        let group = engine
            .handle_mut(gid)
            .expect("newly inserted backpressure group should exist");
        group.set_status(match i % 4 {
            0 => DownloadStatus::Waiting,
            2 => DownloadStatus::Error,
            _ => DownloadStatus::Active,
        });
        group.set_total_length(16 * 1024);
        group.set_completed_length((i % 8) as u64 * 1024);
        group.set_piece_length(1024);
        group.set_download_speed(2_500 + i as u64 * 11);
        group.set_upload_speed(900 + i as u64 * 5);
        group.set_num_connections((scenario.max_connections_per_server.min(8)) as u32);
        group.set_retry_count((i % 3) as u32);
        group.set_piece_state(PieceId(0), PieceState::Verified);
        group.set_piece_state(PieceId(1), PieceState::Pending);
        group.set_piece_state(PieceId(2), PieceState::Downloading);
        group.set_piece_state(PieceId(3), PieceState::Queued);
        group.set_piece_state(PieceId(4), PieceState::Missing);
    }

    for _ in 0..4 {
        let _ = engine.schedule_once();
    }
    engine
}

fn run_runtime_snapshot_pressure(engine: &mut DownloadEngine, rounds: usize) -> usize {
    let gids = engine.registry().handles().collect::<Vec<_>>();
    for round in 0..rounds {
        let gid = gids[round % gids.len()].gid();
        let group = engine
            .handle_mut(gid)
            .expect("benchmark group should still exist");
        group.set_download_speed(600 + round as u64 * 10);
        group.set_upload_speed(200 + round as u64 * 5);
        let _ = engine.schedule_once();
        let runtime = engine.runtime_instrumentation_snapshot();
        assert!(runtime.download_count >= gids.len());
        assert!(runtime.scheduler_counters.schedule_run_count >= 1);
    }
    rounds
}

fn run_backpressure_runtime_pressure(
    engine: &mut DownloadEngine,
    scenario: BackpressureScenario,
) -> usize {
    let gids = engine.registry().handles().collect::<Vec<_>>();
    for round in 0..scenario.rounds {
        let gid = gids[round % gids.len()].gid();
        let group = engine
            .handle_mut(gid)
            .expect("backpressure benchmark group should still exist");
        group.set_retry_count((round % 5) as u32);
        group.set_status(if round % 3 == 0 {
            DownloadStatus::Waiting
        } else {
            DownloadStatus::Active
        });
        group.set_download_speed(4_000 + round as u64 * 40);
        group.set_upload_speed(1_500 + round as u64 * 25);

        let _ = engine.schedule_once();
        let runtime = engine.runtime_instrumentation_snapshot();
        assert_eq!(
            runtime.configured_disk_cache_bytes,
            scenario.disk_cache_bytes
        );
        assert!(runtime.total_active_segments >= 1);
        assert!(runtime.scheduler_counters.schedule_run_count >= 1);
    }
    scenario.rounds
}

pub(super) fn bench_runtime_snapshot_pressure(c: &mut Criterion) {
    let mut group = c.benchmark_group("runtime_snapshot_pressure");
    for task_count in [64usize, 128, 256] {
        group.throughput(Throughput::Elements(task_count as u64));
        group.bench_with_input(
            BenchmarkId::new("runtime_snapshot", task_count),
            &task_count,
            |b, &task_count| {
                b.iter_batched(
                    || seed_instrumented_engine(task_count),
                    |mut engine| {
                        let observations = run_runtime_snapshot_pressure(&mut engine, 8);
                        assert_eq!(observations, 8);
                    },
                    criterion::BatchSize::SmallInput,
                );
            },
        );
    }
    group.finish();
}

pub(super) fn bench_scheduler_backpressure_pressure(c: &mut Criterion) {
    let mut group = c.benchmark_group("scheduler_backpressure_pressure");
    for scenario in [
        BackpressureScenario {
            task_count: 64,
            rounds: 10,
            disk_cache_bytes: 4 * 1024 * 1024,
            split: 4,
            max_connections_per_server: 4,
        },
        BackpressureScenario {
            task_count: 128,
            rounds: 10,
            disk_cache_bytes: 32 * 1024 * 1024,
            split: 8,
            max_connections_per_server: 8,
        },
    ] {
        group.throughput(Throughput::Elements(scenario.task_count as u64));
        group.bench_with_input(
            BenchmarkId::new("backpressure", scenario.task_count),
            &scenario,
            |b, &scenario| {
                b.iter_batched(
                    || seed_backpressure_engine(scenario),
                    |mut engine| {
                        let observations = run_backpressure_runtime_pressure(&mut engine, scenario);
                        assert_eq!(observations, scenario.rounds);
                    },
                    criterion::BatchSize::SmallInput,
                );
            },
        );
    }
    group.finish();
}