use openssl::symm::{Cipher, Crypter, Mode};
use pyo3::prelude::*;
use pyo3::types::{PyByteArray, PyBytes};
use std::borrow::Cow;

#[inline(always)]
fn xor_block(a: &mut [u8; 16], b: &[u8; 16]) {
    for i in 0..16 {
        a[i] ^= b[i];
    }
}

#[inline(always)]
fn increment_iv(iv: &mut [u8; 16]) {
    for j in (0..16).rev() {
        iv[j] = iv[j].wrapping_add(1);
        if iv[j] != 0 {
            break;
        }
    }
}

#[inline(always)]
fn first_state_byte(state: &Bound<'_, PyAny>) -> PyResult<Option<u8>> {
    if let Ok(b) = state.cast::<PyByteArray>() {
        let bytes = unsafe { b.as_bytes() };
        return Ok(bytes.first().copied());
    }
    if let Ok(b) = state.cast::<PyBytes>() {
        return Ok(b.as_bytes().first().copied());
    }
    let bytes: Vec<u8> = state.extract()?;
    Ok(bytes.first().copied())
}

#[inline(always)]
fn iv_to_array(iv: &Bound<'_, PyAny>) -> PyResult<[u8; 16]> {
    if let Ok(b) = iv.cast::<PyByteArray>() {
        let bytes = unsafe { b.as_bytes() };
        if bytes.len() != 16 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("IV must be 16 bytes"));
        }
        let mut out = [0u8; 16];
        out.copy_from_slice(bytes);
        return Ok(out);
    }
    if let Ok(b) = iv.cast::<PyBytes>() {
        let bytes = b.as_bytes();
        if bytes.len() != 16 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("IV must be 16 bytes"));
        }
        let mut out = [0u8; 16];
        out.copy_from_slice(bytes);
        return Ok(out);
    }
    iv.extract()
}

#[inline(always)]
fn bytes_like<'py>(obj: &'py Bound<'py, PyAny>) -> PyResult<Cow<'py, [u8]>> {
    if let Ok(b) = obj.cast::<PyBytes>() {
        return Ok(Cow::Borrowed(b.as_bytes()));
    }
    if let Ok(b) = obj.cast::<PyByteArray>() {
        let bytes = unsafe { b.as_bytes() };
        return Ok(Cow::Borrowed(bytes));
    }
    let bytes: Vec<u8> = obj.extract()?;
    Ok(Cow::Owned(bytes))
}

#[inline(always)]
fn openssl_err(err: openssl::error::ErrorStack) -> PyErr {
    PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("OpenSSL error: {err}"))
}

/// AES-256-IGE Encryption
#[pyfunction]
#[pyo3(signature = (data, key, iv, /))]
pub fn ige256_encrypt<'py>(
    py: Python<'py>,
    data: &[u8],
    key: &[u8],
    iv: &[u8],
) -> PyResult<Bound<'py, PyBytes>> {
    if key.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Key must be 32 bytes"));
    }
    if iv.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("IV must be 32 bytes"));
    }
    if data.len() % 16 != 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Data length must be multiple of 16"));
    }

    let result = py.detach(|| -> PyResult<Vec<u8>> {
        let mut ecb = Crypter::new(Cipher::aes_256_ecb(), Mode::Encrypt, key, None)
            .map_err(openssl_err)?;
        ecb.pad(false);
        let mut iv1: [u8; 16] = iv[..16].try_into().unwrap();
        let mut iv2: [u8; 16] = iv[16..].try_into().unwrap();

        let mut out = Vec::with_capacity(data.len());
        unsafe { out.set_len(data.len()); }

        for (in_chunk, out_chunk) in data.chunks_exact(16).zip(out.chunks_exact_mut(16)) {
            let mut block: [u8; 16] = in_chunk.try_into().unwrap();
            let old_input = block;

            xor_block(&mut block, &iv1);
            let mut enc = [0u8; 32];
            let n = ecb.update(&block, &mut enc).map_err(openssl_err)?;
            debug_assert_eq!(n, 16);
            let mut enc_block: [u8; 16] = enc[..16].try_into().unwrap();
            xor_block(&mut enc_block, &iv2);

            out_chunk.copy_from_slice(&enc_block);
            iv1 = enc_block;
            iv2 = old_input;
        }
        let mut final_buf = [0u8; 32];
        let _ = ecb.finalize(&mut final_buf).map_err(openssl_err)?;
        Ok(out)
    })?;

    Ok(PyBytes::new(py, &result))
}

/// AES-256-IGE Decryption
#[pyfunction]
#[pyo3(signature = (data, key, iv, /))]
pub fn ige256_decrypt<'py>(
    py: Python<'py>,
    data: &[u8],
    key: &[u8],
    iv: &[u8],
) -> PyResult<Bound<'py, PyBytes>> {
    if key.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Key must be 32 bytes"));
    }
    if iv.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("IV must be 32 bytes"));
    }
    if data.len() % 16 != 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Data length must be multiple of 16"));
    }

    let result = py.detach(|| -> PyResult<Vec<u8>> {
        let mut ecb = Crypter::new(Cipher::aes_256_ecb(), Mode::Decrypt, key, None)
            .map_err(openssl_err)?;
        ecb.pad(false);
        let mut iv1: [u8; 16] = iv[16..].try_into().unwrap();
        let mut iv2: [u8; 16] = iv[..16].try_into().unwrap();

        let mut out = Vec::with_capacity(data.len());
        unsafe { out.set_len(data.len()); }

        for (in_chunk, out_chunk) in data.chunks_exact(16).zip(out.chunks_exact_mut(16)) {
            let block: [u8; 16] = in_chunk.try_into().unwrap();
            let mut x = block;
            xor_block(&mut x, &iv1);
            let mut dec = [0u8; 32];
            let n = ecb.update(&x, &mut dec).map_err(openssl_err)?;
            debug_assert_eq!(n, 16);
            let mut dec_block: [u8; 16] = dec[..16].try_into().unwrap();
            xor_block(&mut dec_block, &iv2);

            out_chunk.copy_from_slice(&dec_block);
            iv1 = dec_block;
            iv2 = block;
        }
        let mut final_buf = [0u8; 32];
        let _ = ecb.finalize(&mut final_buf).map_err(openssl_err)?;
        Ok(out)
    })?;

    Ok(PyBytes::new(py, &result))
}

/// AES-256-CTR Encryption/Decryption
#[pyfunction]
#[pyo3(signature = (data, key, iv, state, /))]
pub fn ctr256_encrypt<'py>(
    py: Python<'py>,
    data: &Bound<'py, PyAny>,
    key: &Bound<'py, PyAny>,
    iv: &Bound<'py, PyAny>,
    state: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyBytes>> {
    let key = bytes_like(key)?;
    let key = key.as_ref();
    if key.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Key must be 32 bytes"));
    }

    let data = bytes_like(data)?;
    let mut iv_array = iv_to_array(iv)?;
    let mut ks_pos = first_state_byte(state)?.map(|b| b as usize).unwrap_or(0);

    let (out_bytes, final_iv, final_state) = py.detach(|| {
        let mut ecb = Crypter::new(Cipher::aes_256_ecb(), Mode::Encrypt, key, None)
            .map_err(openssl_err)?;
        ecb.pad(false);

        let data = data.as_ref();
        let mut out = vec![0u8; data.len()];
        let mut offset = 0usize;

        if ks_pos > 0 && !data.is_empty() {
            let mut ks_buf = [0u8; 32];
            let n = ecb.update(&iv_array, &mut ks_buf).map_err(openssl_err)?;
            debug_assert_eq!(n, 16);
            let ks = &ks_buf[..16];
            let take = (16 - ks_pos).min(data.len());
            for i in 0..take {
                out[i] = data[i] ^ ks[ks_pos + i];
            }
            offset += take;
            if ks_pos + take == 16 {
                ks_pos = 0;
                increment_iv(&mut iv_array);
            } else {
                ks_pos += take;
            }
        }

        while offset + 16 <= data.len() {
            let mut ks_buf = [0u8; 32];
            let n = ecb.update(&iv_array, &mut ks_buf).map_err(openssl_err)?;
            debug_assert_eq!(n, 16);
            let ks = &ks_buf[..16];
            for i in 0..16 {
                out[offset + i] = data[offset + i] ^ ks[i];
            }
            offset += 16;
            increment_iv(&mut iv_array);
        }

        if offset < data.len() {
            let mut ks_buf = [0u8; 32];
            let n = ecb.update(&iv_array, &mut ks_buf).map_err(openssl_err)?;
            debug_assert_eq!(n, 16);
            let ks = &ks_buf[..16];
            let tail = data.len() - offset;
            for i in 0..tail {
                out[offset + i] = data[offset + i] ^ ks[i];
            }
            ks_pos = tail;
        }

        let mut final_buf = [0u8; 32];
        let _ = ecb.finalize(&mut final_buf).map_err(openssl_err)?;

        Ok::<(Vec<u8>, [u8; 16], usize), PyErr>((out, iv_array, ks_pos))
    })?;

    if let Ok(state_bytes_obj) = state.cast::<PyByteArray>() {
        if !state_bytes_obj.is_empty() {
            unsafe { state_bytes_obj.as_bytes_mut()[0] = final_state as u8; }
        }
    }
    if let Ok(iv_bytes_obj) = iv.cast::<PyByteArray>() {
        if iv_bytes_obj.len() == 16 {
            unsafe { iv_bytes_obj.as_bytes_mut().copy_from_slice(&final_iv); }
        }
    }

    Ok(PyBytes::new(py, &out_bytes))
}

#[pyfunction]
#[pyo3(signature = (data, key, iv, state, /))]
pub fn ctr256_decrypt<'py>(
    py: Python<'py>,
    data: &Bound<'py, PyAny>,
    key: &Bound<'py, PyAny>,
    iv: &Bound<'py, PyAny>,
    state: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyBytes>> {
    ctr256_encrypt(py, data, key, iv, state)
}

/// AES-256-CBC Encryption
#[pyfunction]
#[pyo3(signature = (data, key, iv, /))]
pub fn cbc256_encrypt<'py>(
    py: Python<'py>,
    data: &[u8],
    key: &[u8],
    iv: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyBytes>> {
    if key.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Key must be 32 bytes"));
    }
    let iv_array: [u8; 16] = iv.extract()?;
    if data.len() % 16 != 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Data length must be multiple of 16"));
    }

    let (out_bytes, final_iv) = py.detach(|| {
        let mut crypter = Crypter::new(Cipher::aes_256_cbc(), Mode::Encrypt, key, Some(&iv_array))
            .map_err(openssl_err)?;
        crypter.pad(false);
        let mut out = vec![0u8; data.len() + 16];
        let mut count = crypter.update(data, &mut out).map_err(openssl_err)?;
        count += crypter.finalize(&mut out[count..]).map_err(openssl_err)?;
        out.truncate(count);
        let final_iv = if out.is_empty() {
            iv_array
        } else {
            out[out.len() - 16..].try_into().unwrap()
        };
        Ok::<(Vec<u8>, [u8; 16]), PyErr>((out, final_iv))
    })?;

    if let Ok(iv_bytes_obj) = iv.cast::<PyByteArray>() {
        if iv_bytes_obj.len() == 16 {
            unsafe { iv_bytes_obj.as_bytes_mut().copy_from_slice(&final_iv); }
        }
    }

    Ok(PyBytes::new(py, &out_bytes))
}

/// AES-256-CBC Decryption
#[pyfunction]
#[pyo3(signature = (data, key, iv, /))]
pub fn cbc256_decrypt<'py>(
    py: Python<'py>,
    data: &[u8],
    key: &[u8],
    iv: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyBytes>> {
    if key.len() != 32 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Key must be 32 bytes"));
    }
    let iv_array: [u8; 16] = iv.extract()?;
    if data.len() % 16 != 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("Data length must be multiple of 16"));
    }

    let (out_bytes, final_iv) = py.detach(|| {
        let mut crypter = Crypter::new(Cipher::aes_256_cbc(), Mode::Decrypt, key, Some(&iv_array))
            .map_err(openssl_err)?;
        crypter.pad(false);
        let mut out = vec![0u8; data.len() + 16];
        let mut count = crypter.update(data, &mut out).map_err(openssl_err)?;
        count += crypter.finalize(&mut out[count..]).map_err(openssl_err)?;
        out.truncate(count);
        let final_iv = if data.is_empty() {
            iv_array
        } else {
            data[data.len() - 16..].try_into().unwrap()
        };
        Ok::<(Vec<u8>, [u8; 16]), PyErr>((out, final_iv))
    })?;

    if let Ok(iv_bytes_obj) = iv.cast::<PyByteArray>() {
        if iv_bytes_obj.len() == 16 {
            unsafe { iv_bytes_obj.as_bytes_mut().copy_from_slice(&final_iv); }
        }
    }

    Ok(PyBytes::new(py, &out_bytes))
}
