use crate::error::{DaveError, DaveResult};
use hkdf::Hkdf;
use pyo3::prelude::*;
use sha2::Sha256;
use std::collections::HashMap;
use std::sync::RwLock;

const KEY_SIZE: usize = 32;
const MAX_GENERATION_GAP: u32 = 256;

/// Key ratchet for forward secrecy in media encryption
///
/// Based on DAVE protocol key ratcheting scheme
#[pyclass]
pub struct KeyRatchet {
    base_key: Vec<u8>,
    current_generation: u32,
    cached_keys: RwLock<HashMap<u32, Vec<u8>>>,
    max_cached: usize,
}

#[pymethods]
impl KeyRatchet {
    #[new]
    #[pyo3(signature = (base_key, max_cached = 16))]
    pub fn new(base_key: Vec<u8>, max_cached: usize) -> PyResult<Self> {
        if base_key.len() < KEY_SIZE {
            return Err(DaveError::Crypto("Base key too short".into()).into());
        }

        Ok(KeyRatchet {
            base_key,
            current_generation: 0,
            cached_keys: RwLock::new(HashMap::new()),
            max_cached,
        })
    }

    #[getter]
    pub fn generation(&self) -> u32 {
        self.current_generation
    }

    pub fn get_key(&self, generation: u32) -> PyResult<Vec<u8>> {
        self.derive_key(generation).map_err(Into::into)
    }

    pub fn current_key(&self) -> PyResult<Vec<u8>> {
        self.derive_key(self.current_generation).map_err(Into::into)
    }

    pub fn ratchet(&mut self) -> PyResult<Vec<u8>> {
        self.current_generation = self.current_generation.wrapping_add(1);
        self.derive_key(self.current_generation).map_err(Into::into)
    }

    pub fn ratchet_to(&mut self, generation: u32) -> PyResult<Vec<u8>> {
        if generation < self.current_generation {
            return Err(DaveError::InvalidState(
                "Cannot ratchet backwards".into()
            ).into());
        }

        if generation - self.current_generation > MAX_GENERATION_GAP {
            return Err(DaveError::InvalidState(
                format!("Generation gap too large: {}", generation - self.current_generation)
            ).into());
        }

        self.current_generation = generation;
        self.derive_key(generation).map_err(Into::into)
    }

    pub fn truncate_cache(&self, min_generation: u32) {
        let mut cache = self.cached_keys.write().unwrap();
        cache.retain(|&gen, _| gen >= min_generation);
    }

    pub fn clear_cache(&self) {
        self.cached_keys.write().unwrap().clear();
    }
}

impl KeyRatchet {
    fn derive_key(&self, generation: u32) -> DaveResult<Vec<u8>> {
        // Check cache first
        {
            let cache = self.cached_keys.read().unwrap();
            if let Some(key) = cache.get(&generation) {
                return Ok(key.clone());
            }
        }

        // Derive key using HKDF
        let hkdf = Hkdf::<Sha256>::new(None, &self.base_key);

        let info = format!("DAVE Ratchet Key {}", generation);
        let mut key = vec![0u8; KEY_SIZE];

        hkdf.expand(info.as_bytes(), &mut key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        // Cache the key
        {
            let mut cache = self.cached_keys.write().unwrap();

            // Evict old keys if cache is full
            if cache.len() >= self.max_cached {
                let min_gen = cache.keys().min().copied();
                if let Some(min) = min_gen {
                    cache.remove(&min);
                }
            }

            cache.insert(generation, key.clone());
        }

        Ok(key)
    }
}

/// Sender key ratchet - manages per-sender key evolution
#[pyclass]
pub struct SenderKeyRatchet {
    sender_id: u64,
    ratchet: KeyRatchet,
}

#[pymethods]
impl SenderKeyRatchet {
    #[new]
    pub fn new(sender_id: u64, base_key: Vec<u8>) -> PyResult<Self> {
        let ratchet = KeyRatchet::new(base_key, 16)?;
        Ok(SenderKeyRatchet { sender_id, ratchet })
    }

    #[getter]
    pub fn sender_id(&self) -> u64 {
        self.sender_id
    }

    #[getter]
    pub fn generation(&self) -> u32 {
        self.ratchet.generation()
    }

    pub fn get_key(&self, generation: u32) -> PyResult<Vec<u8>> {
        self.ratchet.get_key(generation)
    }

    pub fn ratchet(&mut self) -> PyResult<Vec<u8>> {
        self.ratchet.ratchet()
    }
}

/// Manager for multiple sender key ratchets
#[pyclass]
pub struct KeyRatchetManager {
    epoch_base_key: Vec<u8>,
    ratchets: RwLock<HashMap<u64, KeyRatchet>>,
}

#[pymethods]
impl KeyRatchetManager {
    #[new]
    pub fn new(epoch_base_key: Vec<u8>) -> Self {
        KeyRatchetManager {
            epoch_base_key,
            ratchets: RwLock::new(HashMap::new()),
        }
    }

    pub fn get_or_create_ratchet(&self, sender_id: u64) -> PyResult<()> {
        let mut ratchets = self.ratchets.write().unwrap();
        if !ratchets.contains_key(&sender_id) {
            let sender_key = self.derive_sender_base_key(sender_id)?;
            let ratchet = KeyRatchet::new(sender_key, 16)?;
            ratchets.insert(sender_id, ratchet);
        }
        Ok(())
    }

    pub fn get_key(&self, sender_id: u64, generation: u32) -> PyResult<Vec<u8>> {
        self.get_or_create_ratchet(sender_id)?;

        let ratchets = self.ratchets.read().unwrap();
        ratchets
            .get(&sender_id)
            .ok_or_else(|| DaveError::InvalidKeyId(sender_id))?
            .get_key(generation)
    }

    pub fn remove_sender(&self, sender_id: u64) {
        self.ratchets.write().unwrap().remove(&sender_id);
    }

    pub fn clear(&self) {
        self.ratchets.write().unwrap().clear();
    }
}

impl KeyRatchetManager {
    fn derive_sender_base_key(&self, sender_id: u64) -> DaveResult<Vec<u8>> {
        let hkdf = Hkdf::<Sha256>::new(None, &self.epoch_base_key);

        let info = format!("DAVE Sender Key {}", sender_id);
        let mut key = vec![0u8; KEY_SIZE];

        hkdf.expand(info.as_bytes(), &mut key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        Ok(key)
    }
}