"""
Generate Ideas AI Function
Extracted from modules/planner/tasks.py
"""
import logging
import json
from typing import Dict, List, Any
from django.db import transaction
from igny8_core.ai.base import BaseAIFunction
from igny8_core.modules.planner.models import Clusters, ContentIdeas
from igny8_core.ai.ai_core import AICore
from igny8_core.ai.validators import validate_cluster_exists, validate_cluster_limits
from igny8_core.ai.prompts import PromptRegistry
from igny8_core.ai.settings import get_model_config

logger = logging.getLogger(__name__)


class GenerateIdeasFunction(BaseAIFunction):
    """Generate content ideas from clusters using AI"""
    
    def get_name(self) -> str:
        return 'generate_ideas'
    
    def get_metadata(self) -> Dict:
        return {
            'display_name': 'Generate Ideas',
            'description': 'Generate SEO-optimized content ideas from keyword clusters',
            'phases': {
                'INIT': 'Initializing idea generation...',
                'PREP': 'Loading clusters...',
                'AI_CALL': 'Generating ideas with AI...',
                'PARSE': 'Parsing idea data...',
                'SAVE': 'Saving ideas...',
                'DONE': 'Ideas generated!'
            }
        }
    
    def get_max_items(self) -> int:
        return 10  # Max clusters per idea generation
    
    def validate(self, payload: dict, account=None) -> Dict:
        """Validate cluster IDs and plan limits"""
        result = super().validate(payload, account)
        if not result['valid']:
            return result
        
        # Check cluster exists
        cluster_ids = payload.get('ids', [])
        if cluster_ids:
            cluster_id = cluster_ids[0]  # For single cluster idea generation
            cluster_result = validate_cluster_exists(cluster_id, account)
            if not cluster_result['valid']:
                return cluster_result
        
        # Check plan limits
        limit_result = validate_cluster_limits(account, operation_type='idea')
        if not limit_result['valid']:
            return limit_result
        
        return {'valid': True}
    
    def prepare(self, payload: dict, account=None) -> Dict:
        """Load clusters with keywords"""
        cluster_ids = payload.get('ids', [])
        if not cluster_ids:
            raise ValueError("No cluster IDs provided")
        
        # Support multiple clusters (up to get_max_items())
        queryset = Clusters.objects.filter(id__in=cluster_ids)
        if account:
            queryset = queryset.filter(account=account)
        
        clusters = list(queryset.select_related('sector', 'account', 'site', 'sector__site').prefetch_related('keywords'))
        
        if not clusters:
            raise ValueError("No clusters found")
        
        # Get keywords for each cluster
        from igny8_core.modules.planner.models import Keywords
        cluster_data = []
        for cluster in clusters:
            # Get keywords and extract the keyword text from seed_keyword relationship
            keyword_objects = Keywords.objects.filter(cluster=cluster).select_related('seed_keyword')
            keywords = [kw.seed_keyword.keyword for kw in keyword_objects if kw.seed_keyword]
            cluster_data.append({
                'id': cluster.id,
                'name': cluster.name,
                'description': cluster.description or '',
                'keywords': keywords,
            })
        
        # Get account from first cluster if not provided
        account = account or (clusters[0].account if clusters else None)
        
        return {
            'clusters': clusters,  # List of cluster objects
            'cluster_data': cluster_data,  # Formatted data for AI
            'account': account
        }
    
    def build_prompt(self, data: Dict, account=None) -> str:
        """Build ideas generation prompt using registry"""
        cluster_data = data['cluster_data']
        account = account or data.get('account')
        
        # Format clusters text
        clusters_text = '\n'.join([
            f"Cluster ID: {c.get('id', '')} | Name: {c.get('name', '')} | Description: {c.get('description', '')}"
            for c in cluster_data
        ])
        
        # Format cluster keywords
        cluster_keywords_text = '\n'.join([
            f"Cluster ID: {c.get('id', '')} | Name: {c.get('name', '')} | Keywords: {', '.join(c.get('keywords', []))}"
            for c in cluster_data
        ])
        
        # Get prompt from registry with context
        prompt = PromptRegistry.get_prompt(
            function_name='generate_ideas',
            account=account,
            context={
                'CLUSTERS': clusters_text,
                'CLUSTER_KEYWORDS': cluster_keywords_text,
            }
        )
        
        return prompt
    
    def parse_response(self, response: str, step_tracker=None) -> List[Dict]:
        """Parse AI response into idea data"""
        ai_core = AICore(account=self.account if hasattr(self, 'account') else None)
        json_data = ai_core.extract_json(response)
        
        if not json_data or 'ideas' not in json_data:
            error_msg = f"Failed to parse ideas response: {response[:200]}..."
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        return json_data.get('ideas', [])
    
    def save_output(
        self,
        parsed: List[Dict],
        original_data: Dict,
        account=None,
        progress_tracker=None,
        step_tracker=None
    ) -> Dict:
        """Save ideas to database"""
        clusters = original_data['clusters']  # List of cluster objects
        cluster_data = original_data['cluster_data']  # Formatted data for matching
        account = account or original_data.get('account')
        
        if not account:
            raise ValueError("Account is required for idea creation")
        
        ideas_created = 0
        
        with transaction.atomic():
            for idx, idea_data in enumerate(parsed):
                # Find matching cluster by ID or name
                cluster = None
                cluster_id_from_ai = idea_data.get('cluster_id')
                cluster_name = idea_data.get('cluster_name', '')
                
                # Try to match by ID first
                if cluster_id_from_ai:
                    for c in clusters:
                        if c.id == cluster_id_from_ai:
                            cluster = c
                            break
                
                # Fallback to name matching
                if not cluster and cluster_name:
                    for c in clusters:
                        if c.name == cluster_name:
                            cluster = c
                            break
                
                # If still no match, use position-based matching (first idea -> first cluster, etc.)
                if not cluster and len(clusters) > 0:
                    cluster_index = idx % len(clusters)
                    cluster = clusters[cluster_index]
                    logger.warning(f"Cluster not found by ID/name for idea '{idea_data.get('title', 'Untitled')}', using cluster at index {cluster_index}")
                
                if not cluster:
                    logger.warning(f"Cluster not found for idea '{idea_data.get('title', 'Untitled')}', skipping")
                    continue
                
                # Ensure site is available
                site = cluster.site
                if not site and cluster.sector:
                    site = cluster.sector.site
                
                if not site:
                    logger.error(f"Site not found for cluster {cluster.id}, cannot create ContentIdeas")
                    continue
                
                # Handle description - might be dict or string
                description = idea_data.get('description', '')
                if isinstance(description, dict):
                    description = json.dumps(description)
                elif not isinstance(description, str):
                    description = str(description)
                
                # Handle target_keywords
                target_keywords = idea_data.get('covered_keywords', '') or idea_data.get('target_keywords', '')
                
                # Create ContentIdeas record
                ContentIdeas.objects.create(
                    idea_title=idea_data.get('title', 'Untitled Idea'),
                    description=description,
                    content_type=idea_data.get('content_type', 'blog_post'),
                    content_structure=idea_data.get('content_structure', 'supporting_page'),
                    target_keywords=target_keywords,
                    keyword_cluster=cluster,
                    estimated_word_count=idea_data.get('estimated_word_count', 1500),
                    status='new',
                    account=account,
                    site=site,
                    sector=cluster.sector,
                )
                ideas_created += 1
        
        return {
            'count': ideas_created,
            'ideas_created': ideas_created
        }


