use std::io::{Read, Write};

use foyer::{
    BlockEngineConfig, Cache as FoyerCache, CacheBuilder, Code, DeviceBuilder, Error as FoyerError,
    ErrorKind as FoyerErrorKind, FsDeviceBuilder, HybridCache as FoyerHybridCache,
    HybridCacheBuilder, Result as FoyerResult,
};
use pyo3::exceptions::{PyOverflowError, PyRuntimeError, PyTypeError};
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyBytes, PyInt, PyModule, PyString};
use tokio::runtime::{Builder as TokioRuntimeBuilder, Runtime};

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum PyCacheKey {
    Bytes(Vec<u8>),
    Str(String),
    Int(i128),
}

impl Code for PyCacheKey {
    fn encode(&self, writer: &mut impl Write) -> FoyerResult<()> {
        match self {
            Self::Bytes(v) => {
                0u8.encode(writer)?;
                v.encode(writer)
            }
            Self::Str(v) => {
                1u8.encode(writer)?;
                v.encode(writer)
            }
            Self::Int(v) => {
                2u8.encode(writer)?;
                v.encode(writer)
            }
        }
    }

    fn decode(reader: &mut impl Read) -> FoyerResult<Self> {
        let tag = u8::decode(reader)?;
        match tag {
            0 => Ok(Self::Bytes(Vec::<u8>::decode(reader)?)),
            1 => Ok(Self::Str(String::decode(reader)?)),
            2 => Ok(Self::Int(i128::decode(reader)?)),
            _ => Err(FoyerError::new(
                FoyerErrorKind::Parse,
                format!("invalid PyCacheKey tag: {tag}"),
            )),
        }
    }

    fn estimated_size(&self) -> usize {
        match self {
            Self::Bytes(v) => 1 + v.estimated_size(),
            Self::Str(v) => 1 + v.estimated_size(),
            Self::Int(v) => 1 + v.estimated_size(),
        }
    }
}

fn parse_key(key: &Bound<'_, PyAny>) -> PyResult<PyCacheKey> {
    if let Ok(bytes) = key.downcast::<PyBytes>() {
        return Ok(PyCacheKey::Bytes(bytes.as_bytes().to_vec()));
    }
    if key.downcast::<PyString>().is_ok() {
        return Ok(PyCacheKey::Str(key.extract::<String>()?));
    }
    if key.is_instance_of::<PyInt>() {
        return key
            .extract::<i128>()
            .map(PyCacheKey::Int)
            .map_err(|_| PyOverflowError::new_err("int key is out of supported i128 range"));
    }
    Err(PyTypeError::new_err("key must be bytes, str, or int"))
}

fn parse_value(value: &Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
    value
        .downcast::<PyBytes>()
        .map(|bytes| bytes.as_bytes().to_vec())
        .map_err(|_| PyTypeError::new_err("value must be bytes"))
}

fn map_runtime_err(err: impl std::fmt::Display) -> PyErr {
    PyRuntimeError::new_err(err.to_string())
}

#[pyclass(name = "Cache")]
struct PyCache {
    inner: FoyerCache<PyCacheKey, Vec<u8>>,
}

#[pymethods]
impl PyCache {
    #[new]
    fn new(capacity: usize) -> Self {
        Self {
            inner: CacheBuilder::new(capacity).build(),
        }
    }

    fn insert(&self, key: &Bound<'_, PyAny>, value: &Bound<'_, PyAny>) -> PyResult<()> {
        let key = parse_key(key)?;
        let value = parse_value(value)?;
        self.inner.insert(key, value);
        Ok(())
    }

    fn get<'py>(&self, py: Python<'py>, key: &Bound<'py, PyAny>) -> PyResult<Option<Py<PyBytes>>> {
        let key = parse_key(key)?;
        let value = self
            .inner
            .get(&key)
            .map(|entry| PyBytes::new(py, entry.value()).into());
        Ok(value)
    }

    fn remove<'py>(
        &self,
        py: Python<'py>,
        key: &Bound<'py, PyAny>,
    ) -> PyResult<Option<Py<PyBytes>>> {
        let key = parse_key(key)?;
        let value = self
            .inner
            .remove(&key)
            .map(|entry| PyBytes::new(py, entry.value()).into());
        Ok(value)
    }

    fn contains(&self, key: &Bound<'_, PyAny>) -> PyResult<bool> {
        let key = parse_key(key)?;
        Ok(self.inner.contains(&key))
    }

    fn clear(&self) {
        self.inner.clear();
    }

    fn __len__(&self) -> usize {
        self.inner.entries()
    }

    #[getter]
    fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    #[getter]
    fn usage(&self) -> usize {
        self.inner.usage()
    }

    #[getter]
    fn entries(&self) -> usize {
        self.inner.entries()
    }
}

#[pyclass(unsendable)]
struct HybridCache {
    runtime: Runtime,
    inner: FoyerHybridCache<PyCacheKey, Vec<u8>>,
}

#[pymethods]
impl HybridCache {
    #[new]
    #[pyo3(signature = (storage_path, memory_capacity = 64 * 1024 * 1024, storage_capacity = 256 * 1024 * 1024))]
    fn new(
        storage_path: String,
        memory_capacity: usize,
        storage_capacity: usize,
    ) -> PyResult<Self> {
        let runtime = TokioRuntimeBuilder::new_multi_thread()
            .enable_all()
            .build()
            .map_err(map_runtime_err)?;

        let device = FsDeviceBuilder::new(storage_path)
            .with_capacity(storage_capacity)
            .build()
            .map_err(map_runtime_err)?;

        let inner = runtime
            .block_on(async {
                HybridCacheBuilder::new()
                    .memory(memory_capacity)
                    .storage()
                    .with_engine_config(BlockEngineConfig::new(device))
                    .build()
                    .await
            })
            .map_err(map_runtime_err)?;

        Ok(Self { runtime, inner })
    }

    fn insert(&self, key: &Bound<'_, PyAny>, value: &Bound<'_, PyAny>) -> PyResult<()> {
        let key = parse_key(key)?;
        let value = parse_value(value)?;
        self.inner.insert(key, value);
        Ok(())
    }

    fn get<'py>(&self, py: Python<'py>, key: &Bound<'py, PyAny>) -> PyResult<Option<Py<PyBytes>>> {
        let key = parse_key(key)?;
        let value = self
            .runtime
            .block_on(self.inner.get(&key))
            .map_err(map_runtime_err)?
            .map(|entry| PyBytes::new(py, entry.value()).into());
        Ok(value)
    }

    fn remove(&self, key: &Bound<'_, PyAny>) -> PyResult<()> {
        let key = parse_key(key)?;
        self.inner.remove(&key);
        Ok(())
    }

    fn contains(&self, key: &Bound<'_, PyAny>) -> PyResult<bool> {
        let key = parse_key(key)?;
        Ok(self.inner.contains(&key))
    }

    fn clear(&self) -> PyResult<()> {
        self.runtime
            .block_on(self.inner.clear())
            .map_err(map_runtime_err)
    }

    fn close(&self) -> PyResult<()> {
        self.runtime
            .block_on(self.inner.close())
            .map_err(map_runtime_err)
    }

    fn is_hybrid(&self) -> bool {
        self.inner.is_hybrid()
    }

    fn __len__(&self) -> usize {
        self.inner.memory().entries()
    }

    #[getter]
    fn capacity(&self) -> usize {
        self.inner.memory().capacity()
    }

    #[getter]
    fn usage(&self) -> usize {
        self.inner.memory().usage()
    }

    #[getter]
    fn entries(&self) -> usize {
        self.inner.memory().entries()
    }
}

#[pymodule]
fn _foyer(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyCache>()?;
    m.add_class::<HybridCache>()?;
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    Ok(())
}
