from abc import ABC, abstractmethod


class ModelPostProcessor(ABC):
    """
    Abstract base class for model post-processors.

    Subclasses should implement the `apply` method to define how the generated
    text from a model should be modified or cleaned after generation.
    """

    @abstractmethod
    def apply(self, generated_text: str) -> str:
        """
        Apply a post-processing operation to the given generated text.

        Args:
            generated_text (str): The raw text generated by a model.

        Returns:
            str: The post-processed text.
        """
        pass


class RemoveThinkData(ModelPostProcessor):
    """
    Post-processor that removes content enclosed within <think>...</think> tags.

    Useful for stripping out internal reasoning, hidden thoughts, or metadata
    that is not intended to be shown in the final output.
    """

    def __init__(self, start_tag: str = "<think>", end_tag: str = "</think>"):
        """
        Initialize the post-processor with optional custom tags.

        Args:
            start_tag (str): The opening tag indicating the start of the hidden thoughts.
            end_tag (str): The closing tag indicating the end of the hidden thoughts.
        """
        self.start_tag = start_tag
        self.end_tag = end_tag

    def apply(self, generated_text: str) -> str:
        """
        Remove the first instance of text between the specified tags from the generated text.

        Args:
            generated_text (str): The raw text potentially containing <think>...</think>.

        Returns:
            str: The cleaned text with the <think> block removed if present.
        """
        if self.start_tag in generated_text and self.end_tag in generated_text:
            index_start = generated_text.find(self.start_tag)
            index_end = generated_text.find(self.end_tag) + len(self.end_tag)
            return generated_text[:index_start] + generated_text[index_end:]
        else:
            return generated_text
