use crate::error::{DaveError, DaveResult};
use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes128Gcm, Nonce,
};
use hkdf::Hkdf;
use pyo3::prelude::*;
use sha2::Sha256;
use std::sync::atomic::{AtomicU64, Ordering};
use std::collections::HashMap;
use std::sync::RwLock;

const TAG_SIZE: usize = 16;
const KEY_SIZE: usize = 16;
const SALT_SIZE: usize = 12;

/// SFrame header format for DAVE
#[pyclass]
#[derive(Clone)]
pub struct SFrameHeader {
    #[pyo3(get)]
    pub key_id: u64,
    #[pyo3(get)]
    pub counter: u64,
}

#[pymethods]
impl SFrameHeader {
    #[new]
    pub fn new(key_id: u64, counter: u64) -> Self {
        SFrameHeader { key_id, counter }
    }

    pub fn serialize(&self) -> Vec<u8> {
        encode_header(self.key_id, self.counter)
    }

    #[staticmethod]
    pub fn parse(data: &[u8]) -> PyResult<(Self, usize)> {
        let (kid, ctr, consumed) = decode_header(data)?;
        Ok((SFrameHeader { key_id: kid, counter: ctr }, consumed))
    }
}

/// Encodes an SFrame header
fn encode_header(key_id: u64, counter: u64) -> Vec<u8> {
    // Use compact encoding when possible
    if key_id < 8 && counter < 8 {
        // Single byte: 0 | KID(3) | CTR(4)
        vec![((key_id as u8) << 4) | (counter as u8)]
    } else {
        // Extended format
        let kid_bytes = encode_varint(key_id);
        let ctr_bytes = encode_varint(counter);

        let k_len = kid_bytes.len() - 1;
        let ctr_len = ctr_bytes.len() - 1;

        let first_byte = 0x80 | ((k_len as u8) << 4) | (ctr_len as u8);

        let mut result = vec![first_byte];
        result.extend(kid_bytes);
        result.extend(ctr_bytes);
        result
    }
}

/// Decodes an SFrame header
fn decode_header(data: &[u8]) -> DaveResult<(u64, u64, usize)> {
    if data.is_empty() {
        return Err(DaveError::InvalidFrame("Empty header".into()));
    }

    let first_byte = data[0];

    if first_byte & 0x80 == 0 {
        // Compact format
        let kid = ((first_byte >> 4) & 0x07) as u64;
        let ctr = (first_byte & 0x0F) as u64;
        Ok((kid, ctr, 1))
    } else {
        // Extended format
        let k_len = ((first_byte >> 4) & 0x07) as usize + 1;
        let ctr_len = (first_byte & 0x0F) as usize + 1;

        if data.len() < 1 + k_len + ctr_len {
            return Err(DaveError::InvalidFrame("Header too short".into()));
        }

        let kid = decode_varint(&data[1..1 + k_len]);
        let ctr = decode_varint(&data[1 + k_len..1 + k_len + ctr_len]);

        Ok((kid, ctr, 1 + k_len + ctr_len))
    }
}

fn encode_varint(value: u64) -> Vec<u8> {
    if value == 0 {
        return vec![0];
    }

    let mut bytes = Vec::new();
    let mut v = value;

    while v > 0 {
        bytes.push((v & 0xFF) as u8);
        v >>= 8;
    }

    bytes.reverse();
    bytes
}

fn decode_varint(data: &[u8]) -> u64 {
    let mut value = 0u64;
    for &byte in data {
        value = (value << 8) | (byte as u64);
    }
    value
}

#[pyclass]
pub struct SFrameContext {
    #[allow(dead_code)]
    base_key: Vec<u8>,
    key_id: u64,
    counter: AtomicU64,
    cipher_key: Vec<u8>,
    salt: Vec<u8>,
}

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

        let (cipher_key, salt) = derive_sframe_keys(&base_key)?;

        Ok(SFrameContext {
            base_key,
            key_id,
            counter: AtomicU64::new(0),
            cipher_key,
            salt,
        })
    }

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

    #[getter]
    pub fn current_counter(&self) -> u64 {
        self.counter.load(Ordering::SeqCst)
    }

    pub fn encrypt(&self, plaintext: Vec<u8>) -> PyResult<Vec<u8>> {
        let counter = self.counter.fetch_add(1, Ordering::SeqCst);
        self.encrypt_with_counter(plaintext, counter)
    }

    pub fn encrypt_with_counter(&self, plaintext: Vec<u8>, counter: u64) -> PyResult<Vec<u8>> {
        let header = encode_header(self.key_id, counter);
        let nonce = self.compute_nonce(counter);

        let cipher = Aes128Gcm::new_from_slice(&self.cipher_key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let nonce = Nonce::from_slice(&nonce);

        // AAD is the header
        let ciphertext = cipher
            .encrypt(nonce, aes_gcm::aead::Payload {
                msg: &plaintext,
                aad: &header,
            })
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        // Output: header || ciphertext (includes tag)
        let mut output = header;
        output.extend(ciphertext);

        Ok(output)
    }

    pub fn decrypt(&self, ciphertext: Vec<u8>) -> PyResult<(Vec<u8>, SFrameHeader)> {
        let (header, header_len) = SFrameHeader::parse(&ciphertext)?;

        if header.key_id != self.key_id {
            return Err(DaveError::InvalidKeyId(header.key_id).into());
        }

        let encrypted = &ciphertext[header_len..];
        if encrypted.len() < TAG_SIZE {
            return Err(DaveError::InvalidFrame("Ciphertext too short".into()).into());
        }

        let nonce = self.compute_nonce(header.counter);
        let header_bytes = &ciphertext[..header_len];

        let cipher = Aes128Gcm::new_from_slice(&self.cipher_key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let nonce = Nonce::from_slice(&nonce);

        let plaintext = cipher
            .decrypt(nonce, aes_gcm::aead::Payload {
                msg: encrypted,
                aad: header_bytes,
            })
            .map_err(|e| DaveError::DecryptionFailed(format!("{:?}", e)))?;

        Ok((plaintext, header))
    }

    pub fn set_counter(&self, counter: u64) {
        self.counter.store(counter, Ordering::SeqCst);
    }

    pub fn __repr__(&self) -> String {
        format!(
            "SFrameContext(key_id={}, counter={})",
            self.key_id,
            self.counter.load(Ordering::SeqCst)
        )
    }
}

impl SFrameContext {
    fn compute_nonce(&self, counter: u64) -> Vec<u8> {
        let mut nonce = self.salt.clone();
        let counter_bytes = counter.to_be_bytes();

        // XOR counter into the end of the salt
        let offset = nonce.len() - counter_bytes.len();
        for (i, &b) in counter_bytes.iter().enumerate() {
            nonce[offset + i] ^= b;
        }

        nonce
    }
}

fn derive_sframe_keys(base_key: &[u8]) -> DaveResult<(Vec<u8>, Vec<u8>)> {
    let hkdf = Hkdf::<Sha256>::new(None, base_key);

    let mut cipher_key = vec![0u8; KEY_SIZE];
    let mut salt = vec![0u8; SALT_SIZE];

    hkdf.expand(b"SFrame 1.0 AES-GCM Key", &mut cipher_key)
        .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

    hkdf.expand(b"SFrame 1.0 AES-GCM Salt", &mut salt)
        .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

    Ok((cipher_key, salt))
}

/// Multi-key SFrame manager for handling multiple senders
#[pyclass]
pub struct SFrameMultiContext {
    contexts: RwLock<HashMap<u64, SFrameContextInner>>,
    local_key_id: u64,
}

struct SFrameContextInner {
    cipher_key: Vec<u8>,
    salt: Vec<u8>,
    counter: AtomicU64,
}

#[pymethods]
impl SFrameMultiContext {
    #[new]
    pub fn new(local_key_id: u64) -> Self {
        SFrameMultiContext {
            contexts: RwLock::new(HashMap::new()),
            local_key_id,
        }
    }

    pub fn add_sender(&self, key_id: u64, base_key: Vec<u8>) -> PyResult<()> {
        let (cipher_key, salt) = derive_sframe_keys(&base_key)?;
        let inner = SFrameContextInner {
            cipher_key,
            salt,
            counter: AtomicU64::new(0),
        };
        self.contexts.write().unwrap().insert(key_id, inner);
        Ok(())
    }

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

    pub fn encrypt(&self, plaintext: Vec<u8>) -> PyResult<Vec<u8>> {
        let contexts = self.contexts.read().unwrap();
        let ctx = contexts
            .get(&self.local_key_id)
            .ok_or_else(|| DaveError::InvalidKeyId(self.local_key_id))?;

        let counter = ctx.counter.fetch_add(1, Ordering::SeqCst);
        let header = encode_header(self.local_key_id, counter);
        let nonce = compute_nonce_inner(&ctx.salt, counter);

        let cipher = Aes128Gcm::new_from_slice(&ctx.cipher_key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let nonce = Nonce::from_slice(&nonce);

        let ciphertext = cipher
            .encrypt(nonce, aes_gcm::aead::Payload {
                msg: &plaintext,
                aad: &header,
            })
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let mut output = header;
        output.extend(ciphertext);

        Ok(output)
    }

    pub fn decrypt(&self, ciphertext: Vec<u8>) -> PyResult<(Vec<u8>, SFrameHeader)> {
        // Parse header to get key_id
        let (header, header_len) = SFrameHeader::parse(&ciphertext)?;

        let contexts = self.contexts.read().unwrap();
        let ctx = contexts
            .get(&header.key_id)
            .ok_or_else(|| DaveError::InvalidKeyId(header.key_id))?;

        let encrypted = &ciphertext[header_len..];
        if encrypted.len() < TAG_SIZE {
            return Err(DaveError::InvalidFrame("Ciphertext too short".into()).into());
        }

        let nonce = compute_nonce_inner(&ctx.salt, header.counter);
        let header_bytes = &ciphertext[..header_len];

        let cipher = Aes128Gcm::new_from_slice(&ctx.cipher_key)
            .map_err(|e| DaveError::Crypto(format!("{:?}", e)))?;

        let nonce = Nonce::from_slice(&nonce);

        let plaintext = cipher
            .decrypt(nonce, aes_gcm::aead::Payload {
                msg: encrypted,
                aad: header_bytes,
            })
            .map_err(|e| DaveError::DecryptionFailed(format!("{:?}", e)))?;

        Ok((plaintext, header))
    }

    pub fn has_sender(&self, key_id: u64) -> bool {
        self.contexts.read().unwrap().contains_key(&key_id)
    }

    pub fn sender_count(&self) -> usize {
        self.contexts.read().unwrap().len()
    }
}

fn compute_nonce_inner(salt: &[u8], counter: u64) -> Vec<u8> {
    let mut nonce = salt.to_vec();
    let counter_bytes = counter.to_be_bytes();

    let offset = nonce.len() - counter_bytes.len();
    for (i, &b) in counter_bytes.iter().enumerate() {
        nonce[offset + i] ^= b;
    }

    nonce
}