# src/interpretability/api/wrapper.py
"""
Clean API wrapper for external applications.
Translates simple inputs to the complex internal pipeline.
"""
import os
import tempfile
import logging
from typing import Optional, List, Dict, Any
from pathlib import Path

from ..models.llava_onevision import LlavaOnevisionHandler
from ..models.qwen2_llm import Qwen2LLMHandler
from ..models.qwen_25_vl import QwenVLHandler
from ..tasks.web_api_task import WebAPITask  # NEW: Use web-specific task
from ..analysis.extractor_aggregated import AttentionExtractor


class InterpretabilityAnalyzer:
    """
    Simplified interface for running interpretability analysis.
    Designed to be called from web apps, APIs, etc.
    """
    
    def __init__(self, model_type: str, model_id: str, device: str = 'auto'):
        """
        Initialize the analyzer with minimal configuration.
        
        Args:
            model_type: "llava_onevision", "qwen2_llm", or "qwen_25_vl"
            model_id: HuggingFace model identifier
        """
        self.logger = logging.getLogger(__name__)
        
        # Build minimal config programmatically
        self.model_config = {
            'type': model_type,
            'model_id': model_id,
            'device': device,
            'torch_dtype': 'bfloat16',
            'attn_implementation': 'eager',
        }
        
        # Initialize model handler
        if model_type == "llava_onevision":
            self.model_handler = LlavaOnevisionHandler(self.model_config)
        elif model_type == "qwen2_llm":
            self.model_handler = Qwen2LLMHandler(self.model_config)
        elif model_type == "qwen_25_vl":
            self.model_handler = QwenVLHandler(self.model_config)
        else:
            raise ValueError(f"Unknown model type: {model_type}")
        
        self.logger.info(f"Loading model: {model_id}")
        self.model_handler.load_model()
        self.logger.info("Model loaded successfully")
        
        # Initialize web API task (no specific task type needed)
        self.task = WebAPITask(
            config={'name': 'web_api'},
            prompt_templates={}  # Not needed for web API
        )
    
    def analyze(
        self,
        image_path: Optional[str] = None,
        task_text: str = "",
        instruction: str = "",
        blocking_mode: str = "none",
        custom_steps: Optional[List[int]] = None,
        num_tokens: int = 20,
        image_type: str = "actual"  # "actual", "blank", or "gaussian"
    ) -> Dict[str, Any]:
        """
        Run interpretability analysis with simple inputs.
        
        Args:
            image_path: Path to image file (optional)
            task_text: The main task/question text
            instruction: Additional instruction text
            blocking_mode: "none", "previous", or "custom"
            custom_steps: List of steps to block (for custom mode)
            num_tokens: Number of tokens to generate
            image_type: Type of image - "actual", "blank", or "gaussian"
            
        Returns:
            Dictionary with:
            - input_tokens: List of input token strings
            - output_tokens: List of generated token strings
            - series: Dict of attention series (Image, Text, Instruction, Previous)
            - metadata: Additional info
        """
        # Build extractor config
        extractor_config = {
            'save_dir': tempfile.gettempdir(),
            'experiment_name': 'web_api',
            'generation': {'num_tokens_to_generate': num_tokens},
            'data': {
                'image_base_dir': '',
                'black_image_path': self._get_blank_image_path(),
                'gaussian_noise_image_path': self._get_gaussian_image_path()
            },
            'model': self.model_config,
            'task': {
                'name': 'web_api',
                'prompt_template': 'web_api',  # ADD THIS LINE
                'image_type': image_type        # ADD THIS LINE TOO
            }
        }
        
        # Prepare the item (like a CSV row, but for single request)
        item = {
            # User inputs
            'task_text': task_text,
            'instruction': instruction,
            'has_image': image_path is not None and image_type == "actual",
            'image_type': image_type,
            'image_path': image_path if image_path else '',
            
            # Dummy values for compatibility with get_filename_for_experiment
            'fruit_category': 'web',
            'sport_category': 'api',
            'math_answer': 'request',
            'query_item': task_text[:20] if task_text else 'query',
            'image_count': 1 if image_path else 0,
            'text_count': len(task_text.split()) if task_text else 0
        }
        
        # Build blocking config
        blocking_dict = self._build_blocking_config(blocking_mode, custom_steps)
        
        # Create extractor
        self.logger.info("Creating attention extractor...")
        extractor = AttentionExtractor(
            self.model_handler,
            self.task,
            extractor_config
        )
        
        # Process the sample
        self.logger.info("Processing sample...")
        extractor.process_sample(item, sample_idx=0, blocking_dict=blocking_dict)
        
        # Return the results
        result = extractor.output_dict
        result['metadata'] = {
            'blocking_mode': blocking_mode,
            'custom_steps': custom_steps,
            'image_used': image_path,
            'image_type': image_type
        }
        
        self.logger.info("Analysis complete!")
        return result
    
    def _build_blocking_config(
        self,
        blocking_mode: str,
        custom_steps: Optional[List[int]]
    ) -> Optional[Dict]:
        """Convert simple blocking mode to internal config."""
        if blocking_mode == "none":
            return None
        elif blocking_mode == "previous":
            return {
                'block_type': 'layer_collapsed',
                'step_to_block': list(range(2, 25)),
                'chunk': 'previously_generating_tokens'
            }
        elif blocking_mode == "custom" and custom_steps:
            return {
                'block_type': 'layer_collapsed',
                'step_to_block': custom_steps,
                'chunk': 'previously_generating_tokens'
            }
        return None
    
    def _get_blank_image_path(self) -> str:
        """Get or create a blank image for testing."""
        # You can generate this on-the-fly or point to a static file
        blank_path = os.path.join(tempfile.gettempdir(), 'blank_image.png')
        if not os.path.exists(blank_path):
            from PIL import Image
            img = Image.new('RGB', (224, 224), color='black')
            img.save(blank_path)
        return blank_path
    
    def _get_gaussian_image_path(self) -> str:
        """Get or create a gaussian noise image for testing."""
        import numpy as np
        noise_path = os.path.join(tempfile.gettempdir(), 'gaussian_noise.png')
        if not os.path.exists(noise_path):
            from PIL import Image
            noise = np.random.normal(128, 50, (224, 224, 3)).astype(np.uint8)
            img = Image.fromarray(noise)
            img.save(noise_path)
        return noise_path


# src/interpretability/api/wrapper.py

if __name__ == "__main__":
    # Example usage
    logging.basicConfig(level=logging.INFO)
    
    analyzer = InterpretabilityAnalyzer(
        model_type="llava_onevision",
        model_id="llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
    )
    
    # Option 1: Test without image
    result = analyzer.analyze(
        image_path=None,  # No image
        task_text="What is 2 + 2?",
        instruction="Answer with just the number.",
        blocking_mode="none",
        num_tokens=10,
        image_type="blank"  # Use blank image
    )
    
    print("\n=== Results ===")
    print(f"Input tokens: {result['input_tokens'][:20]}...")  # First 20
    print(f"Output tokens: {result['output_tokens']}")
    print(f"Series keys: {list(result['series'].keys())}")
    
    # Option 2: Test with actual image (uncomment and use real path)
    # result = analyzer.analyze(
    #     image_path='src/interpretability/data/images/Fruits_Clean_12/apple/Image_4.jpg',  # Use actual path
    #     task_text="What fruit is in the image?",
    #     instruction="Answer briefly.",
    #     blocking_mode="none",
    #     num_tokens=15
    # )

    # print("\n=== Results ===")
    # print(f"Input tokens: {result['input_tokens'][:20]}...")  # First 20
    # print(f"Output tokens: {result['output_tokens']}")
    # print(f"Series keys: {list(result['series'].keys())}")