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

use crate::topology::CpuGeneration;

pub const CHUNK_2_MB: usize = 2 * 1024 * 1024;
pub const CHUNK_4_MB: usize = 4 * 1024 * 1024;
pub const CHUNK_8_MB: usize = 8 * 1024 * 1024;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChunkPolicy {
    Fixed2M,
    Fixed4M,
    Fixed8M,
    Adaptive8To4,
}

#[derive(Debug, Clone, Copy)]
pub struct RuntimeMetrics {
    pub put_p99_ms: f64,
    pub median_queue_depth: f64,
}

#[derive(Debug, Clone)]
pub struct ChunkController {
    policy: ChunkPolicy,
    current_chunk_bytes: usize,
    downshift_streak: u32,
    upshift_streak: u32,
}

#[derive(Debug, Clone)]
pub struct EncryptedChunk {
    pub nonce_id: u64,
    pub crc32c: u32,
    pub plain_len: usize,
    pub cipher_len: usize,
    pub bytes: Vec<u8>,
}

#[derive(Debug, Clone)]
pub struct PipelineOutput {
    pub chunks: Vec<EncryptedChunk>,
    pub total_ciphertext_bytes: usize,
}

impl ChunkPolicy {
    pub fn from_str(value: &str) -> Result<Self> {
        match value {
            "fixed_2m" => Ok(Self::Fixed2M),
            "fixed_4m" => Ok(Self::Fixed4M),
            "fixed_8m" => Ok(Self::Fixed8M),
            "adaptive_8_to_4" => Ok(Self::Adaptive8To4),
            _ => anyhow::bail!("unsupported chunk policy: {value}"),
        }
    }
}

impl ChunkController {
    pub fn new(cpu_generation: CpuGeneration, policy: ChunkPolicy) -> Self {
        let current_chunk_bytes = match policy {
            ChunkPolicy::Fixed2M => CHUNK_2_MB,
            ChunkPolicy::Fixed4M => CHUNK_4_MB,
            ChunkPolicy::Fixed8M => CHUNK_8_MB,
            ChunkPolicy::Adaptive8To4 => match cpu_generation {
                CpuGeneration::Zen2 => CHUNK_2_MB,
                CpuGeneration::Zen3 | CpuGeneration::Unknown => CHUNK_8_MB,
            },
        };

        Self {
            policy,
            current_chunk_bytes,
            downshift_streak: 0,
            upshift_streak: 0,
        }
    }

    pub fn current_chunk_bytes(&self) -> usize {
        self.current_chunk_bytes
    }

    pub fn policy(&self) -> ChunkPolicy {
        self.policy
    }

    pub fn observe(&mut self, metrics: RuntimeMetrics) {
        if self.policy != ChunkPolicy::Adaptive8To4 {
            return;
        }

        let should_downshift = metrics.put_p99_ms > 200.0 && metrics.median_queue_depth > 64.0;
        let should_upshift = metrics.put_p99_ms < 120.0 && metrics.median_queue_depth < 32.0;

        if self.current_chunk_bytes == CHUNK_8_MB {
            if should_downshift {
                self.downshift_streak += 1;
            } else {
                self.downshift_streak = 0;
            }
            self.upshift_streak = 0;
            if self.downshift_streak >= 3 {
                self.current_chunk_bytes = CHUNK_4_MB;
                self.downshift_streak = 0;
            }
            return;
        }

        if self.current_chunk_bytes == CHUNK_4_MB {
            if should_upshift {
                self.upshift_streak += 1;
            } else {
                self.upshift_streak = 0;
            }
            self.downshift_streak = 0;
            if self.upshift_streak >= 30 {
                self.current_chunk_bytes = CHUNK_8_MB;
                self.upshift_streak = 0;
            }
        }
    }
}

pub fn process_payload_crc_zstd_aes(
    payload: &[u8],
    chunk_bytes: usize,
    key_bytes: &[u8; 32],
    nonce_seed: u64,
) -> Result<PipelineOutput> {
    if chunk_bytes == 0 {
        anyhow::bail!("chunk_bytes must be > 0");
    }

    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);
    let mut compressor = Compressor::new(1).context("failed initializing zstd compressor")?;

    let mut compressed_buffer = vec![0_u8; zstd::zstd_safe::compress_bound(chunk_bytes)];
    let mut output = PipelineOutput {
        chunks: Vec::new(),
        total_ciphertext_bytes: 0,
    };

    for (chunk_idx, chunk) in payload.chunks(chunk_bytes).enumerate() {
        let crc = crc32c::crc32c(chunk);
        let compressed_len = compressor
            .compress_to_buffer(chunk, &mut compressed_buffer)
            .with_context(|| format!("zstd compression failed at chunk {chunk_idx}"))?;

        let nonce_id = nonce_seed + chunk_idx as u64 + 1;
        let nonce = nonce_for_chunk(nonce_id);
        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 {chunk_idx}"))?;

        let mut bytes = Vec::with_capacity(compressed_len + tag.as_ref().len());
        bytes.extend_from_slice(&compressed_buffer[..compressed_len]);
        bytes.extend_from_slice(tag.as_ref());

        output.total_ciphertext_bytes += bytes.len();
        output.chunks.push(EncryptedChunk {
            nonce_id,
            crc32c: crc,
            plain_len: chunk.len(),
            cipher_len: bytes.len(),
            bytes,
        });
    }

    Ok(output)
}

pub fn restore_chunk_zstd_aes(
    encrypted_chunk: &[u8],
    expected_crc32c: u32,
    expected_plain_len: usize,
    key_bytes: &[u8; 32],
    nonce_id: u64,
) -> Result<Vec<u8>> {
    if encrypted_chunk.len() < aead::AES_256_GCM.tag_len() {
        anyhow::bail!(
            "encrypted chunk is too small: {} bytes",
            encrypted_chunk.len()
        );
    }

    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);
    let nonce = nonce_for_chunk(nonce_id);

    let mut sealed = encrypted_chunk.to_vec();
    let decrypted = key
        .open_in_place(nonce, Aad::from(expected_crc32c.to_le_bytes()), &mut sealed)
        .map_err(|_| anyhow::anyhow!("AES-256-GCM decrypt failed"))?;

    let restored = zstd::bulk::decompress(decrypted, expected_plain_len)
        .context("zstd decompression failed during restore")?;
    let crc = crc32c::crc32c(&restored);
    if crc != expected_crc32c {
        anyhow::bail!(
            "crc32c mismatch after decrypt/decompress: expected {}, got {}",
            expected_crc32c,
            crc
        );
    }

    Ok(restored)
}

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 adaptive_policy_downshifts_after_three_windows() {
        let mut controller = ChunkController::new(CpuGeneration::Zen3, ChunkPolicy::Adaptive8To4);
        for _ in 0..2 {
            controller.observe(RuntimeMetrics {
                put_p99_ms: 250.0,
                median_queue_depth: 100.0,
            });
            assert_eq!(controller.current_chunk_bytes(), CHUNK_8_MB);
        }
        controller.observe(RuntimeMetrics {
            put_p99_ms: 250.0,
            median_queue_depth: 100.0,
        });
        assert_eq!(controller.current_chunk_bytes(), CHUNK_4_MB);
    }

    #[test]
    fn adaptive_policy_upshifts_after_thirty_windows() {
        let mut controller = ChunkController::new(CpuGeneration::Zen3, ChunkPolicy::Adaptive8To4);
        for _ in 0..3 {
            controller.observe(RuntimeMetrics {
                put_p99_ms: 300.0,
                median_queue_depth: 120.0,
            });
        }
        assert_eq!(controller.current_chunk_bytes(), CHUNK_4_MB);
        for _ in 0..30 {
            controller.observe(RuntimeMetrics {
                put_p99_ms: 80.0,
                median_queue_depth: 20.0,
            });
        }
        assert_eq!(controller.current_chunk_bytes(), CHUNK_8_MB);
    }

    #[test]
    fn process_payload_emits_multiple_chunks() {
        let payload = vec![7_u8; CHUNK_2_MB + 1024];
        let output = process_payload_crc_zstd_aes(&payload, CHUNK_2_MB, &[0x42_u8; 32], 100)
            .expect("pipeline processing should succeed");
        assert_eq!(output.chunks.len(), 2);
        assert!(output.total_ciphertext_bytes > 0);
    }

    #[test]
    fn restore_chunk_round_trip_works() {
        let payload = vec![9_u8; CHUNK_2_MB + 321];
        let output = process_payload_crc_zstd_aes(&payload, CHUNK_2_MB, &[0xAA_u8; 32], 42)
            .expect("pipeline processing should succeed");
        let mut restored = Vec::new();
        for chunk in output.chunks {
            let plain = restore_chunk_zstd_aes(
                &chunk.bytes,
                chunk.crc32c,
                chunk.plain_len,
                &[0xAA_u8; 32],
                chunk.nonce_id,
            )
            .expect("restore should succeed");
            restored.extend_from_slice(&plain);
        }
        assert_eq!(restored, payload);
    }
}
