import torch
from transformers import LogitsProcessor
from xml_schema_validator import XmlSchemaValidator

class XmlLogitsProcessor(LogitsProcessor):
    """
    A LogitsProcessor that enforces valid XML according to a schema.
    
    This processor masks tokens that would lead to invalid XML
    by setting their logits to negative infinity.
    """
    
    def __init__(self, tokenizer, schema_text: str):
        """
        Initialize the processor with a schema and tokenizer.
        
        Args:
            tokenizer: The tokenizer to use for decoding tokens
            schema_text: The XSD schema text to validate against
        """
        self.tokenizer = tokenizer
        self.schema_validator = XmlSchemaValidator(schema_text)
        self.prompt_length = None
        self.is_first_call = True

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        """
        Process logits to mask invalid XML tokens.
        
        Args:
            input_ids: Current input token IDs
            scores: Current scores (logits) for next token prediction
            
        Returns:
            Processed scores with invalid tokens masked
        """
        batch_size, _ = input_ids.shape
        
        # If this is the first call, store the prompt length
        if self.is_first_call:
            self.prompt_length = input_ids.shape[1]
            self.is_first_call = False
        
        # For each sequence in the batch
        for batch_idx in range(batch_size):
            # Get only the generated portion of the text
            current_ids = input_ids[batch_idx]
            generated_ids = current_ids[self.prompt_length:]
            generated_text = self.tokenizer.decode(generated_ids)
            
            # Get all possible next tokens
            vocab_size = scores.shape[-1]
            
            # Create a mask to track which tokens are valid
            valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
            batch_validator = self.schema_validator.copy()

            print("evaluate tokens continuing:", generated_text)
            
            if generated_text:
                valid, msg = batch_validator.append(generated_text)
                if not valid:
                    print("current generated text invalid, only accept eos_token_id")
                    eos_token_id = self.tokenizer.eos_token_id
                    if eos_token_id is not None:
                        valid_tokens_mask[eos_token_id] = True
                        invalid_tokens_mask = ~valid_tokens_mask
                        scores[batch_idx, invalid_tokens_mask] = float('-inf')
                        continue
            
            # Rest of the method remains the same
            for token_idx in range(vocab_size):
                token_validator = batch_validator.copy()
                token_text = self.tokenizer.decode([token_idx])
                valid, msg = token_validator.append(token_text)
                #print("token:", token_text, "valid:", valid)
                if valid:
                    #print("token:", generated_text + token_text)
                    valid_tokens_mask[token_idx] = True
            
            invalid_tokens_mask = ~valid_tokens_mask
            scores[batch_idx, invalid_tokens_mask] = float('-inf')
        
        return scores