"""
Generate Content AI Function
Extracted from modules/writer/tasks.py
"""
import logging
import re
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, Content as TaskContent
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 GenerateContentFunction(BaseAIFunction):
    """Generate content for tasks using AI"""
    
    def get_name(self) -> str:
        return 'generate_content'
    
    def get_metadata(self) -> Dict:
        return {
            'display_name': 'Generate Content',
            'description': 'Generate article content from task ideas',
            'phases': {
                'INIT': 'Initializing content generation...',
                'PREP': 'Loading tasks and building prompts...',
                'AI_CALL': 'Generating content with AI...',
                'PARSE': 'Processing content...',
                'SAVE': 'Saving content...',
                'DONE': 'Content generated!'
            }
        }
    
    def get_max_items(self) -> int:
        return 50  # Max tasks per batch
    
    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) -> List:
        """Load tasks with all relationships"""
        task_ids = payload.get('ids', [])
        
        queryset = Tasks.objects.filter(id__in=task_ids)
        if account:
            queryset = queryset.filter(account=account)
        
        # Preload all relationships to avoid N+1 queries
        tasks = list(queryset.select_related(
            'account', 'site', 'sector', 'cluster', 'idea'
        ))
        
        if not tasks:
            raise ValueError("No tasks found")
        
        return tasks
    
    def build_prompt(self, data: Any, account=None) -> str:
        """Build content generation prompt for a single task using registry"""
        if isinstance(data, list):
            # For now, handle single task (will be called per task)
            if not data:
                raise ValueError("No tasks provided")
            task = data[0]
        else:
            task = data
        
        account = account or task.account
        
        # Build idea data string
        idea_data = f"Title: {task.title or 'Untitled'}\n"
        if task.description:
            idea_data += f"Description: {task.description}\n"
        
        # Handle idea description (might be JSON or plain text)
        if task.idea and task.idea.description:
            description = task.idea.description
            try:
                import json
                parsed_desc = json.loads(description)
                if isinstance(parsed_desc, dict):
                    formatted_desc = "Content Outline:\n\n"
                    if 'H2' in parsed_desc:
                        for h2_section in parsed_desc['H2']:
                            formatted_desc += f"## {h2_section.get('heading', '')}\n"
                            if 'subsections' in h2_section:
                                for h3_section in h2_section['subsections']:
                                    formatted_desc += f"### {h3_section.get('subheading', '')}\n"
                                    formatted_desc += f"Content Type: {h3_section.get('content_type', '')}\n"
                                    formatted_desc += f"Details: {h3_section.get('details', '')}\n\n"
                    description = formatted_desc
            except (json.JSONDecodeError, TypeError):
                pass  # Use as plain text
            
            idea_data += f"Outline: {description}\n"
        
        if task.idea:
            idea_data += f"Structure: {task.idea.content_structure or task.content_structure or 'blog_post'}\n"
            idea_data += f"Type: {task.idea.content_type or task.content_type or 'blog_post'}\n"
            if task.idea.estimated_word_count:
                idea_data += f"Estimated Word Count: {task.idea.estimated_word_count}\n"
        
        # Build cluster data string
        cluster_data = ''
        if task.cluster:
            cluster_data = f"Cluster Name: {task.cluster.name or ''}\n"
            if task.cluster.description:
                cluster_data += f"Description: {task.cluster.description}\n"
            cluster_data += f"Status: {task.cluster.status or 'active'}\n"
        
        # Build keywords string
        keywords_data = task.keywords or ''
        if not keywords_data and task.idea:
            keywords_data = task.idea.target_keywords or ''
        
        # Get prompt from registry with context
        prompt = PromptRegistry.get_prompt(
            function_name='generate_content',
            account=account,
            task=task,
            context={
                'IDEA': idea_data,
                'CLUSTER': cluster_data,
                'KEYWORDS': keywords_data,
            }
        )
        
        return prompt
    
    def parse_response(self, response: str, step_tracker=None) -> Dict:
        """Parse content response - can be JSON or plain text"""
        import json
        
        # Try to parse as JSON first
        try:
            parsed_json = json.loads(response.strip())
            if isinstance(parsed_json, dict):
                # It's a JSON object with structured data
                return parsed_json
        except (json.JSONDecodeError, ValueError):
            pass
        
        # If not JSON, treat as plain content and normalize
        try:
            from igny8_core.utils.content_normalizer import normalize_content
            normalized = normalize_content(response)
            content_text = normalized['normalized_content']
            # Return as dict with content field for consistency
            return {'content': content_text}
        except Exception as e:
            logger.warning(f"Content normalization failed: {e}, using original")
            return {'content': response}
    
    def save_output(
        self,
        parsed: Any,
        original_data: Any,
        account=None,
        progress_tracker=None,
        step_tracker=None
    ) -> Dict:
        """Save content to task - handles both JSON and plain text responses"""
        if isinstance(original_data, list):
            task = original_data[0] if original_data else None
        else:
            task = original_data
        
        if not task:
            raise ValueError("No task provided for saving")
        
        # Handle parsed response - can be dict (JSON) or string (plain text)
        if isinstance(parsed, dict):
            # JSON response with structured fields
            content_html = parsed.get('content', '')
            title = parsed.get('title') or task.title
            meta_title = parsed.get('meta_title') or title or task.title
            meta_description = parsed.get('meta_description', '')
            word_count = parsed.get('word_count', 0)
            primary_keyword = parsed.get('primary_keyword', '')
            secondary_keywords = parsed.get('secondary_keywords', [])
            tags = parsed.get('tags', [])
            categories = parsed.get('categories', [])
            # Content status should always be 'draft' for newly generated content
            # Status can only be changed manually to 'review' or 'publish'
            content_status = 'draft'
        else:
            # Plain text response (legacy)
            content_html = str(parsed)
            title = task.title
            meta_title = task.meta_title or task.title
            meta_description = task.meta_description or (task.description or '')[:160] if task.description else ''
            word_count = 0
            primary_keyword = ''
            secondary_keywords = []
            tags = []
            categories = []
            content_status = 'draft'
        
        # Calculate word count if not provided
        if not word_count and content_html:
            text_for_counting = re.sub(r'<[^>]+>', '', content_html)
            word_count = len(text_for_counting.split())

        # Ensure related content record exists
        content_record, _created = TaskContent.objects.get_or_create(
            task=task,
            defaults={
                'account': task.account,
                'site': task.site,
                'sector': task.sector,
                'html_content': content_html or '',
                'word_count': word_count or 0,
                'status': 'draft',
            },
        )

        # Update content fields
        if content_html:
            content_record.html_content = content_html
        content_record.word_count = word_count or content_record.word_count or 0
        content_record.title = title
        content_record.meta_title = meta_title
        content_record.meta_description = meta_description
        content_record.primary_keyword = primary_keyword or ''
        if isinstance(secondary_keywords, list):
            content_record.secondary_keywords = secondary_keywords
        elif secondary_keywords:
            content_record.secondary_keywords = [secondary_keywords]
        else:
            content_record.secondary_keywords = []
        if isinstance(tags, list):
            content_record.tags = tags
        elif tags:
            content_record.tags = [tags]
        else:
            content_record.tags = []
        if isinstance(categories, list):
            content_record.categories = categories
        elif categories:
            content_record.categories = [categories]
        else:
            content_record.categories = []

        # Always set status to 'draft' for newly generated content
        # Status can only be: draft, review, published (changed manually)
        content_record.status = 'draft'

        # Merge any extra fields into metadata (non-standard keys)
        if isinstance(parsed, dict):
            excluded_keys = {
                'content',
                'title',
                'meta_title',
                'meta_description',
                'primary_keyword',
                'secondary_keywords',
                'tags',
                'categories',
                'word_count',
                'status',
            }
            extra_meta = {k: v for k, v in parsed.items() if k not in excluded_keys}
            existing_meta = content_record.metadata or {}
            existing_meta.update(extra_meta)
            content_record.metadata = existing_meta

        # Align foreign keys to ensure consistency
        content_record.account = task.account
        content_record.site = task.site
        content_record.sector = task.sector
        content_record.task = task

        content_record.save()

        # Update task status - keep task data intact but mark as completed
        task.status = 'completed'
        task.save(update_fields=['status', 'updated_at'])

        return {
            'count': 1,
            'tasks_updated': 1,
            'word_count': content_record.word_count,
        }


