import json
from outlines_core.json_schema import build_regex_from_schema
from outlines_core.fsm.guide import Guide, Index, Vocabulary

class GuideWrapper:
    """Wrapper for the Guide class to handle different guide types."""

    def __init__(self, schema: dict, tokenizer_model: str):
        self.regex = build_regex_from_schema(json.dumps(schema))
        self.vocabulary = Vocabulary.from_pretrained(tokenizer_model)
        self.index = Index(self.regex, self.vocabulary)
        self.guide = Guide(self.index)

    def get_current_state(self):
        """Get current state of the Guide."""
        return self.guide.get_state()

    def get_allowed_tokens(self):
        """Get allowed tokens for the current state of the Guide."""
        return self.guide.get_tokens()

    def advance(self, token_id: int):
        """Advance Guide to the next state via some token_id and return allowed tokens for that new state."""
        return self.guide.advance(token_id)

    def is_finished(self):
        """Check if Guide is finished."""
        return self.guide.is_finished()

    def get_eos_token_id(self):
        """Get the EOS token ID from the vocabulary."""
        return self.vocabulary.get_eos_token_id()

# Example usage
schema = {
    "title": "Foo",
    "type": "object",
    "properties": {"date": {"type": "string", "format": "date"}}
}
tokenizer_model = "openai-community/gpt2"
guide_wrapper = GuideWrapper(schema, tokenizer_model)

# Get current state of the Guide:
current_state = guide_wrapper.get_current_state()

# Get allowed tokens for the current state of the Guide:
allowed_tokens = guide_wrapper.get_allowed_tokens()

# Advance Guide to the next state via some token_id and return allowed tokens for that new state:
next_allowed_tokens = guide_wrapper.advance(allowed_tokens[-1])

# To check if Guide is finished:
guide_finished = guide_wrapper.is_finished()

# If it's finished then this assertion holds:
assert guide_wrapper.get_allowed_tokens() == [guide_wrapper.get_eos_token_id()]