use std::sync::{Arc, RwLock};

use crate::pre_tokenizers::from_string;
use crate::tokenizer::PyTokenizer;
use crate::utils::PyPattern;
use pyo3::exceptions;
use pyo3::prelude::*;
use pyo3::types::*;
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use tk::decoders::bpe::BPEDecoder;
use tk::decoders::byte_fallback::ByteFallback;
use tk::decoders::byte_level::ByteLevel;
use tk::decoders::ctc::CTC;
use tk::decoders::fuse::Fuse;
use tk::decoders::metaspace::{Metaspace, PrependScheme};
use tk::decoders::sequence::Sequence;
use tk::decoders::strip::Strip;
use tk::decoders::wordpiece::WordPiece;
use tk::decoders::DecoderWrapper;
use tk::normalizers::replace::Replace;
use tk::Decoder;
use tokenizers as tk;

use super::error::ToPyResult;

/// Base class for all decoders
///
/// This class is not supposed to be instantiated directly. Instead, any implementation of
/// a Decoder will return an instance of this class when instantiated.
#[pyclass(
    dict,
    module = "tokenizers.decoders",
    name = "Decoder",
    subclass,
    from_py_object
)]
#[derive(Clone, Deserialize, Serialize)]
#[serde(transparent)]
pub struct PyDecoder {
    pub(crate) decoder: PyDecoderWrapper,
}

impl PyDecoder {
    pub(crate) fn new(decoder: PyDecoderWrapper) -> Self {
        PyDecoder { decoder }
    }

    pub(crate) fn get_as_subtype(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let base = self.clone();
        Ok(match &self.decoder {
            PyDecoderWrapper::Custom(_) => Py::new(py, base)?.into_any(),
            PyDecoderWrapper::Wrapped(inner) => match &*inner.as_ref().read().unwrap() {
                DecoderWrapper::Metaspace(_) => Py::new(py, (PyMetaspaceDec {}, base))?.into_any(),
                DecoderWrapper::WordPiece(_) => Py::new(py, (PyWordPieceDec {}, base))?.into_any(),
                DecoderWrapper::ByteFallback(_) => {
                    Py::new(py, (PyByteFallbackDec {}, base))?.into_any()
                }
                DecoderWrapper::Strip(_) => Py::new(py, (PyStrip {}, base))?.into_any(),
                DecoderWrapper::Fuse(_) => Py::new(py, (PyFuseDec {}, base))?.into_any(),
                DecoderWrapper::ByteLevel(_) => Py::new(py, (PyByteLevelDec {}, base))?.into_any(),
                DecoderWrapper::Replace(_) => Py::new(py, (PyReplaceDec {}, base))?.into_any(),
                DecoderWrapper::BPE(_) => Py::new(py, (PyBPEDecoder {}, base))?.into_any(),
                DecoderWrapper::CTC(_) => Py::new(py, (PyCTCDecoder {}, base))?.into_any(),
                DecoderWrapper::Sequence(_) => {
                    Py::new(py, (PySequenceDecoder {}, base))?.into_any()
                }
            },
        })
    }
}

impl Decoder for PyDecoder {
    fn decode_chain(&self, tokens: Vec<String>) -> tk::Result<Vec<String>> {
        self.decoder.decode_chain(tokens)
    }
}

#[pymethods]
impl PyDecoder {
    #[staticmethod]
    #[pyo3(text_signature = "(decoder)")]
    fn custom(decoder: Py<PyAny>) -> Self {
        let decoder = PyDecoderWrapper::Custom(Arc::new(RwLock::new(CustomDecoder::new(decoder))));
        PyDecoder::new(decoder)
    }

    fn __getstate__(&self, py: Python) -> PyResult<Py<PyAny>> {
        let data = serde_json::to_string(&self.decoder).map_err(|e| {
            exceptions::PyException::new_err(format!(
                "Error while attempting to pickle Decoder: {e}"
            ))
        })?;
        Ok(PyBytes::new(py, data.as_bytes()).into())
    }

    fn __setstate__(&mut self, py: Python, state: Py<PyAny>) -> PyResult<()> {
        match state.extract::<&[u8]>(py) {
            Ok(s) => {
                self.decoder = serde_json::from_slice(s).map_err(|e| {
                    exceptions::PyException::new_err(format!(
                        "Error while attempting to unpickle Decoder: {e}"
                    ))
                })?;
                Ok(())
            }
            Err(e) => Err(e.into()),
        }
    }

    /// Decode the given list of tokens to a final string
    ///
    /// Args:
    ///     tokens (:obj:`List[str]`):
    ///         The list of tokens to decode
    ///
    /// Returns:
    ///     :obj:`str`: The decoded string
    #[pyo3(text_signature = "(self, tokens)")]
    fn decode(&self, tokens: Vec<String>) -> PyResult<String> {
        ToPyResult(self.decoder.decode(tokens)).into()
    }

    fn __repr__(&self) -> PyResult<String> {
        crate::utils::serde_pyo3::repr(self)
            .map_err(|e| exceptions::PyException::new_err(e.to_string()))
    }

    fn __str__(&self) -> PyResult<String> {
        crate::utils::serde_pyo3::to_string(self)
            .map_err(|e| exceptions::PyException::new_err(e.to_string()))
    }
}

macro_rules! getter {
    ($self: ident, $variant: ident, $($name: tt)+) => {{
        let super_ = $self.as_ref();
        if let PyDecoderWrapper::Wrapped(ref wrap) = super_.decoder {
            if let DecoderWrapper::$variant(ref dec) = *wrap.read().unwrap() {
                dec.$($name)+
            } else {
                unreachable!()
            }
        } else {
            unreachable!()
        }
    }};
}

macro_rules! setter {
    ($self: ident, $variant: ident, $name: ident, $value: expr) => {{
        let super_ = $self.as_ref();
        if let PyDecoderWrapper::Wrapped(ref wrap) = super_.decoder {
            if let DecoderWrapper::$variant(ref mut dec) = *wrap.write().unwrap() {
                dec.$name = $value;
            }
        }
    }};
    ($self: ident, $variant: ident, @$name: ident, $value: expr) => {{
        let super_ = $self.as_ref();
        if let PyDecoderWrapper::Wrapped(ref wrap) = super_.decoder {
            if let DecoderWrapper::$variant(ref mut dec) = *wrap.write().unwrap() {
                dec.$name($value);
            }
        }
    }};
}

/// ByteLevel Decoder
///
/// This decoder is to be used in tandem with the
/// :class:`~tokenizers.pre_tokenizers.ByteLevel` pre-tokenizer. It reverses the
/// byte-to-unicode mapping applied during pre-tokenization, converting the special
/// Unicode characters back into the original bytes to reconstruct the original string.
///
/// Example::
///
///     >>> from tokenizers.decoders import ByteLevel
///     >>> decoder = ByteLevel()
///     >>> decoder.decode(["ĠHello", "Ġworld"])
///     ' Hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "ByteLevel")]
pub struct PyByteLevelDec {}
#[pymethods]
impl PyByteLevelDec {
    #[new]
    #[pyo3(signature = (**_kwargs), text_signature = "(self)")]
    fn new(_kwargs: Option<&Bound<'_, PyDict>>) -> (Self, PyDecoder) {
        (PyByteLevelDec {}, ByteLevel::default().into())
    }
}

/// Replace Decoder
///
/// This decoder is to be used in tandem with the
/// :class:`~tokenizers.normalizers.Replace` normalizer or a similar replace operation.
/// It reverses a string replacement by substituting the replacement content back
/// with the original pattern.
///
/// Args:
///     pattern (:obj:`str` or :class:`~tokenizers.Regex`):
///         The pattern that was used as the replacement target during encoding.
///
///     content (:obj:`str`):
///         The string to replace each match of the pattern with during decoding.
///
/// Example::
///
///     >>> from tokenizers.decoders import Replace
///     >>> decoder = Replace("▁", " ")
///     >>> decoder.decode(["▁Hello", "▁world"])
///     ' Hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "Replace")]
pub struct PyReplaceDec {}
#[pymethods]
impl PyReplaceDec {
    #[new]
    #[pyo3(text_signature = "(self, pattern, content)")]
    fn new(pattern: PyPattern, content: String) -> PyResult<(Self, PyDecoder)> {
        Ok((
            PyReplaceDec {},
            ToPyResult(Replace::new(pattern, content)).into_py()?.into(),
        ))
    }
}

/// WordPiece Decoder
///
/// Args:
///     prefix (:obj:`str`, `optional`, defaults to :obj:`##`):
///         The prefix to use for subwords that are not a beginning-of-word
///
///     cleanup (:obj:`bool`, `optional`, defaults to :obj:`True`):
///         Whether to cleanup some tokenization artifacts. Mainly spaces before punctuation,
///         and some abbreviated english forms.
///
/// Example::
///
///     >>> from tokenizers.decoders import WordPiece
///     >>> decoder = WordPiece()
///     >>> decoder.decode(["Hello", ",", "##world", "!"])
///     'Hello, world!'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "WordPiece")]
pub struct PyWordPieceDec {}
#[pymethods]
impl PyWordPieceDec {
    #[getter]
    fn get_prefix(self_: PyRef<Self>) -> String {
        getter!(self_, WordPiece, prefix.clone())
    }

    #[setter]
    fn set_prefix(self_: PyRef<Self>, prefix: String) {
        setter!(self_, WordPiece, prefix, prefix);
    }

    #[getter]
    fn get_cleanup(self_: PyRef<Self>) -> bool {
        getter!(self_, WordPiece, cleanup)
    }

    #[setter]
    fn set_cleanup(self_: PyRef<Self>, cleanup: bool) {
        setter!(self_, WordPiece, cleanup, cleanup);
    }

    #[new]
    #[pyo3(signature = (prefix = String::from("##"), cleanup = true), text_signature = "(self, prefix=\"##\", cleanup=True)")]
    fn new(prefix: String, cleanup: bool) -> (Self, PyDecoder) {
        (PyWordPieceDec {}, WordPiece::new(prefix, cleanup).into())
    }
}

/// ByteFallback Decoder
///
/// ByteFallback is a decoder that handles tokens representing raw bytes in the
/// ``<0xNN>`` format (e.g., ``<0x61>`` for the byte ``0x61`` = ``'a'``). It converts
/// such tokens to their corresponding bytes and attempts to decode the resulting byte
/// sequence as UTF-8. This is used in LLaMA/SentencePiece models that use byte fallback
/// for unknown characters. Inconvertible byte tokens are replaced with the Unicode
/// replacement character (U+FFFD).
///
/// Example::
///
///     >>> from tokenizers.decoders import ByteFallback, Fuse, Sequence
///     >>> decoder = Sequence([ByteFallback(), Fuse()])
///     >>> decoder.decode(["<0x48>", "<0x65>", "<0x6C>", "<0x6C>", "<0x6F>"])
///     'Hello'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "ByteFallback")]
pub struct PyByteFallbackDec {}
#[pymethods]
impl PyByteFallbackDec {
    #[new]
    #[pyo3(signature = (), text_signature = "(self)")]
    fn new() -> (Self, PyDecoder) {
        (PyByteFallbackDec {}, ByteFallback::new().into())
    }
}

/// Fuse Decoder
///
/// Fuse simply concatenates every token into a single string without any separator.
/// This is typically the last step in a decoder chain when other decoders need to
/// operate on individual tokens before they are joined together.
///
/// Example::
///
///     >>> from tokenizers.decoders import Fuse
///     >>> decoder = Fuse()
///     >>> decoder.decode(["Hello", ",", " ", "world", "!"])
///     'Hello, world!'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "Fuse")]
pub struct PyFuseDec {}
#[pymethods]
impl PyFuseDec {
    #[new]
    #[pyo3(signature = (), text_signature = "(self)")]
    fn new() -> (Self, PyDecoder) {
        (PyFuseDec {}, Fuse::new().into())
    }
}

/// Strip Decoder
///
/// Strips a given number of occurrences of a character from the left and/or right
/// side of each token. This is useful for removing padding characters or special
/// prefix/suffix markers added during tokenization.
///
/// Args:
///     content (:obj:`str`, defaults to :obj:`" "`):
///         The character to strip from each token.
///
///     left (:obj:`int`, defaults to :obj:`0`):
///         The number of occurrences of :obj:`content` to remove from the left
///         side of each token.
///
///     right (:obj:`int`, defaults to :obj:`0`):
///         The number of occurrences of :obj:`content` to remove from the right
///         side of each token.
///
/// Example::
///
///     >>> from tokenizers.decoders import Strip
///     >>> decoder = Strip(content="▁", left=1)
///     >>> decoder.decode(["▁Hello", "▁world"])
///     'Hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "Strip")]
pub struct PyStrip {}
#[pymethods]
impl PyStrip {
    #[getter]
    fn get_start(self_: PyRef<Self>) -> usize {
        getter!(self_, Strip, start)
    }

    #[setter]
    fn set_start(self_: PyRef<Self>, start: usize) {
        setter!(self_, Strip, start, start)
    }

    #[getter]
    fn get_stop(self_: PyRef<Self>) -> usize {
        getter!(self_, Strip, stop)
    }

    #[setter]
    fn set_stop(self_: PyRef<Self>, stop: usize) {
        setter!(self_, Strip, stop, stop)
    }

    #[getter]
    fn get_content(self_: PyRef<Self>) -> char {
        getter!(self_, Strip, content)
    }

    #[setter]
    fn set_content(self_: PyRef<Self>, content: char) {
        setter!(self_, Strip, content, content)
    }

    #[new]
    #[pyo3(
        signature = (content=' ', left=0, right=0),
        text_signature = "(self, content=' ', left=0, right=0)"
    )]
    fn new(content: char, left: usize, right: usize) -> (Self, PyDecoder) {
        (PyStrip {}, Strip::new(content, left, right).into())
    }
}

/// Metaspace Decoder
///
/// Args:
///     replacement (:obj:`str`, `optional`, defaults to :obj:`▁`):
///         The replacement character. Must be exactly one character. By default we
///         use the `▁` (U+2581) meta symbol (Same as in SentencePiece).
///
///     prepend_scheme (:obj:`str`, `optional`, defaults to :obj:`"always"`):
///         Whether to add a space to the first word if there isn't already one. This
///         lets us treat `hello` exactly like `say hello`.
///         Choices: "always", "never", "first". First means the space is only added on the first
///         token (relevant when special tokens are used or other pre_tokenizer are used).
///
/// Example::
///
///     >>> from tokenizers.decoders import Metaspace
///     >>> decoder = Metaspace()
///     >>> decoder.decode(["▁Hello", "▁my", "▁friend"])
///     'Hello my friend'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "Metaspace")]
pub struct PyMetaspaceDec {}
#[pymethods]
impl PyMetaspaceDec {
    #[getter]
    fn get_replacement(self_: PyRef<Self>) -> String {
        getter!(self_, Metaspace, get_replacement().to_string())
    }

    #[setter]
    fn set_replacement(self_: PyRef<Self>, replacement: char) {
        setter!(self_, Metaspace, @set_replacement, replacement);
    }

    #[getter]
    fn get_split(self_: PyRef<Self>) -> bool {
        getter!(self_, Metaspace, get_split())
    }

    #[setter]
    fn set_split(self_: PyRef<Self>, split: bool) {
        setter!(self_, Metaspace, @set_split, split);
    }

    #[getter]
    fn get_prepend_scheme(self_: PyRef<Self>) -> String {
        // Assuming Metaspace has a method to get the prepend_scheme as a string
        let scheme: PrependScheme = getter!(self_, Metaspace, get_prepend_scheme());
        match scheme {
            PrependScheme::First => "first",
            PrependScheme::Never => "never",
            PrependScheme::Always => "always",
        }
        .to_string()
    }

    #[setter]
    fn set_prepend_scheme(self_: PyRef<Self>, prepend_scheme: String) -> PyResult<()> {
        let scheme = from_string(prepend_scheme)?;
        setter!(self_, Metaspace, @set_prepend_scheme, scheme);
        Ok(())
    }

    #[new]
    #[pyo3(signature = (replacement = '▁', prepend_scheme = String::from("always"), split = true), text_signature = "(self, replacement = \"▁\",  prepend_scheme = \"always\", split = True)")]
    fn new(replacement: char, prepend_scheme: String, split: bool) -> PyResult<(Self, PyDecoder)> {
        let prepend_scheme = from_string(prepend_scheme)?;
        Ok((
            PyMetaspaceDec {},
            Metaspace::new(replacement, prepend_scheme, split).into(),
        ))
    }
}

/// BPEDecoder Decoder
///
/// Args:
///     suffix (:obj:`str`, `optional`, defaults to :obj:`</w>`):
///         The suffix that was used to characterize an end-of-word. This suffix will
///         be replaced by whitespaces during the decoding
///
/// Example::
///
///     >>> from tokenizers.decoders import BPEDecoder
///     >>> decoder = BPEDecoder()
///     >>> decoder.decode(["Hello</w>", "world</w>"])
///     'Hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "BPEDecoder")]
pub struct PyBPEDecoder {}
#[pymethods]
impl PyBPEDecoder {
    #[getter]
    fn get_suffix(self_: PyRef<Self>) -> String {
        getter!(self_, BPE, suffix.clone())
    }

    #[setter]
    fn set_suffix(self_: PyRef<Self>, suffix: String) {
        setter!(self_, BPE, suffix, suffix);
    }

    #[new]
    #[pyo3(signature = (suffix = String::from("</w>")), text_signature = "(self, suffix=\"</w>\")")]
    fn new(suffix: String) -> (Self, PyDecoder) {
        (PyBPEDecoder {}, BPEDecoder::new(suffix).into())
    }
}

/// CTC Decoder
///
/// Args:
///     pad_token (:obj:`str`, `optional`, defaults to :obj:`<pad>`):
///         The pad token used by CTC to delimit a new token.
///     word_delimiter_token (:obj:`str`, `optional`, defaults to :obj:`|`):
///         The word delimiter token. It will be replaced by a <space>
///     cleanup (:obj:`bool`, `optional`, defaults to :obj:`True`):
///         Whether to cleanup some tokenization artifacts.
///         Mainly spaces before punctuation, and some abbreviated english forms.
///
/// Example::
///
///     >>> from tokenizers.decoders import CTC
///     >>> decoder = CTC()
///     >>> decoder.decode(["h", "e", "e", "<pad>", "l", "l", "o", "|", "w", "o", "r", "l", "d"])
///     'hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name = "CTC")]
pub struct PyCTCDecoder {}
#[pymethods]
impl PyCTCDecoder {
    #[getter]
    fn get_pad_token(self_: PyRef<Self>) -> String {
        getter!(self_, CTC, pad_token.clone())
    }

    #[setter]
    fn set_pad_token(self_: PyRef<Self>, pad_token: String) {
        setter!(self_, CTC, pad_token, pad_token);
    }

    #[getter]
    fn get_word_delimiter_token(self_: PyRef<Self>) -> String {
        getter!(self_, CTC, word_delimiter_token.clone())
    }

    #[setter]
    fn set_word_delimiter_token(self_: PyRef<Self>, word_delimiter_token: String) {
        setter!(self_, CTC, word_delimiter_token, word_delimiter_token);
    }

    #[getter]
    fn get_cleanup(self_: PyRef<Self>) -> bool {
        getter!(self_, CTC, cleanup)
    }

    #[setter]
    fn set_cleanup(self_: PyRef<Self>, cleanup: bool) {
        setter!(self_, CTC, cleanup, cleanup);
    }

    #[new]
    #[pyo3(signature = (
        pad_token = String::from("<pad>"),
        word_delimiter_token = String::from("|"),
        cleanup = true
    ),
        text_signature = "(self, pad_token=\"<pad>\", word_delimiter_token=\"|\", cleanup=True)")]
    fn new(pad_token: String, word_delimiter_token: String, cleanup: bool) -> (Self, PyDecoder) {
        (
            PyCTCDecoder {},
            CTC::new(pad_token, word_delimiter_token, cleanup).into(),
        )
    }
}

/// Sequence Decoder
///
/// Chains multiple decoders together, applying them in order. Each decoder in the
/// sequence processes the output of the previous one, allowing complex decoding
/// pipelines to be built from simpler components.
///
/// Args:
///     decoders (:obj:`List[Decoder]`):
///         The list of decoders to chain together.
///
/// Example::
///
///     >>> from tokenizers.decoders import ByteFallback, Fuse, Metaspace, Sequence
///     >>> decoder = Sequence([ByteFallback(), Fuse(), Metaspace()])
///     >>> decoder.decode(["▁Hello", "▁world"])
///     'Hello world'
///
#[pyclass(extends=PyDecoder, module = "tokenizers.decoders", name="Sequence")]
pub struct PySequenceDecoder {}
#[pymethods]
impl PySequenceDecoder {
    #[new]
    #[pyo3(signature = (decoders_py), text_signature = "(self, decoders)")]
    fn new(decoders_py: &Bound<'_, PyList>) -> PyResult<(Self, PyDecoder)> {
        let mut decoders: Vec<DecoderWrapper> = Vec::with_capacity(decoders_py.len());
        for decoder_py in decoders_py.iter() {
            let decoder: PyRef<PyDecoder> = decoder_py.extract()?;
            let decoder = match &decoder.decoder {
                PyDecoderWrapper::Wrapped(inner) => inner,
                PyDecoderWrapper::Custom(_) => unimplemented!(),
            };
            decoders.push(decoder.read().unwrap().clone());
        }
        Ok((PySequenceDecoder {}, Sequence::new(decoders).into()))
    }

    fn __getnewargs__<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyTuple>> {
        PyTuple::new(py, [PyList::empty(py)])
    }
}

pub(crate) struct CustomDecoder {
    inner: Py<PyAny>,
}

impl CustomDecoder {
    pub(crate) fn new(inner: Py<PyAny>) -> Self {
        CustomDecoder { inner }
    }
}

impl Decoder for CustomDecoder {
    fn decode(&self, tokens: Vec<String>) -> tk::Result<String> {
        Python::attach(|py| {
            let decoded = self
                .inner
                .call_method(py, "decode", (tokens,), None)?
                .extract(py)?;
            Ok(decoded)
        })
    }

    fn decode_chain(&self, tokens: Vec<String>) -> tk::Result<Vec<String>> {
        Python::attach(|py| {
            let decoded = self
                .inner
                .call_method(py, "decode_chain", (tokens,), None)?
                .extract(py)?;
            Ok(decoded)
        })
    }
}

impl Serialize for CustomDecoder {
    fn serialize<S>(&self, _serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        Err(serde::ser::Error::custom(
            "Custom PyDecoder cannot be serialized",
        ))
    }
}

impl<'de> Deserialize<'de> for CustomDecoder {
    fn deserialize<D>(_deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        Err(D::Error::custom("PyDecoder cannot be deserialized"))
    }
}

#[derive(Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub(crate) enum PyDecoderWrapper {
    Custom(Arc<RwLock<CustomDecoder>>),
    Wrapped(Arc<RwLock<DecoderWrapper>>),
}

impl<I> From<I> for PyDecoderWrapper
where
    I: Into<DecoderWrapper>,
{
    fn from(norm: I) -> Self {
        PyDecoderWrapper::Wrapped(Arc::new(RwLock::new(norm.into())))
    }
}

impl<I> From<I> for PyDecoder
where
    I: Into<DecoderWrapper>,
{
    fn from(dec: I) -> Self {
        PyDecoder {
            decoder: dec.into().into(),
        }
    }
}

impl Decoder for PyDecoderWrapper {
    fn decode_chain(&self, tokens: Vec<String>) -> tk::Result<Vec<String>> {
        match self {
            PyDecoderWrapper::Wrapped(inner) => inner.read().unwrap().decode_chain(tokens),
            PyDecoderWrapper::Custom(inner) => inner.read().unwrap().decode_chain(tokens),
        }
    }
}

/// Decoders Module
#[pymodule(gil_used = false)]
pub mod decoders {
    #[pymodule_export]
    pub use super::PyBPEDecoder;
    #[pymodule_export]
    pub use super::PyByteFallbackDec;
    #[pymodule_export]
    pub use super::PyByteLevelDec;
    #[pymodule_export]
    pub use super::PyCTCDecoder;
    #[pymodule_export]
    pub use super::PyDecodeStream;
    #[pymodule_export]
    pub use super::PyDecoder;
    #[pymodule_export]
    pub use super::PyFuseDec;
    #[pymodule_export]
    pub use super::PyMetaspaceDec;
    #[pymodule_export]
    pub use super::PyReplaceDec;
    #[pymodule_export]
    pub use super::PySequenceDecoder;
    #[pymodule_export]
    pub use super::PyStrip;
    #[pymodule_export]
    pub use super::PyWordPieceDec;
}

/// Provides incremental decoding of token IDs as they are generated, yielding
/// decoded text chunks as soon as they are available.
///
/// Unlike batch decoding, streaming decode is designed for use with autoregressive
/// generation — tokens arrive one at a time and the decoder needs to handle
/// multi-byte sequences (e.g., UTF-8 characters split across token boundaries) and
/// byte-fallback tokens gracefully.
///
/// The decoder internally buffers tokens until it can produce a valid UTF-8 string
/// chunk, then yields that chunk and advances its internal state. This means
/// individual calls to :meth:`~tokenizers.decoders.DecodeStream.step` may return
/// :obj:`None` when the current token completes a partial sequence that cannot yet
/// be decoded.
///
/// Args:
///     skip_special_tokens (:obj:`bool`, defaults to :obj:`False`):
///         Whether to skip special tokens (e.g. ``[CLS]``, ``[SEP]``, ``<s>``) when
///         decoding.
///
/// Example::
///
///     >>> from tokenizers import Tokenizer
///     >>> from tokenizers.decoders import DecodeStream
///     >>> tokenizer = Tokenizer.from_pretrained("gpt2")
///     >>> stream = DecodeStream(skip_special_tokens=True)
///     >>> # Simulate streaming token-by-token generation
///     >>> token_ids = tokenizer.encode("Hello, streaming world!").ids
///     >>> for token_id in token_ids:
///     ...     chunk = stream.step(tokenizer, token_id)
///     ...     if chunk is not None:
///     ...         print(chunk, end="", flush=True)
///
#[pyclass(module = "tokenizers.decoders", name = "DecodeStream", from_py_object)]
#[derive(Clone)]
pub struct PyDecodeStream {
    /// Regular decode option that is kept throughout.
    skip_special_tokens: bool,
    /// A temporary buffer of the necessary token_ids needed
    /// to produce valid string chunks.
    /// This typically contains 3 parts:
    ///  - read
    ///  - prefix
    ///  - rest
    ///
    /// Read is the bit necessary to surround the prefix
    /// so decoding the whole ids produces a valid prefix.
    /// Prefix is the previously produced string, kept around to trim off of
    /// the next valid chunk
    ids: Vec<u32>,
    /// The previously returned chunk that needs to be discarded from the
    /// decoding of the current ids to produce the next chunk
    prefix: String,
    /// The index within the ids corresponding to the prefix so we can drain
    /// correctly
    prefix_index: usize,
}

#[derive(Clone)]
enum StreamInput {
    Id(u32),
    Ids(Vec<u32>),
}

impl<'a, 'py> FromPyObject<'a, 'py> for StreamInput {
    type Error = PyErr;

    fn extract(obj: Borrowed<'a, 'py, PyAny>) -> Result<Self, Self::Error> {
        if let Ok(id) = obj.extract::<u32>() {
            Ok(StreamInput::Id(id))
        } else if let Ok(ids) = obj.extract::<Vec<u32>>() {
            Ok(StreamInput::Ids(ids))
        } else {
            Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                "StreamInput must be either an integer or a list of integers",
            ))
        }
    }
}

#[pymethods]
impl PyDecodeStream {
    #[new]
    #[pyo3(signature = (ids=None, skip_special_tokens=false), text_signature = "(self, ids=None, skip_special_tokens=False)")]
    fn new(ids: Option<Vec<u32>>, skip_special_tokens: Option<bool>) -> Self {
        PyDecodeStream {
            skip_special_tokens: skip_special_tokens.unwrap_or(false),
            ids: ids.unwrap_or_default(),
            prefix: String::new(),
            prefix_index: 0,
        }
    }

    /// Add the next token ID (or list of IDs) to the stream and return the next
    /// decoded text chunk if one is available.
    ///
    /// Because some characters span multiple tokens (e.g. multi-byte UTF-8
    /// sequences or byte-fallback tokens), this method may return :obj:`None`
    /// when the provided token does not yet complete a decodable unit. Callers
    /// should simply continue feeding tokens until a non-:obj:`None` value is
    /// returned.
    ///
    /// Args:
    ///     tokenizer (:class:`~tokenizers.Tokenizer`):
    ///         The tokenizer whose decoder pipeline will be used.
    ///
    ///     id (:obj:`int` or :obj:`List[int]`):
    ///         The next token ID, or a list of token IDs to append to the stream.
    ///
    /// Returns:
    ///     :obj:`Optional[str]`: The next decoded text chunk if enough tokens have
    ///         accumulated, or :obj:`None` if more tokens are still needed.
    #[pyo3(signature = (tokenizer, id), text_signature = "(self, tokenizer, id)")]
    fn step(&mut self, tokenizer: &PyTokenizer, id: StreamInput) -> PyResult<Option<String>> {
        let id: Vec<u32> = match id {
            StreamInput::Id(id) => vec![id],
            StreamInput::Ids(ids) => ids,
        };
        let tokenizer_guard = tokenizer.read_inner()?;
        ToPyResult(tk::tokenizer::step_decode_stream(
            &tokenizer_guard,
            id,
            self.skip_special_tokens,
            &mut self.ids,
            &mut self.prefix,
            &mut self.prefix_index,
        ))
        .into()
    }
    fn __copy__(&self) -> Self {
        self.clone()
    }

    fn __deepcopy__(&self, _memo: &Bound<'_, PyDict>) -> Self {
        self.clone()
    }
}

#[cfg(test)]
mod test {
    use std::sync::{Arc, RwLock};

    use pyo3::prelude::*;
    use tk::decoders::metaspace::Metaspace;
    use tk::decoders::DecoderWrapper;

    use crate::decoders::{CustomDecoder, PyDecoder, PyDecoderWrapper};

    #[test]
    fn get_subtype() {
        Python::attach(|py| {
            let py_dec = PyDecoder::new(Metaspace::default().into());
            let py_meta = py_dec.get_as_subtype(py).unwrap();
            assert_eq!("Metaspace", py_meta.bind(py).get_type().qualname().unwrap());
        })
    }

    #[test]
    fn serialize() {
        let py_wrapped: PyDecoderWrapper = Metaspace::default().into();
        let py_ser = serde_json::to_string(&py_wrapped).unwrap();
        let rs_wrapped = DecoderWrapper::Metaspace(Metaspace::default());
        let rs_ser = serde_json::to_string(&rs_wrapped).unwrap();
        assert_eq!(py_ser, rs_ser);
        let py_dec: PyDecoder = serde_json::from_str(&rs_ser).unwrap();
        match py_dec.decoder {
            PyDecoderWrapper::Wrapped(msp) => match *msp.as_ref().read().unwrap() {
                DecoderWrapper::Metaspace(_) => {}
                _ => panic!("Expected Metaspace"),
            },
            _ => panic!("Expected wrapped, not custom."),
        }

        let obj = Python::attach(|py| {
            let py_msp = PyDecoder::new(Metaspace::default().into());
            Py::new(py, py_msp).unwrap().into_any()
        });
        let py_seq = PyDecoderWrapper::Custom(Arc::new(RwLock::new(CustomDecoder::new(obj))));
        assert!(serde_json::to_string(&py_seq).is_err());
    }
}
