"""
RiadDiseaseGPT Tokenizer - SentencePiece based
"""

import os
import json
from typing import List, Dict, Optional
from transformers import PreTrainedTokenizer

try:
    import sentencepiece as spm
except ImportError:
    raise ImportError("Please install sentencepiece: pip install sentencepiece")


class RiadDiseaseGPTTokenizer(PreTrainedTokenizer):
    """SentencePiece tokenizer for RiadDiseaseGPT"""
    
    model_input_names = ["input_ids", "attention_mask"]
    
    def __init__(
        self,
        vocab_file: Optional[str] = None,
        model_file: Optional[str] = None,
        unk_token="<unk>",
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
        **kwargs
    ):
        self.sp_model = None
        self.vocab_file = vocab_file
        
        # Try to find tokenizer.model in same directory as vocab_file
        if vocab_file and not model_file:
            model_file = vocab_file.replace('tokenizer.json', 'tokenizer.model')
            if not os.path.exists(model_file):
                # Try current directory
                model_file = os.path.join(os.path.dirname(vocab_file or '.'), 'tokenizer.model')
        
        if model_file and os.path.exists(model_file):
            self.sp_model = spm.SentencePieceProcessor()
            self.sp_model.Load(model_file)
        
        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            **kwargs
        )
    
    @property
    def vocab_size(self) -> int:
        if self.sp_model:
            return self.sp_model.GetPieceSize()
        return 2500
    
    def get_vocab(self) -> Dict[str, int]:
        if not self.sp_model:
            return {}
        return {self.sp_model.IdToPiece(i): i for i in range(self.vocab_size)}
    
    def _tokenize(self, text: str) -> List[str]:
        if not self.sp_model:
            return text.split()
        return self.sp_model.EncodeAsPieces(text)
    
    def _convert_token_to_id(self, token: str) -> int:
        if not self.sp_model:
            return 0
        return self.sp_model.PieceToId(token)
    
    def _convert_id_to_token(self, index: int) -> str:
        if not self.sp_model:
            return "<unk>"
        return self.sp_model.IdToPiece(index)
    
    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        if not self.sp_model:
            return " ".join(tokens)
        return self.sp_model.DecodePieces(tokens)
    
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple:
        """Save tokenizer vocabulary"""
        if not os.path.isdir(save_directory):
            os.makedirs(save_directory)
        
        # Note: tokenizer.model should be copied manually
        out_vocab_file = os.path.join(
            save_directory, 
            (filename_prefix + "-" if filename_prefix else "") + "tokenizer.model"
        )
        
        return (out_vocab_file,)