use std::time::Instant;

use anyhow::{Context, Result};
use ring::aead::{self, Aad, LessSafeKey, Nonce, UnboundKey};
use zstd::bulk::Compressor;

use crate::alloc_counter;

#[derive(Debug, Clone)]
pub struct ModuleAConfig {
    pub payload_bytes: usize,
    pub chunk_bytes: usize,
}

#[derive(Debug, Clone)]
pub struct ModuleAStats {
    pub payload_bytes: usize,
    pub chunk_bytes: usize,
    pub chunks: usize,
    pub elapsed_ms: f64,
    pub throughput_mb_s: f64,
    pub total_output_bytes: usize,
    pub hot_loop_allocations: u64,
    pub crc_fold: u32,
}

impl ModuleAStats {
    pub fn to_json(&self) -> String {
        format!(
            "{{\"module\":\"A\",\"payload_bytes\":{},\"chunk_bytes\":{},\"chunks\":{},\"elapsed_ms\":{:.3},\"throughput_mb_s\":{:.3},\"total_output_bytes\":{},\"hot_loop_allocations\":{},\"crc_fold\":{}}}",
            self.payload_bytes,
            self.chunk_bytes,
            self.chunks,
            self.elapsed_ms,
            self.throughput_mb_s,
            self.total_output_bytes,
            self.hot_loop_allocations,
            self.crc_fold
        )
    }
}

pub fn run(config: ModuleAConfig) -> Result<ModuleAStats> {
    if config.chunk_bytes == 0 {
        anyhow::bail!("chunk_bytes must be > 0");
    }
    if config.payload_bytes == 0 {
        anyhow::bail!("payload_bytes must be > 0");
    }
    if config.payload_bytes % config.chunk_bytes != 0 {
        anyhow::bail!("payload_bytes must be divisible by chunk_bytes");
    }

    let chunks = config.payload_bytes / config.chunk_bytes;
    let mut payload = vec![0_u8; config.payload_bytes];
    fill_semi_compressible(&mut payload);

    let chunk_output_bound = zstd::zstd_safe::compress_bound(config.chunk_bytes)
        .saturating_add(aead::AES_256_GCM.tag_len());
    let sink_capacity = chunk_output_bound
        .checked_mul(chunks)
        .context("sink capacity overflow")?;

    let mut compressed_buffer = vec![0_u8; zstd::zstd_safe::compress_bound(config.chunk_bytes)];
    let mut blackhole = vec![0_u8; sink_capacity];
    let mut compressor = Compressor::new(1).context("failed to initialize zstd compressor")?;

    let key_bytes = [0x42_u8; 32];
    let unbound = UnboundKey::new(&aead::AES_256_GCM, &key_bytes)
        .map_err(|_| anyhow::anyhow!("failed to initialize AES-256-GCM key"))?;
    let key = LessSafeKey::new(unbound);

    // Warm-up once to avoid allocator noise from lazy initialization.
    let warm_chunk = &payload[..config.chunk_bytes];
    let warm_compressed = compressor
        .compress_to_buffer(warm_chunk, &mut compressed_buffer[..])
        .context("warm-up zstd compression failed")?;
    let warm_nonce = nonce_for_chunk(0);
    let warm_tag = key
        .seal_in_place_separate_tag(
            warm_nonce,
            Aad::empty(),
            &mut compressed_buffer[..warm_compressed],
        )
        .map_err(|_| anyhow::anyhow!("warm-up AES-256-GCM encryption failed"))?;
    let warm_total = warm_compressed + warm_tag.as_ref().len();
    blackhole[..warm_compressed].copy_from_slice(&compressed_buffer[..warm_compressed]);
    blackhole[warm_compressed..warm_total].copy_from_slice(warm_tag.as_ref());

    let alloc_before = alloc_counter::snapshot();
    let start = Instant::now();
    let mut sink_offset = 0usize;
    let mut crc_fold = 0_u32;

    for idx in 0..chunks {
        let start_idx = idx * config.chunk_bytes;
        let end_idx = start_idx + config.chunk_bytes;
        let chunk = &payload[start_idx..end_idx];

        let crc = crc32c::crc32c(chunk);
        crc_fold ^= crc;

        let compressed_len = compressor
            .compress_to_buffer(chunk, &mut compressed_buffer[..])
            .with_context(|| format!("zstd compression failed at chunk {}", idx))?;

        let nonce = nonce_for_chunk(idx as u64 + 1);
        let tag = key
            .seal_in_place_separate_tag(
                nonce,
                Aad::from(crc.to_le_bytes()),
                &mut compressed_buffer[..compressed_len],
            )
            .map_err(|_| anyhow::anyhow!("AES-256-GCM encryption failed at chunk {}", idx))?;

        let encrypted_len = compressed_len + tag.as_ref().len();
        let end_offset = sink_offset + encrypted_len;
        if end_offset > blackhole.len() {
            anyhow::bail!("blackhole buffer overflow");
        }

        blackhole[sink_offset..sink_offset + compressed_len]
            .copy_from_slice(&compressed_buffer[..compressed_len]);
        blackhole[sink_offset + compressed_len..end_offset].copy_from_slice(tag.as_ref());
        sink_offset = end_offset;
    }

    let elapsed = start.elapsed();
    let hot_loop_allocations = alloc_counter::allocations_since(alloc_before);
    let elapsed_ms = elapsed.as_secs_f64() * 1_000.0;
    let throughput_mb_s = (config.payload_bytes as f64 / (1024.0 * 1024.0)) / elapsed.as_secs_f64();

    Ok(ModuleAStats {
        payload_bytes: config.payload_bytes,
        chunk_bytes: config.chunk_bytes,
        chunks,
        elapsed_ms,
        throughput_mb_s,
        total_output_bytes: sink_offset,
        hot_loop_allocations,
        crc_fold,
    })
}

fn fill_semi_compressible(buffer: &mut [u8]) {
    let pattern = b"TRACER_BULLET_PATTERN_";
    for (idx, byte) in buffer.iter_mut().enumerate() {
        if idx % 64 < pattern.len() {
            *byte = pattern[idx % pattern.len()];
        } else {
            let mixed = (idx as u64)
                .wrapping_mul(6364136223846793005_u64)
                .wrapping_add(1);
            *byte = (mixed >> 32) as u8;
        }
    }
}

fn nonce_for_chunk(chunk: u64) -> Nonce {
    let mut nonce = [0_u8; 12];
    nonce[4..].copy_from_slice(&chunk.to_be_bytes());
    Nonce::assume_unique_for_key(nonce)
}

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

    #[test]
    fn rejects_non_divisible_payload() {
        let err = run(ModuleAConfig {
            payload_bytes: 17,
            chunk_bytes: 8,
        })
        .expect_err("expected configuration validation failure");
        assert!(err.to_string().contains("divisible"));
    }

    #[test]
    fn processes_small_payload() {
        let stats = run(ModuleAConfig {
            payload_bytes: 8 * 1024 * 1024,
            chunk_bytes: 1024 * 1024,
        })
        .expect("module-a small payload run should succeed");
        assert_eq!(stats.chunks, 8);
        assert!(stats.total_output_bytes > 0);
        assert!(stats.elapsed_ms >= 0.0);
        assert!(stats.throughput_mb_s > 0.0);
    }

    #[test]
    fn rejects_zero_chunk() {
        let err = run(ModuleAConfig {
            payload_bytes: 1024,
            chunk_bytes: 0,
        })
        .expect_err("expected zero chunk rejection");
        assert!(err.to_string().contains("chunk_bytes"));
    }

    #[test]
    fn json_output_contains_module_label() {
        let stats = run(ModuleAConfig {
            payload_bytes: 1024 * 1024,
            chunk_bytes: 256 * 1024,
        })
        .expect("module-a run should succeed");
        let json = stats.to_json();
        assert!(json.contains("\"module\":\"A\""));
        assert!(json.contains("\"chunks\":4"));
    }
}
