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)
    
    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
        
        # For each sequence in the batch
        for batch_idx in range(batch_size):
            # Get the current text generated so far
            current_ids = input_ids[batch_idx]
            current_text = self.tokenizer.decode(current_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)
            
            # For each possible next token
            for token_idx in range(vocab_size):
                # Create a copy of the validator to test this token
                validator_copy = self.schema_validator.copy()
                
                # Decode the token and test if appending it would be valid
                token_text = self.tokenizer.decode([token_idx])
                if validator_copy.append(token_text):
                    valid_tokens_mask[token_idx] = True
            
            # Mask out invalid tokens by setting their scores to negative infinity
            invalid_tokens_mask = ~valid_tokens_mask
            scores[batch_idx, invalid_tokens_mask] = float('-inf')
        
        return scores