"""
Generate Images AI Function
Extracted from modules/writer/tasks.py
"""
import logging
from typing import Dict, List, Any
from django.db import transaction
from igny8_core.ai.base import BaseAIFunction
from igny8_core.modules.writer.models import Tasks, Images
from igny8_core.ai.ai_core import AICore
from igny8_core.ai.validators import validate_tasks_exist
from igny8_core.ai.prompts import PromptRegistry
from igny8_core.ai.settings import get_model_config

logger = logging.getLogger(__name__)


class GenerateImagesFunction(BaseAIFunction):
    """Generate images for tasks using AI"""
    
    def get_name(self) -> str:
        return 'generate_images'
    
    def get_metadata(self) -> Dict:
        return {
            'display_name': 'Generate Images',
            'description': 'Generate featured and in-article images for tasks',
            'phases': {
                'INIT': 'Initializing image generation...',
                'PREP': 'Extracting image prompts...',
                'AI_CALL': 'Generating images with AI...',
                'PARSE': 'Processing image URLs...',
                'SAVE': 'Saving images...',
                'DONE': 'Images generated!'
            }
        }
    
    def get_max_items(self) -> int:
        # No limit - return None
        return None
    
    def validate(self, payload: dict, account=None) -> Dict:
        """Validate task IDs"""
        result = super().validate(payload, account)
        if not result['valid']:
            return result
        
        # Check tasks exist
        task_ids = payload.get('ids', [])
        if task_ids:
            task_result = validate_tasks_exist(task_ids, account)
            if not task_result['valid']:
                return task_result
        
        return {'valid': True}
    
    def prepare(self, payload: dict, account=None) -> Dict:
        """Load tasks and image generation settings"""
        task_ids = payload.get('ids', [])
        
        queryset = Tasks.objects.filter(id__in=task_ids)
        if account:
            queryset = queryset.filter(account=account)
        
        tasks = list(queryset.select_related('account', 'sector', 'site'))
        
        if not tasks:
            raise ValueError("No tasks found")
        
        # Get image generation settings
        image_settings = {}
        if account:
            try:
                from igny8_core.modules.system.models import IntegrationSettings
                integration = IntegrationSettings.objects.get(
                    account=account,
                    integration_type='image_generation',
                    is_active=True
                )
                image_settings = integration.config or {}
            except Exception:
                pass
        
        # Extract settings with defaults
        provider = image_settings.get('provider') or image_settings.get('service', 'openai')
        if provider == 'runware':
            model = image_settings.get('model') or image_settings.get('runwareModel', 'runware:97@1')
        else:
            model = image_settings.get('model', 'dall-e-3')
        
        return {
            'tasks': tasks,
            'account': account,
            'provider': provider,
            'model': model,
            'image_type': image_settings.get('image_type', 'realistic'),
            'max_in_article_images': int(image_settings.get('max_in_article_images', 2)),
            'desktop_enabled': image_settings.get('desktop_enabled', True),
            'mobile_enabled': image_settings.get('mobile_enabled', True),
        }
    
    def build_prompt(self, data: Dict, account=None) -> Dict:
        """Extract image prompts from task content"""
        task = data.get('task')
        max_images = data.get('max_in_article_images', 2)
        
        if not task or not task.content:
            raise ValueError("Task has no content")
        
        # Use AI to extract image prompts
        ai_core = AICore(account=account or data.get('account'))
        account_obj = account or data.get('account')
        
        # Get prompt from registry
        prompt = PromptRegistry.get_prompt(
            function_name='extract_image_prompts',
            account=account_obj,
            context={
                'title': task.title,
                'content': task.content[:5000],  # Limit content length
                'max_images': max_images
            }
        )
        
        # Get model config
        model_config = get_model_config('extract_image_prompts')
        
        # Call AI to extract prompts using centralized request handler
        result = ai_core.run_ai_request(
            prompt=prompt,
            model=model_config.get('model'),
            max_tokens=model_config.get('max_tokens'),
            temperature=model_config.get('temperature'),
            response_format=model_config.get('response_format'),
            function_name='extract_image_prompts'
        )
        
        if result.get('error'):
            raise ValueError(f"Failed to extract image prompts: {result['error']}")
        
        # Parse JSON response
        json_data = ai_core.extract_json(result['content'])
        
        if not json_data:
            raise ValueError("Failed to parse image prompts response")
        
        return {
            'featured_prompt': json_data.get('featured_prompt', ''),
            'in_article_prompts': json_data.get('in_article_prompts', [])
        }
    
    def parse_response(self, response: Dict, step_tracker=None) -> Dict:
        """Parse image generation response (already parsed, just return)"""
        return response
    
    def save_output(
        self,
        parsed: Dict,
        original_data: Dict,
        account=None,
        progress_tracker=None,
        step_tracker=None
    ) -> Dict:
        """Save images to database"""
        task = original_data.get('task')
        image_url = parsed.get('url')
        image_type = parsed.get('image_type')  # 'featured', 'desktop', 'mobile'
        
        if not task or not image_url:
            raise ValueError("Missing task or image URL")
        
        # Create Images record
        image = Images.objects.create(
            task=task,
            image_url=image_url,
            image_type=image_type,
            account=account or task.account,
            site=task.site,
            sector=task.sector,
        )
        
        return {
            'count': 1,
            'images_created': 1,
            'image_id': image.id
        }


def generate_images_core(task_ids: List[int], account_id: int = None, progress_callback=None):
    """
    Core logic for generating images (legacy function signature for backward compatibility).
    Can be called with or without Celery.
    
    Args:
        task_ids: List of task IDs
        account_id: Account ID for account isolation
        progress_callback: Optional function to call for progress updates
    
    Returns:
        Dict with 'success', 'images_created', 'message', etc.
    """
    try:
        from igny8_core.auth.models import Account
        
        account = None
        if account_id:
            account = Account.objects.get(id=account_id)
        
        # Use the new function class
        fn = GenerateImagesFunction()
        fn.account = account
        
        # Prepare payload
        payload = {'ids': task_ids}
        
        # Validate
        validated = fn.validate(payload, account)
        if not validated['valid']:
            return {'success': False, 'error': validated['error']}
        
        # Prepare data
        data = fn.prepare(payload, account)
        tasks = data['tasks']
        
        # Get prompts from registry
        image_prompt_template = PromptRegistry.get_image_prompt_template(account)
        negative_prompt = PromptRegistry.get_negative_prompt(account)
        
        ai_core = AICore(account=account)
        images_created = 0
        
        # Process each task
        for task in tasks:
            if not task.content:
                continue
            
            # Extract image prompts
            prompts_data = fn.build_prompt({'task': task, **data}, account)
            featured_prompt = prompts_data['featured_prompt']
            in_article_prompts = prompts_data['in_article_prompts']
            
            # Format featured prompt
            formatted_featured = image_prompt_template.format(
                image_type=data['image_type'],
                post_title=task.title,
                image_prompt=featured_prompt
            )
            
            # Generate featured image using centralized handler
            featured_result = ai_core.generate_image(
                prompt=formatted_featured,
                provider=data['provider'],
                model=data['model'],
                negative_prompt=negative_prompt,
                function_name='generate_images'
            )
            
            if not featured_result.get('error') and featured_result.get('url'):
                fn.save_output(
                    {'url': featured_result['url'], 'image_type': 'featured'},
                    {'task': task, **data},
                    account
                )
                images_created += 1
            
            # Generate in-article images (desktop/mobile if enabled)
            # ... (simplified for now, full logic in tasks.py)
        
        return {
            'success': True,
            'images_created': images_created,
            'message': f'Image generation complete: {images_created} images created'
        }
        
    except Exception as e:
        logger.error(f"Error in generate_images_core: {str(e)}", exc_info=True)
        return {'success': False, 'error': str(e)}

