"""Streamers in MLC LLM."""

from typing import List, Union  # noqa: UP035

import tvm_ffi
from tvm.runtime import Object
from tvm_ffi import Shape

from . import _ffi_api
from .tokenizers import Tokenizer


@tvm_ffi.register_object("mlc.TextStreamer")
class TextStreamer(Object):
    """The class that streams back validated utf-8 text strings
    that generated by tokenizer.
    """

    def __init__(self, tokenizer: Tokenizer) -> None:
        """Create the text streamer from tokenizer"""
        self.__init_handle_by_constructor__(
            _ffi_api.TextStreamer,
            tokenizer,
        )

    def put(self, delta_tokens: Union[List[int], Shape]) -> str:  # noqa: UP006
        """Put new delta tokens into the streamer, and get the UTF-8-valid
        delta string. The text streamer may hold some of the input delta tokens
        which cannot decode into valid UTF-8 strings. The returned string
        is always guaranteed to be UTF-8 valid.

        Parameters
        ----------
        delta_tokens : Union[List[int], Shape]
            The new tokens to put into the streamer.

        Returns
        -------
        delta_text : str
            The decoded delta string after putting the input new tokens.
        """
        if isinstance(delta_tokens, list):
            delta_tokens = Shape(delta_tokens)
        return _ffi_api.TextStreamerPut(self, delta_tokens)

    def finish(self) -> str:
        """Return the string decoded by remaining tokens."""
        return _ffi_api.TextStreamerFinish(self)


@tvm_ffi.register_object("mlc.StopStrHandler")
class StopStrHandler(Object):
    """The stop string handler in MLC LLM, which takes input delta tokens
    one at a time, and return the output delta token before stopping due to
    stop strings."""

    def __init__(
        self,
        stop_strs: List[str],  # noqa: UP006
        tokenizer: Tokenizer,
    ) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.StopStrHandler,
            stop_strs,
            tokenizer,
        )

    def put(self, token_id: int) -> List[int]:  # noqa: UP006
        """Add new input delta token to the handler, return output
        delta tokens before stopping. The stop string handler may hold
        some of the input delta token which may be part of a stop string.
        The returned tokens are always guaranteed not to be part of stop string.
        """
        return list(_ffi_api.StopStrHandlerPut(self, token_id))

    def finish(self) -> List[int]:  # noqa: UP006
        """Stop string handling has finished, return remaining cached token ids."""
        return list(_ffi_api.StopStringHandlerFinish(self))

    @property
    def stop_triggered(self) -> bool:
        """Check if the generation has stopped due to stop string."""
        return _ffi_api.StopStrHandlerStopTriggered(self)