# src/interpretability/tasks/web_api_task.py
"""
Simplified task handler for web API requests.
Doesn't rely on CSV data structure.
"""
from typing import Dict, Optional

class WebAPITask:
    """
    A generic task handler for web API requests.
    Constructs prompts dynamically from user inputs.
    """
    
    def __init__(self, config: Dict, prompt_templates: Dict):
        self.config = config
        self.prompt_templates = prompt_templates
    
    def get_task_content(self, item: Dict) -> Dict:
        """
        Build task content from a web API request item.
        
        Expected item structure:
        {
            'task_text': str,        # Main question/task
            'instruction': str,       # System/instruction prompt
            'has_image': bool,
            'image_path': str or None
        }
        """
        # Extract user inputs
        task_text = item.get('task_text', '')
        instruction = item.get('instruction', '')
        has_image = item.get('has_image', False)
        image_type = item.get('image_type', 'actual')  # 'actual', 'blank', or 'gaussian'
        
        # Build the chunks
        # For simplicity, we'll use a basic template structure
        # You can customize this based on your needs
        
        if task_text:
            text_chunk = f"{task_text}"
        else:
            text_chunk = ""
        
        if instruction:
            inst_chunk = f"{instruction}"
        else:
            inst_chunk = ""
        
        # Combine into full prompt
        full_prompt_text = text_chunk + inst_chunk
        
        # Build chunks for attention tracking
        chunks_for_attention = {}
        if text_chunk:
            chunks_for_attention['text'] = text_chunk
        if inst_chunk:
            chunks_for_attention['instruction'] = inst_chunk
        
        # Determine image configuration
        if image_type == 'blank':
            return {
                "full_prompt": full_prompt_text,
                "has_image": False,
                "has_blank_image": True,
                "has_gaussian_noise_image": False,
                "chunks_for_attention": chunks_for_attention
            }
        elif image_type == 'gaussian':
            return {
                "full_prompt": full_prompt_text,
                "has_image": False,
                "has_blank_image": False,
                "has_gaussian_noise_image": True,
                "chunks_for_attention": chunks_for_attention
            }
        else:  # 'actual' or has_image=True
            return {
                "full_prompt": full_prompt_text,
                "has_image": has_image,
                "has_blank_image": False,
                "has_gaussian_noise_image": False,
                "chunks_for_attention": chunks_for_attention
            }