import torch, json

class CharTokenizer:
    def __init__(self, corpus=None, vocab=None):
        if vocab is not None:
            self.vocab = vocab
        elif corpus is not None:
            self.vocab = self._build_vocab(corpus)
        else:
            raise Exception("Either corpus or vocab has to be supplied")
        self.id2vocab = [char for char, index in sorted(self.vocab.items(), key=lambda item: item[1])]
        
    def _tokenize(self, text):
        return list(text)
        
    def __call__(self, prompt, text=None, add_eos_token=False):
        token_ids = [self.vocab.get(token, 0) for token in self._tokenize(prompt)]
        if text is not None:
            text_token_ids = [self.vocab.get(token, 0) for token in self._tokenize(text)]
            token_ids = token_ids + [self.vocab["<bos>"]] + text_token_ids
        if add_eos_token:
            token_ids = token_ids + [self.vocab["<eos>"]]
        input_ids_tensor = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0)
        attention_masks = torch.ones_like(input_ids_tensor)
        return {"input_ids": input_ids_tensor, "attention_mask": attention_masks}
        
    def _build_vocab(self, corpus):
        vocab = {"<pad>": 0}
        for verse_lengths in range(3, 10):
            vocab[str(verse_lengths)] = len(vocab)
        for doc in corpus:
            chars = self._tokenize(doc)
            for char in chars:
                if char not in vocab:
                    vocab[char] = len(vocab)
        vocab["<bos>"] = len(vocab)
        vocab["<eos>"] = len(vocab)
        return vocab
    
    def decode(self, token_ids):
        chars = [self.id2vocab[token_id] for token_id in token_ids.flatten().tolist()]
        filtered_chars = [char for char in chars if char not in ["<eos>", "<bos>", "<pad>"]]
        return "".join(filtered_chars)
    
    def save(self, filepath):
        with open(filepath, "w") as f:
            json.dump(self.vocab, f)
    
    @classmethod
    def load(cls, filepath):
        with open(filepath) as f:
            vocab = json.load(f)
        return cls(vocab=vocab)