use crate::buffer::PyBufferExt;

use pyo3::buffer::PyBuffer;
use pyo3::intern;
use pyo3::pyclass;
use pyo3::pyfunction;
use pyo3::pymethods;
use pyo3::types::PyAnyMethods;
use pyo3::types::PyDict;
use pyo3::Bound;
use pyo3::IntoPyObjectExt;
use pyo3::PyAny;
use pyo3::PyResult;
use pyo3::Python;

use std::fs::File;
use std::io::Read;
use std::io::Seek;
use std::mem::ManuallyDrop;

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;

#[cfg(unix)]
use std::os::unix::io::FromRawFd;

#[cfg(windows)]
use std::os::windows::io::FromRawHandle;

#[cfg(windows)]
unsafe extern "C" {
    fn _get_osfhandle(fd: i32) -> isize;
}

trait HexDigest {
    fn hexdigest(self) -> String;
}

impl HexDigest for u32 {
    #[cfg(target_arch = "x86_64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(8);

        unsafe {
            let table = _mm_setr_epi8(
                b'0' as i8, b'1' as i8, b'2' as i8, b'3' as i8, b'4' as i8, b'5' as i8, b'6' as i8, b'7' as i8,
                b'8' as i8, b'9' as i8, b'a' as i8, b'b' as i8, b'c' as i8, b'd' as i8, b'e' as i8, b'f' as i8,
            );

            let input = _mm_cvtsi32_si128(self as i32);
            let mask = _mm_set1_epi8(0x0F);
            let lo = _mm_and_si128(input, mask);
            let hi = _mm_and_si128(_mm_srli_epi16(input, 4), mask);
            _mm_storel_epi64(
                hex.as_mut_ptr().cast(),
                _mm_shuffle_epi8(table, _mm_unpacklo_epi8(hi, lo)),
            );

            hex.set_len(8);
            String::from_utf8_unchecked(hex)
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(8);

        unsafe {
            let table = vld1q_u8(b"0123456789abcdef".as_ptr());
            let input = vcombine_u8(vcreate_u8(self as u64), vcreate_u8(0));
            let hi = vshrq_n_u8(input, 4);
            let lo = vandq_u8(input, vdupq_n_u8(0x0F));
            vst1_u8(hex.as_mut_ptr(), vget_low_u8(vqtbl1q_u8(table, vzip1q_u8(hi, lo))));

            hex.set_len(8);
            String::from_utf8_unchecked(hex)
        }
    }
}

impl HexDigest for u64 {
    #[cfg(target_arch = "x86_64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(16);

        unsafe {
            let table = _mm_setr_epi8(
                b'0' as i8, b'1' as i8, b'2' as i8, b'3' as i8, b'4' as i8, b'5' as i8, b'6' as i8, b'7' as i8,
                b'8' as i8, b'9' as i8, b'a' as i8, b'b' as i8, b'c' as i8, b'd' as i8, b'e' as i8, b'f' as i8,
            );

            let input = _mm_cvtsi64_si128(self as i64);
            let mask = _mm_set1_epi8(0x0F);
            let lo = _mm_and_si128(input, mask);
            let hi = _mm_and_si128(_mm_srli_epi16(input, 4), mask);

            _mm_storeu_si128(
                hex.as_mut_ptr().cast(),
                _mm_shuffle_epi8(table, _mm_unpacklo_epi8(hi, lo)),
            );

            hex.set_len(16);
            String::from_utf8_unchecked(hex)
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(16);

        unsafe {
            let table = vld1q_u8(b"0123456789abcdef".as_ptr());
            let input = vcombine_u8(vcreate_u8(self), vcreate_u8(0));
            let hi = vshrq_n_u8(input, 4);
            let lo = vandq_u8(input, vdupq_n_u8(0x0F));
            vst1q_u8(hex.as_mut_ptr(), vqtbl1q_u8(table, vzip1q_u8(hi, lo)));

            hex.set_len(16);
            String::from_utf8_unchecked(hex)
        }
    }
}

impl HexDigest for u128 {
    #[cfg(target_arch = "x86_64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(32);

        unsafe {
            let table = _mm_setr_epi8(
                b'0' as i8, b'1' as i8, b'2' as i8, b'3' as i8, b'4' as i8, b'5' as i8, b'6' as i8, b'7' as i8,
                b'8' as i8, b'9' as i8, b'a' as i8, b'b' as i8, b'c' as i8, b'd' as i8, b'e' as i8, b'f' as i8,
            );

            let input = _mm_set_epi64x((self >> 64) as i64, self as i64);
            let mask = _mm_set1_epi8(0x0F);
            let lo = _mm_and_si128(input, mask);
            let hi = _mm_and_si128(_mm_srli_epi16(input, 4), mask);
            let buffer = hex.as_mut_ptr();
            _mm_storeu_si128(buffer.cast(), _mm_shuffle_epi8(table, _mm_unpacklo_epi8(hi, lo)));
            _mm_storeu_si128(
                buffer.add(16).cast(),
                _mm_shuffle_epi8(table, _mm_unpackhi_epi8(hi, lo)),
            );

            hex.set_len(32);
            String::from_utf8_unchecked(hex)
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[inline(always)]
    fn hexdigest(self) -> String {
        let mut hex = Vec::<u8>::with_capacity(32);

        unsafe {
            let table = vld1q_u8(b"0123456789abcdef".as_ptr());
            let input = vcombine_u8(vcreate_u8(self as u64), vcreate_u8((self >> 64) as u64));
            let hi = vshrq_n_u8(input, 4);
            let lo = vandq_u8(input, vdupq_n_u8(0x0F));
            let buffer = hex.as_mut_ptr();
            vst1q_u8(buffer, vqtbl1q_u8(table, vzip1q_u8(hi, lo)));
            vst1q_u8(buffer.add(16), vqtbl1q_u8(table, vzip2q_u8(hi, lo)));

            hex.set_len(32);
            String::from_utf8_unchecked(hex)
        }
    }
}

#[cfg_attr(Py_3_10, pyclass(name = "HASH", module = "_hashlib", immutable_type))]
#[cfg_attr(not(Py_3_10), pyclass(name = "HASH", module = "_hashlib"))]
pub(crate) struct GxHashLib32 {
    seed: i64,
    buffer: PyBuffer<u8>,
}

#[cfg_attr(Py_3_10, pyclass(name = "HASH", module = "_hashlib", immutable_type))]
#[cfg_attr(not(Py_3_10), pyclass(name = "HASH", module = "_hashlib"))]
pub(crate) struct GxHashLib64 {
    seed: i64,
    buffer: PyBuffer<u8>,
}

#[cfg_attr(Py_3_10, pyclass(name = "HASH", module = "_hashlib", immutable_type))]
#[cfg_attr(not(Py_3_10), pyclass(name = "HASH", module = "_hashlib"))]
pub(crate) struct GxHashLib128 {
    seed: i64,
    buffer: PyBuffer<u8>,
}

macro_rules! impl_hashlib {
    ($name:ident, $function_name:ident, $digest_size:expr, $hasher:path) => {
        #[pymethods]
        impl $name {
            #[getter(__class__)]
            fn __class__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
                py.import(intern!(py, "_hashlib"))?.getattr(intern!(py, "HASH"))
            }

            #[getter]
            fn name(&self) -> &'static str {
                stringify!($function_name)
            }

            #[getter]
            fn digest_size(&self) -> usize {
                $digest_size
            }

            #[getter]
            fn block_size(&self) -> usize {
                1
            }

            fn digest(&self) -> [u8; $digest_size] {
                $hasher(self.buffer.as_bytes(), self.seed).to_le_bytes()
            }

            fn hexdigest(&self) -> String {
                $hasher(self.buffer.as_bytes(), self.seed).hexdigest()
            }

            fn update(&mut self, py: Python, data: PyBuffer<u8>) -> PyResult<()> {
                let slice = self.buffer.as_bytes();
                let new_slice = data.as_bytes();

                let mut combined = Vec::with_capacity(slice.len() + new_slice.len());
                combined.extend_from_slice(slice);
                combined.extend_from_slice(new_slice);
                self.buffer = PyBuffer::get(&combined.into_bound_py_any(py)?)?;

                Ok(())
            }

            fn copy(&self, py: Python) -> PyResult<Self> {
                let new_hashlib = Self {
                    seed: self.seed,
                    buffer: PyBuffer::get(&self.buffer.as_bytes().into_bound_py_any(py)?)?,
                };

                Ok(new_hashlib)
            }
        }

        #[pyfunction]
        #[pyo3(signature = (data = None, *, seed = 0, **_kwargs))]
        fn $function_name(
            py: Python<'_>,
            data: Option<PyBuffer<u8>>,
            seed: i64,
            _kwargs: Option<Bound<'_, PyDict>>,
        ) -> PyResult<$name> {
            let buffer = data.map_or_else(|| PyBuffer::get(&b"".into_bound_py_any(py)?), Ok)?;
            Ok($name { seed, buffer })
        }
    };
}

#[pyfunction]
#[pyo3(signature = (name, data = None, *, seed = 0, **kwargs))]
fn new<'py>(
    py: Python<'py>,
    name: &str,
    data: Option<PyBuffer<u8>>,
    seed: i64,
    kwargs: Option<Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
    match name {
        "gxhash32" => gxhash32(py, data, seed, kwargs)?.into_bound_py_any(py),
        "gxhash64" => gxhash64(py, data, seed, kwargs)?.into_bound_py_any(py),
        "gxhash128" => gxhash128(py, data, seed, kwargs)?.into_bound_py_any(py),
        _ => Err(pyo3::exceptions::PyValueError::new_err(format!(
            "unsupported hash type: {name}",
        ))),
    }
}

#[pyfunction]
#[pyo3(signature = (fileobj, digest, /, *, seed = 0, **kwargs))]
fn file_digest<'py>(
    py: Python<'py>,
    fileobj: Bound<'py, PyAny>,
    digest: Bound<'py, PyAny>,
    seed: i64,
    kwargs: Option<Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
    let data = fileobj.call_method0(intern!(py, "getbuffer")).or_else(|_| {
        let mut file = ManuallyDrop::new(unsafe {
            let fileno = fileobj.call_method0(intern!(py, "fileno"))?.extract::<i32>()?;

            #[cfg(unix)]
            {
                File::from_raw_fd(fileno)
            }
            #[cfg(windows)]
            {
                File::from_raw_handle(_get_osfhandle(fileno) as *mut std::ffi::c_void)
            }
        });

        let bytes = pyo3::types::PyBytes::new_with(
            py,
            (file.metadata()?.len() - file.stream_position()?) as usize,
            |buffer| py.detach(|| (&*file).read_exact(buffer).map_err(Into::into)),
        );

        bytes.map(Bound::into_any)
    })?;

    match digest.cast::<pyo3::types::PyString>() {
        Ok(name) => new(py, name.extract()?, Some(PyBuffer::get(&data)?), seed, kwargs),
        _ => digest
            .call0()
            .and_then(|hasher| hasher.call_method1(intern!(py, "update"), (&data,)).map(|_| hasher)),
    }
}

impl_hashlib!(GxHashLib32, gxhash32, 4, gxhash_core::gxhash32);
impl_hashlib!(GxHashLib64, gxhash64, 8, gxhash_core::gxhash64);
impl_hashlib!(GxHashLib128, gxhash128, 16, gxhash_core::gxhash128);

/// hashlib-compatible GxHash API
///
/// This module contains the hashlib-compatible API for GxHash.
///
/// * gxhash32  - a hashlib-compatible class for computing 32-bit hashes
/// * gxhash64  - a hashlib-compatible class for computing 64-bit hashes
/// * gxhash128 - a hashlib-compatible class for computing 128-bit hashes
///
/// The functions provide a compatible interface with Python's built-in hashlib module.
///
/// * gxhash32(data: Buffer = b"", *, seed: int = 0, usedforsecurity: bool = False) -> HASH
/// * gxhash64(data: Buffer = b"", *, seed: int = 0, usedforsecurity: bool = False) -> HASH
/// * gxhash128(data: Buffer = b"", *, seed: int = 0, usedforsecurity: bool = False) -> HASH
/// * new(name: str, data: str | Buffer = b"", *, seed: int = 0, usedforsecurity: bool = False) -> HASH
/// * file_digest(fileobj: BytesIOLike | FileLike, digest: str | Callable[[], HASH], /, *, seed: int = 0) -> HASH
///
/// The HASH objects returned by these functions again provide the standard HASH methods and properties.
///
/// * name -> str
/// * digest_size -> int
/// * block_size -> int
/// * digest() -> bytes
/// * hexdigest() -> str
/// * update(data: bytes) -> None
/// * copy() -> HASH
///
#[pyo3::pymodule(submodule, name = "gxhashlib", gil_used = false)]
pub mod hashlib_module {
    use pyo3::types;
    use pyo3::types::PyModuleMethods;

    #[pymodule_export]
    use super::file_digest;
    #[pymodule_export]
    use super::gxhash128;
    #[pymodule_export]
    use super::gxhash32;
    #[pymodule_export]
    use super::gxhash64;
    #[pymodule_export]
    use super::new;

    #[pymodule_init]
    fn init(m: &pyo3::Bound<'_, types::PyModule>) -> pyo3::PyResult<()> {
        let py = m.py();
        let algorithms_available = types::PySet::new(py, ["gxhash32", "gxhash64", "gxhash128"])?;

        m.add("algorithms_available", &algorithms_available)?;
        m.add("algorithms_guaranteed", &algorithms_available)
    }
}
