"""
Generate Image Prompts AI Function
Extracts image prompts from content using AI
"""
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 Content, Images
from igny8_core.ai.ai_core import AICore
from igny8_core.ai.validators import validate_ids
from igny8_core.ai.prompts import PromptRegistry

logger = logging.getLogger(__name__)


class GenerateImagePromptsFunction(BaseAIFunction):
    """Generate image prompts from content using AI"""
    
    def get_name(self) -> str:
        return 'generate_image_prompts'
    
    def get_metadata(self) -> Dict:
        return {
            'display_name': 'Generate Image Prompts',
            'description': 'Extract image prompts from content (title, intro, H2 headings)',
            'phases': {
                'INIT': 'Initializing prompt generation...',
                'PREP': 'Loading content and extracting elements...',
                'AI_CALL': 'Generating prompts with AI...',
                'PARSE': 'Parsing prompt data...',
                'SAVE': 'Saving prompts...',
                'DONE': 'Prompts generated!'
            }
        }
    
    def get_max_items(self) -> int:
        return 50  # Max content records per batch
    
    def validate(self, payload: dict, account=None) -> Dict:
        """Validate content IDs exist"""
        result = validate_ids(payload, max_items=self.get_max_items())
        if not result['valid']:
            return result
        
        # Check content records exist
        content_ids = payload.get('ids', [])
        if content_ids:
            queryset = Content.objects.filter(id__in=content_ids)
            if account:
                queryset = queryset.filter(account=account)
            
            if queryset.count() == 0:
                return {'valid': False, 'error': 'No content records found'}
        
        return {'valid': True}
    
    def prepare(self, payload: dict, account=None) -> List:
        """Load content records and extract elements for prompt generation"""
        content_ids = payload.get('ids', [])
        
        queryset = Content.objects.filter(id__in=content_ids)
        if account:
            queryset = queryset.filter(account=account)
        
        contents = list(queryset.select_related('task', 'account', 'site', 'sector'))
        
        if not contents:
            raise ValueError("No content records found")
        
        # Get max_in_article_images from IntegrationSettings
        max_images = self._get_max_in_article_images(account)
        
        # Extract content elements for each content record
        extracted_data = []
        for content in contents:
            extracted = self._extract_content_elements(content, max_images)
            extracted_data.append({
                'content': content,
                'extracted': extracted,
                'max_images': max_images,
            })
        
        return extracted_data
    
    def build_prompt(self, data: Any, account=None) -> str:
        """Build prompt using PromptRegistry - handles list of content items"""
        # Handle list of content items (from prepare)
        if isinstance(data, list):
            if not data:
                raise ValueError("No content items provided")
            # For now, process first item (can be extended to batch process all)
            data = data[0]
        
        extracted = data['extracted']
        max_images = data.get('max_images', 2)
        
        # Format content for prompt
        content_text = self._format_content_for_prompt(extracted)
        
        # Get prompt from PromptRegistry - same as other functions
        prompt = PromptRegistry.get_prompt(
            function_name='generate_image_prompts',
            account=account,
            context={
                'title': extracted['title'],
                'content': content_text,
                'max_images': max_images,
            }
        )
        
        return prompt
    
    def parse_response(self, response: str, step_tracker=None) -> Dict:
        """Parse AI response - same pattern as other functions"""
        ai_core = AICore(account=getattr(self, 'account', None))
        json_data = ai_core.extract_json(response)
        
        if not json_data:
            raise ValueError(f"Failed to parse image prompts response: {response[:200]}...")
        
        # Validate structure
        if 'featured_prompt' not in json_data:
            raise ValueError("Missing 'featured_prompt' in AI response")
        
        if 'in_article_prompts' not in json_data:
            raise ValueError("Missing 'in_article_prompts' in AI response")
        
        return json_data
    
    def save_output(
        self,
        parsed: Dict,
        original_data: Any,
        account=None,
        progress_tracker=None,
        step_tracker=None
    ) -> Dict:
        """Save prompts to Images model - handles list of content items"""
        # Handle list of content items (from prepare)
        if isinstance(original_data, list):
            if not original_data:
                raise ValueError("No content items provided")
            # For now, process first item (can be extended to batch process all)
            original_data = original_data[0]
        
        content = original_data['content']
        extracted = original_data['extracted']
        max_images = original_data.get('max_images', 2)
        
        prompts_created = 0
        
        with transaction.atomic():
            # Save featured image prompt - use content instead of task
            Images.objects.update_or_create(
                content=content,
                image_type='featured',
                defaults={
                    'prompt': parsed['featured_prompt'],
                    'status': 'pending',
                    'position': 0,
                }
            )
            prompts_created += 1
            
            # Save in-article image prompts
            in_article_prompts = parsed.get('in_article_prompts', [])
            h2_headings = extracted.get('h2_headings', [])
            
            for idx, prompt_text in enumerate(in_article_prompts[:max_images]):
                heading = h2_headings[idx] if idx < len(h2_headings) else f"Section {idx + 1}"
                
                Images.objects.update_or_create(
                    content=content,
                    image_type='in_article',
                    position=idx + 1,
                    defaults={
                        'prompt': prompt_text,
                        'status': 'pending',
                    }
                )
                prompts_created += 1
        
        return {
            'count': prompts_created,
            'prompts_created': prompts_created,
        }
    
    # Helper methods
    def _get_max_in_article_images(self, account) -> int:
        """Get max_in_article_images from IntegrationSettings"""
        try:
            from igny8_core.modules.system.models import IntegrationSettings
            settings = IntegrationSettings.objects.get(
                account=account,
                integration_type='image_generation'
            )
            return settings.config.get('max_in_article_images', 2)
        except IntegrationSettings.DoesNotExist:
            return 2  # Default
    
    def _extract_content_elements(self, content: Content, max_images: int) -> Dict:
        """Extract title, intro paragraphs, and H2 headings from content HTML"""
        from bs4 import BeautifulSoup
        
        html_content = content.html_content or ''
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # Extract title
        title = content.title or content.task.title or ''
        
        # Extract first 1-2 intro paragraphs (skip italic hook if present)
        paragraphs = soup.find_all('p')
        intro_paragraphs = []
        for p in paragraphs[:3]:  # Check first 3 paragraphs
            text = p.get_text(strip=True)
            # Skip italic hook (usually 30-40 words)
            if len(text.split()) > 50:  # Real paragraph, not hook
                intro_paragraphs.append(text)
                if len(intro_paragraphs) >= 2:
                    break
        
        # Extract first N H2 headings
        h2_tags = soup.find_all('h2')
        h2_headings = [h2.get_text(strip=True) for h2 in h2_tags[:max_images]]
        
        return {
            'title': title,
            'intro_paragraphs': intro_paragraphs,
            'h2_headings': h2_headings,
        }
    
    def _format_content_for_prompt(self, extracted: Dict) -> str:
        """Format extracted content for prompt input"""
        lines = []
        
        if extracted.get('intro_paragraphs'):
            lines.append("ARTICLE INTRODUCTION:")
            for para in extracted['intro_paragraphs']:
                lines.append(para)
            lines.append("")
        
        if extracted.get('h2_headings'):
            lines.append("ARTICLE HEADINGS (for in-article images):")
            for idx, heading in enumerate(extracted['h2_headings'], 1):
                lines.append(f"{idx}. {heading}")
        
        return "\n".join(lines)

