"""
Standalone instruction classifier module for prompt injection defense
Integrates the instruction classifier model to sanitize tool outputs
"""

import os
import re
import json
import tempfile
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModel
import importlib.util
from pathlib import Path
import logging
from typing import List, Tuple, Dict, Any
import numpy as np

try:
    from huggingface_hub import hf_hub_download
except ImportError:
    hf_hub_download = None

try:
    import spaces
except ImportError:
    # Create a no-op decorator if spaces is not available
    def spaces_gpu_decorator(func):
        return func
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()

# Import required components from utils.py
from utils import (
    TransformerInstructionClassifier, 
    InstructionDataset, 
    collate_fn,
    get_device
)

class InstructionClassifierSanitizer:
    """
    Uses a trained instruction classifier model to detect and remove prompt injections
    from tool outputs by identifying instruction tokens and removing them.
    """
    
    def __init__(
        self,
        model_path: str = None,
        model_repo_id: str = "ddas/instruction-classifier-model",  # CHANGE THIS!
        model_filename: str = "best_instruction_classifier.pth",
        model_name: str = "xlm-roberta-base",
        threshold: float = 0.01,
        token_threshold: float = 0.4,
        max_length: int = 512,
        overlap: int = 256,
        use_local_model: bool = False  # Set to False to use HF Hub
    ):
        """
        Initialize the instruction classifier sanitizer
        
        Args:
            model_path: Path to local model file (if use_local_model=True)
            model_repo_id: Hugging Face model repository ID (if use_local_model=False)
            model_filename: Filename of the model in the HF repository
            model_name: Base transformer model name
            threshold: Document-level threshold - proportion of tokens that must be INSTRUCTION to classify document as injection
            token_threshold: Token-level threshold - probability threshold for classifying individual tokens as INSTRUCTION (0.0-1.0, lower = more aggressive)
            max_length: Maximum sequence length for sliding windows
            overlap: Overlap between sliding windows
            use_local_model: Whether to use local model file or download from HF Hub
        """
        self.model_name = model_name
        self.threshold = threshold
        self.token_threshold = token_threshold
        self.max_length = max_length
        self.overlap = overlap
        self.use_local_model = use_local_model
        self.model_repo_id = model_repo_id
        self.model_filename = model_filename
        
        # Initialize device - always use CPU for initialization in ZeroGPU environments
        # GPU operations will be handled within @spaces.GPU decorated methods
        self.device = torch.device('cpu')
        self.target_device = get_device()  # Store target device for later use
        print(f"🔧 Device configuration: init_device={self.device}, target_device={self.target_device}")
        
        # Map friendly names to actual model names
        model_mapping = {
            'modern-bert-base': 'answerdotai/ModernBERT-base',
            'xlm-roberta-base': 'xlm-roberta-base'
        }
        actual_model_name = model_mapping.get(model_name, model_name)
        
        # Load tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(actual_model_name)
        
        # Load model
        self.model = TransformerInstructionClassifier(
            model_name=actual_model_name,
            num_labels=2,
            dropout=0.1
        )
        
        # Load trained weights
        if self.use_local_model:
            # Use local model file
            if model_path is None:
                model_path = "models/best_instruction_classifier.pth"
            
            if os.path.exists(model_path):
                checkpoint = torch.load(model_path, map_location='cpu')
                self._load_model_weights(checkpoint)
                print(f"✅ Loaded instruction classifier model from {model_path}")
                print(f"   Model loaded on {self.device} for ZeroGPU compatibility")
            else:
                raise FileNotFoundError(f"Model file not found: {model_path}")
        else:
            # Download from Hugging Face Hub
            try:
                if hf_hub_download is None:
                    raise ImportError("huggingface_hub is not installed")
                
                print(f"🚀 Starting model download from {self.model_repo_id}")
                print(f"   Device: {self.device}")
                print(f"   Model name: {self.model_name}")
                
                # Use HF_TOKEN from environment for private repositories
                token = os.getenv('HF_TOKEN')
                if token:
                    print(f"📥 Downloading private model from {self.model_repo_id}...")
                    print(f"   Using HF_TOKEN: {token[:8]}...{token[-8:] if len(token) > 16 else 'short'}")
                else:
                    print(f"📥 Downloading public model from {self.model_repo_id}...")
                    print("   No HF_TOKEN found - using public access")
                
                # Download the model file (returns file path, not model object)
                print(f"   Downloading {self.model_filename}...")
                model_path = hf_hub_download(
                    repo_id=self.model_repo_id,
                    filename=self.model_filename,
                    cache_dir="./model_cache",
                    token=token  # Will be None for public repos
                )
                print(f"✅ Model file downloaded to: {model_path}")
                
                # Check file size
                file_size = os.path.getsize(model_path) / (1024**3)  # GB
                print(f"   File size: {file_size:.2f} GB")
                
                # Load the checkpoint from the downloaded file - always use CPU for ZeroGPU compatibility
                print("🔄 Loading checkpoint into memory...")
                checkpoint = torch.load(model_path, map_location='cpu')
                print(f"   Checkpoint keys: {len(checkpoint.keys())}")
                
                self._load_model_weights(checkpoint)
                print(f"✅ Model weights loaded from {self.model_repo_id}")
                print(f"   Model parameter count: {sum(p.numel() for p in self.model.parameters())}")
                print(f"   Model loaded on {self.device} for ZeroGPU compatibility")
                
            except Exception as e:
                print(f"❌ CRITICAL ERROR: Failed to download model from {self.model_repo_id}")
                print(f"   Error type: {type(e).__name__}")
                print(f"   Error message: {e}")
                print("   Full error details:")
                import traceback
                traceback.print_exc()
                print("   Environment info:")
                print(f"     HF_TOKEN set: {'Yes' if os.getenv('HF_TOKEN') else 'No'}")
                print(f"     Device: {self.device}")
                print(f"     PyTorch version: {torch.__version__}")
                raise RuntimeError(f"Failed to download model from {self.model_repo_id}: {e}")
    
    def _load_model_weights(self, checkpoint):
        """Helper method to load model weights with filtering"""
        # Filter out keys that don't belong to the model (like loss function weights)
        model_state_dict = {}
        for key, value in checkpoint.items():
            if not key.startswith('loss_fct'):  # Skip loss function weights
                model_state_dict[key] = value
        
        # Load the filtered state dict - keep on CPU for ZeroGPU compatibility
        self.model.load_state_dict(model_state_dict, strict=False)
        self.model.to(self.device)  # Keep on CPU during initialization
        self.model.eval()
    
    
    @spaces.GPU
    def sanitize_with_annotations(self, tool_output: str) -> Tuple[str, List[Dict[str, any]], str]:
        """
        Sanitization function that also returns annotation data for flagged content.
        
        Args:
            tool_output: The raw tool output string
            
        Returns:
            Tuple of (sanitized_output, annotations, merged_tagged_text) where:
            - sanitized_output: cleaned text with instruction content removed
            - annotations: position information for content flagged by classifier  
            - merged_tagged_text: text with <instruction> tags showing detected content
        """
        if not tool_output or not tool_output.strip():
            return tool_output, [], tool_output
        
        # Move model to target device (GPU) within @spaces.GPU decorated method
        if self.device != self.target_device:
            print(f"🚀 Moving model from {self.device} to {self.target_device} within @spaces.GPU context")
            self.model.to(self.target_device)
            self.device = self.target_device
        
        try:
            # Step 1: Detect if the tool output contains instructions
            is_injection, confidence_score, tagged_text = self._detect_injection(tool_output)
            
            print(f"🔍 Instruction detection: injection={is_injection}, confidence={confidence_score:.3f}")
            
            if not is_injection:
                print("✅ No injection detected - returning original output")
                return tool_output, [], tool_output
            
            print(f"🚨 Injection detected! Processing with extensions and annotations...")
            
            # Step 2: Extract annotation positions from original tagged text
            annotations = self._extract_annotations_from_tagged_text(tagged_text, tool_output)
            print(f"📝 Original tagged text: {tagged_text}")
            
            # Step 3: Extend instruction tags by one token on each side
            extended_tagged_text = self._extend_instruction_tags(tagged_text)
            print(f"🔄 Extended tagged text: {extended_tagged_text}")
            
            # Step 4: Merge close instruction tags  
            merged_tagged_text = self._merge_close_instruction_tags(extended_tagged_text, min_words_between=4)
            print(f"🔗 Merged tagged text: {merged_tagged_text}")
            
            # Step 5: Remove instruction tags and their content
            sanitized_output = self._remove_instruction_tags(merged_tagged_text)
            print(f"🔒 Sanitized output: {sanitized_output}")
            print(f"🔍 DEBUG SANITIZER: Returning merged_tagged_text: '{merged_tagged_text}'")
            return sanitized_output, annotations, merged_tagged_text
            
        except Exception as e:
            print(f"❌ Error in instruction classifier sanitization: {e}")
            # Return original output if sanitization fails
            return tool_output, [], tool_output
    
    def _extract_annotations_from_tagged_text(self, tagged_text: str, original_text: str) -> List[Dict[str, any]]:
        """
        Extract annotation positions from tagged text.
        
        Args:
            tagged_text: Text with <instruction> tags
            original_text: Original untagged text
            
        Returns:
            List of annotation dictionaries with content and address fields
        """
        import re
        
        annotations = []
        
        # Find all instruction tags in the tagged text
        pattern = r'<instruction>(.*?)</instruction>'
        matches = re.finditer(pattern, tagged_text, re.DOTALL)
        
        for match in matches:
            flagged_content = match.group(1).strip()
            
            # Find the position of this content in the original text
            start_pos = original_text.find(flagged_content)
            if start_pos != -1:
                end_pos = start_pos + len(flagged_content)
                
                annotation = {
                    "content": f"Instruction injection detected: '{flagged_content[:50]}{'...' if len(flagged_content) > 50 else ''}'",
                    "address": f"content:{start_pos}-{end_pos}",
                    "extra_metadata": {
                        "source": "instruction-classifier",
                        "flagged_text": flagged_content,
                        "detection_type": "instruction_injection"
                    }
                }
                annotations.append(annotation)
        
        return annotations
    
    def _detect_injection(self, tool_output: str) -> Tuple[bool, float, str]:
        """
        Detect if the tool output contains instructions that could indicate prompt injection.
        
        Returns:
            tuple: (is_injection, confidence_score, tagged_text) where:
                - is_injection: boolean indicating if injection was detected
                - confidence_score: proportion of tokens classified as instructions  
                - tagged_text: original text with <instruction> tags for debugging
        """
        if not tool_output.strip():
            return False, 0.0, ""
            
        try:
            # Use InstructionDataset sliding window logic for raw text inference
            predictions, original_tokens = self._predict_with_sliding_windows(tool_output)
            
            if not predictions:
                return False, 0.0, ""
            
            # Calculate the proportion of tokens classified as instructions (label 1)
            instruction_tokens = sum(1 for pred in predictions if pred == 1)
            total_tokens = len(predictions)
            confidence_score = instruction_tokens / total_tokens if total_tokens > 0 else 0.0
            
            # Determine if this is considered an injection based on threshold
            is_injection = confidence_score > self.threshold
            
            # Only reconstruct with tags if injection detected
            if is_injection:
                tagged_text = self._reconstruct_text_with_tags(original_tokens, predictions)
            else:
                tagged_text = tool_output
            
            return is_injection, confidence_score, tagged_text   
            
        except Exception as e:
            print(f"Error in instruction classifier detection: {e}")
            return False, 0.0, ""

    def _predict_with_sliding_windows(self, text: str) -> Tuple[List[int], List[str]]:
        """
        Simplified prediction using the predict_instructions function from utils.py
        This is more direct and avoids complex aggregation logic.
        """
        from utils import predict_instructions
        
        try:
            # Use the predict_instructions function directly with token-level threshold
            tokens, predictions = predict_instructions(self.model, self.tokenizer, text, self.device, self.token_threshold)
            return predictions, tokens
        except Exception as e:
            print(f"⚠️ FALLBACK TRIGGERED: Error in predict_instructions: {e}")
            print(f"   Using _simple_predict as fallback (still uses threshold={self.token_threshold})")
            # Fallback to simple tokenization if the complex method fails
            return self._simple_predict(text)
    
    def _simple_predict(self, text: str) -> Tuple[List[int], List[str]]:
        """
        Simple fallback prediction method without sliding windows
        """
        words = text.split()
        if not words:
            return [], []
        
        # Tokenize with word alignment
        encoded = self.tokenizer(
            words,
            is_split_into_words=True,
            add_special_tokens=True,
            truncation=True,
            padding=True,
            max_length=self.max_length,
            return_tensors='pt'
        )
        
        # Move to device
        input_ids = encoded['input_ids'].to(self.device)
        attention_mask = encoded['attention_mask'].to(self.device)
        
        # Predict
        self.model.eval()
        with torch.no_grad():
            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
            # Use threshold approach (same as main prediction) instead of argmax
            probs = torch.softmax(outputs['logits'], dim=-1)
            predictions = (probs[:, :, 1] > self.token_threshold).long()
        
        # Convert back to word-level predictions
        word_ids = encoded.word_ids()
        word_predictions = []
        prev_word_id = None
        
        for i, word_id in enumerate(word_ids):
            if word_id is not None and word_id != prev_word_id:
                if word_id < len(words):
                    pred_idx = min(i, predictions.shape[1] - 1)
                    word_predictions.append(predictions[0, pred_idx].item())
                prev_word_id = word_id
        
        # Ensure same length
        while len(word_predictions) < len(words):
            word_predictions.append(0)
        
        return word_predictions[:len(words)], words
    
    def _convert_subword_to_word_predictions(self, subword_tokens, subword_predictions, original_text):
        """Convert aggregated subword predictions back to word-level predictions"""
        # Simple approach: re-tokenize original text and align
        original_words = original_text.split()
        
        # Use tokenizer to get word alignment
        encoded = self.tokenizer(
            original_words,
            is_split_into_words=True,
            add_special_tokens=True,
            truncation=False,
            padding=False,
            return_tensors='pt'
        )
        
        word_ids = encoded.word_ids()
        word_predictions = []
        
        # Extract word-level predictions using BERT approach
        prev_word_id = None
        subword_idx = 0
        
        for i, word_id in enumerate(word_ids):
            if word_id is not None and word_id != prev_word_id:
                # First subtoken of new word - use its prediction
                if subword_idx < len(subword_predictions) and word_id < len(original_words):
                    word_predictions.append(subword_predictions[subword_idx])
                prev_word_id = word_id
            if subword_idx < len(subword_predictions):
                subword_idx += 1
        
        # Ensure same length
        while len(word_predictions) < len(original_words):
            word_predictions.append(0)
        
        return word_predictions[:len(original_words)], original_words

    def _reconstruct_text_with_tags(self, tokens, predictions):
        """Reconstruct text from tokens and predictions, adding instruction tags"""
        if len(tokens) != len(predictions):
            print(f"Length mismatch: tokens ({len(tokens)}) vs predictions ({len(predictions)})")
            # Truncate to the shorter length to avoid crashes
            min_length = min(len(tokens), len(predictions))
            tokens = tokens[:min_length]
            predictions = predictions[:min_length]
        
        result_parts = []
        current_instruction = []
        
        for token, pred in zip(tokens, predictions):
            if pred == 1:  # INSTRUCTION
                current_instruction.append(token)
            else:  # OTHER
                # If we were building an instruction, close it
                if current_instruction:
                    instruction_text = ' '.join(current_instruction)
                    result_parts.append(f'<instruction>{instruction_text}</instruction>')
                    current_instruction = []
                
                # Add the non-instruction token
                result_parts.append(token)
        
        # Handle case where text ends with an instruction
        if current_instruction:
            instruction_text = ' '.join(current_instruction)
            result_parts.append(f'<instruction>{instruction_text}</instruction>')
        
        # Join with spaces
        result = ' '.join(result_parts)
        
        return result

    def _extend_instruction_tags(self, tagged_text: str) -> str:
        """
        Extend each <instruction>...</instruction> block by one word token on each side,
        but only if the adjacent token is not already instruction-tagged.
        
        This prevents overlaps between instruction blocks while extending them
        to capture more context around detected instruction content.
        
        Args:
            tagged_text: Text with <instruction>...</instruction> tags
            
        Returns:
            Text with extended instruction tags
        """
        if not tagged_text.strip():
            return tagged_text
        
        # Find all instruction regions first to avoid overlaps
        instruction_regions = []
        pattern = r'<instruction>(.*?)</instruction>'
        
        for match in re.finditer(pattern, tagged_text, re.DOTALL):
            instruction_regions.append({
                'start': match.start(),
                'end': match.end(),
                'content': match.group(1)
            })
        
        if not instruction_regions:
            return tagged_text
        
        # Split into words while preserving positions
        words = tagged_text.split()
        
        # Build word-to-character position mapping
        word_positions = []
        char_pos = 0
        for word in words:
            start_pos = tagged_text.find(word, char_pos)
            end_pos = start_pos + len(word)
            word_positions.append({
                'word': word,
                'start': start_pos,
                'end': end_pos
            })
            char_pos = end_pos
        
        # Find which words are currently inside instruction tags
        instruction_word_indices = set()
        
        for region in instruction_regions:
            for i, word_info in enumerate(word_positions):
                # Check if word overlaps with instruction region
                if (word_info['start'] < region['end'] and word_info['end'] > region['start']):
                    instruction_word_indices.add(i)
        
        # Find instruction blocks by consecutive instruction words
        instruction_blocks = []
        current_block = None
        
        for i in range(len(words)):
            if i in instruction_word_indices:
                if current_block is None:
                    current_block = {'start': i, 'end': i}
                else:
                    current_block['end'] = i
            else:
                if current_block is not None:
                    instruction_blocks.append(current_block)
                    current_block = None
        
        # Don't forget the last block
        if current_block is not None:
            instruction_blocks.append(current_block)
        
        # Plan extensions with proper overlap prevention
        extensions = []
        planned_tagged_words = set(instruction_word_indices)  # Start with currently tagged words
        
        for block in instruction_blocks:
            start_idx = block['start']
            end_idx = block['end']
            
            extend_left = False
            extend_right = False
            
            # Try extend left (if not at beginning and previous token not planned to be tagged)
            if start_idx > 0 and (start_idx - 1) not in planned_tagged_words:
                extend_left = True
                planned_tagged_words.add(start_idx - 1)  # Reserve this word
            
            # Try extend right (if not at end and next token not planned to be tagged)
            if end_idx < len(words) - 1 and (end_idx + 1) not in planned_tagged_words:
                extend_right = True
                planned_tagged_words.add(end_idx + 1)  # Reserve this word
            
            extensions.append({
                'original_start': start_idx,
                'original_end': end_idx,
                'new_start': start_idx - (1 if extend_left else 0),
                'new_end': end_idx + (1 if extend_right else 0),
                'extend_left': extend_left,
                'extend_right': extend_right
            })
        
        # Create a mapping of which extension block each word belongs to
        word_to_block = {}
        for block_idx, ext in enumerate(extensions):
            for i in range(ext['new_start'], ext['new_end'] + 1):
                word_to_block[i] = block_idx
        
        # Reconstruct the text with separate instruction blocks
        result_parts = []
        current_block = None
        
        for i, word in enumerate(words):
            # Clean the word of existing tags
            clean_word = word.replace('<instruction>', '').replace('</instruction>', '')
            
            # Skip empty words (from empty instruction tags)
            if not clean_word.strip():
                continue
                
            word_block = word_to_block.get(i, None)
            
            if word_block is not None and current_block != word_block:
                # Close previous block if needed
                if current_block is not None:
                    result_parts[-1] += '</instruction>'
                
                # Start new instruction block  
                result_parts.append(f'<instruction>{clean_word}')
                current_block = word_block
                
            elif word_block is not None and current_block == word_block:
                # Continue current instruction block
                result_parts.append(clean_word)
                
            elif word_block is None and current_block is not None:
                # End instruction block and add normal word
                result_parts[-1] += '</instruction>'
                result_parts.append(clean_word)
                current_block = None
                
            else:
                # Normal word (not in any instruction block)
                result_parts.append(clean_word)
        
        # Close instruction if we ended inside one
        if current_block is not None:
            result_parts[-1] += '</instruction>'
        
        return ' '.join(result_parts)

    def _merge_close_instruction_tags(self, text, min_words_between=3):
        """
        Merge <instruction>...</instruction> segments that are separated by less than min_words_between words
        """
        pattern = re.compile(r"(</instruction>)(\s+)([^<]+?)(\s+)(<instruction>)", re.DOTALL)

        def should_merge(between_text):
            # Count words in between_text
            words = re.findall(r"\b\w+\b", between_text)
            return len(words) < min_words_between

        # Keep merging until no more merges are possible
        changed = True
        while changed:
            changed = False
            # Find all potential merge points in the current text
            matches = list(pattern.finditer(text))
            
            # Process matches from right to left to avoid position shifts
            for match in reversed(matches):
                between_text = match.group(3)
                if should_merge(between_text):
                    # Merge: remove the tags between, include the in-between text inside the instruction tags
                    text = (
                        text[: match.start(1)]  # Text before </instruction>
                        + match.group(2)        # Whitespace after </instruction>
                        + between_text          # Text between tags
                        + match.group(4)        # Whitespace before <instruction>
                        + text[match.end(5):]   # Text after <instruction>
                    )
                    changed = True
                    break  # Start over since we changed the text
        
        return text

    def _remove_instruction_tags(self, text: str) -> str:
        """Remove all <instruction>...</instruction> tags and their content from text"""
        # Pattern to match <instruction>...</instruction> tags (including nested content)
        # Using non-greedy matching to handle multiple instruction blocks
        pattern = r'<instruction>.*?</instruction>'
        
        # Remove all instruction tags and their content
        cleaned_text = re.sub(pattern, '', text, flags=re.DOTALL | re.IGNORECASE)
        
        # Clean up any extra whitespace that might be left
        cleaned_text = re.sub(r'\s+', ' ', cleaned_text).strip()
        
        return cleaned_text


# Global instance of the sanitizer
_sanitizer_instance = None

def get_sanitizer():
    """Get or create the global sanitizer instance"""
    global _sanitizer_instance
    if _sanitizer_instance is None:
        try:
            # For Hugging Face Spaces deployment, use external model hosting
            # The model_repo_id is already set to "ddas/instruction-classifier-model"
            print("🚀 Initializing instruction classifier from Hugging Face Hub...")
            _sanitizer_instance = InstructionClassifierSanitizer(
                use_local_model=False,
                model_repo_id="ddas/instruction-classifier-model"
            )
            print("✅ Instruction classifier initialized successfully!")
        except Exception as e:
            print(f"❌ Failed to initialize instruction classifier from HF Hub: {e}")
            print("🔄 Falling back to local model if available...")
            try:
                _sanitizer_instance = InstructionClassifierSanitizer(use_local_model=True)
                print("✅ Local model initialized as fallback!")
            except Exception as e2:
                print(f"❌ Local model also failed: {e2}")
                print("⚠️  Instruction classifier disabled - sanitization will be bypassed")
                return None
    return _sanitizer_instance


def sanitize_tool_output_with_annotations(tool_output, defense_enabled=True):
    """
    Enhanced sanitization function that also returns annotation data for flagged content.
    
    Args:
        tool_output: The raw tool output string
        defense_enabled: Whether defense is enabled (passed from agent)
        
    Returns:
        Tuple of (sanitized_output, annotations, merged_tagged_text) where:
        - sanitized_output: cleaned text with instruction content removed
        - annotations: position information for content flagged by classifier
        - merged_tagged_text: text with <instruction> tags showing detected content
    """
    print(f"🔍 sanitize_tool_output_with_annotations called with: {tool_output[:100]}...")
    
    # If defense is disabled globally, return original output with no annotations
    if not defense_enabled:
        print("⚠️ Defense disabled - returning original output without processing")
        return tool_output, [], tool_output
    
    sanitizer = get_sanitizer()
    if sanitizer is None:
        print("⚠️  Instruction classifier not available, returning original output")
        return tool_output, [], tool_output
    
    print("✅ Sanitizer found, processing with annotations...")
    sanitized_output, annotations, merged_tagged_text = sanitizer.sanitize_with_annotations(tool_output)
    print(f"🔒 Sanitization complete, result: {sanitized_output[:100]}...")
    print(f"📝 Found {len(annotations)} annotations")
    return sanitized_output, annotations, merged_tagged_text