"""
Automation Service
Core orchestrator that executes AI function stages sequentially
"""
import logging
import time
from datetime import datetime, timedelta
from typing import Dict, Optional, List
from django.db import transaction
from django.db.models import Count, Q, F
from django.core.cache import cache
from celery.result import AsyncResult

from igny8_core.business.automation.models import AutomationRun, AutomationConfig
from igny8_core.business.automation.services.automation_logger import AutomationLogger
from django.conf import settings
from igny8_core.auth.models import Account, Site
from igny8_core.modules.planner.models import Keywords, Clusters, ContentIdeas
from igny8_core.modules.writer.models import Tasks, Content, Images
from igny8_core.ai.models import AITaskLog
from igny8_core.ai.engine import AIEngine

# AI Functions
from igny8_core.ai.functions.auto_cluster import AutoClusterFunction
from igny8_core.ai.functions.generate_ideas import GenerateIdeasFunction
from igny8_core.ai.functions.generate_content import GenerateContentFunction
from igny8_core.ai.functions.generate_image_prompts import GenerateImagePromptsFunction
from igny8_core.ai.tasks import process_image_generation_queue

logger = logging.getLogger(__name__)


class AutomationService:
    """Orchestrates AI functions into automated pipeline"""
    
    def __init__(self, account: Account, site: Site):
        self.account = account
        self.site = site
        # Initialize AutomationLogger; allow optional shared log dir from settings
        shared_dir = getattr(settings, 'AUTOMATION_SHARED_LOG_DIR', None)
        self.logger = AutomationLogger(shared_log_dir=shared_dir)
        self.run = None
        self.config = None
        
        # Load or create config
        self.config, _ = AutomationConfig.objects.get_or_create(
            account=account,
            site=site,
            defaults={
                'is_enabled': False,
                'frequency': 'daily',
                'scheduled_time': '02:00',
            }
        )
    
    @classmethod
    def from_run_id(cls, run_id: str) -> 'AutomationService':
        """Create service instance from run_id"""
        run = AutomationRun.objects.get(run_id=run_id)
        service = cls(run.account, run.site)
        service.run = run
        return service
    
    def _check_should_stop(self) -> tuple[bool, str]:
        """
        Check if automation should stop (paused or cancelled)
        
        Returns:
            (should_stop, reason)
        """
        if not self.run:
            return False, ""
        
        # Refresh run from database
        self.run.refresh_from_db()
        
        if self.run.status == 'paused':
            return True, "paused"
        elif self.run.status == 'cancelled':
            return True, "cancelled"
        
        return False, ""
    
    def start_automation(self, trigger_type: str = 'manual') -> str:
        """
        Start automation run
        
        Returns:
            run_id
        """
        # Check for concurrent run
        if AutomationRun.objects.filter(site=self.site, status='running').exists():
            raise ValueError("Automation already running for this site")
        
        # Acquire distributed lock
        lock_key = f'automation_lock_{self.site.id}'
        if not cache.add(lock_key, 'locked', timeout=21600):  # 6 hours
            raise ValueError("Automation already running for this site (cache lock)")
        
        try:
            # Estimate credits needed
            estimated_credits = self.estimate_credits()
            
            # Check credit balance (with 20% buffer)
            required_credits = int(estimated_credits * 1.2)
            if self.account.credits < required_credits:
                raise ValueError(f"Insufficient credits. Need ~{required_credits}, you have {self.account.credits}")
            
            # Create run_id and log files
            run_id = self.logger.start_run(self.account.id, self.site.id, trigger_type)
            
            # Create AutomationRun record
            self.run = AutomationRun.objects.create(
                run_id=run_id,
                account=self.account,
                site=self.site,
                trigger_type=trigger_type,
                status='running',
                current_stage=1,
            )
            
            # Log start
            self.logger.log_stage_progress(
                run_id, self.account.id, self.site.id, 0,
                f"Automation started (trigger: {trigger_type})"
            )
            self.logger.log_stage_progress(
                run_id, self.account.id, self.site.id, 0,
                f"Credit check: Account has {self.account.credits} credits, estimated need: {estimated_credits} credits"
            )
            
            logger.info(f"[AutomationService] Started run: {run_id}")
            return run_id
            
        except Exception as e:
            # Release lock on failure
            cache.delete(lock_key)
            raise
    
    def run_stage_1(self):
        """Stage 1: Keywords → Clusters"""
        stage_number = 1
        stage_name = "Keywords → Clusters (AI)"
        start_time = time.time()
        
        # Query pending keywords
        pending_keywords = Keywords.objects.filter(
            site=self.site,
            status='new',
            cluster__isnull=True,
            disabled=False
        )
        
        total_count = pending_keywords.count()
        
        # NEW: Pre-stage validation for minimum keywords
        from igny8_core.ai.validators.cluster_validators import validate_minimum_keywords
        
        keyword_ids_for_validation = list(pending_keywords.values_list('id', flat=True))
        
        min_validation = validate_minimum_keywords(
            keyword_ids=keyword_ids_for_validation,
            account=self.account,
            min_required=5
        )
        
        if not min_validation['valid']:
            # Log validation failure
            self.logger.log_stage_start(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, stage_name, total_count
            )
            
            error_msg = min_validation['error']
            self.logger.log_stage_error(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, error_msg
            )
            
            # Skip stage with proper result
            self.run.stage_1_result = {
                'keywords_processed': 0,
                'clusters_created': 0,
                'batches_run': 0,
                'skipped': True,
                'skip_reason': error_msg,
                'credits_used': 0
            }
            self.run.current_stage = 2
            self.run.save()
            
            logger.warning(f"[AutomationService] Stage 1 skipped: {error_msg}")
            return
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No keywords to process - skipping stage"
            )
            self.run.stage_1_result = {'keywords_processed': 0, 'clusters_created': 0, 'batches_run': 0, 'credits_used': 0}
            self.run.current_stage = 2
            self.run.save()
            return
        
        # Process in batches with dynamic sizing
        batch_size = self.config.stage_1_batch_size
        # FIXED: Use min() for dynamic batch sizing
        actual_batch_size = min(total_count, batch_size)
        
        keywords_processed = 0
        clusters_created = 0
        batches_run = 0
        credits_before = self._get_credits_used()
        
        keyword_ids = list(pending_keywords.values_list('id', flat=True))
        
        for i in range(0, len(keyword_ids), actual_batch_size):
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({keywords_processed} keywords processed)"
                )
                # Save current progress
                credits_used = self._get_credits_used() - credits_before
                time_elapsed = self._format_time_elapsed(start_time)
                self.run.stage_1_result = {
                    'keywords_processed': keywords_processed,
                    'clusters_created': clusters_created,
                    'batches_run': batches_run,
                    'credits_used': credits_used,
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.total_credits_used += credits_used
                self.run.save()
                return
            
            try:
                batch = keyword_ids[i:i + actual_batch_size]
                batch_num = (i // actual_batch_size) + 1
                total_batches = (len(keyword_ids) + actual_batch_size - 1) // actual_batch_size
                
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Processing batch {batch_num}/{total_batches} ({len(batch)} keywords)"
                )
                
                # Call AI function via AIEngine
                engine = AIEngine(account=self.account)
                result = engine.execute(
                    fn=AutoClusterFunction(),
                    payload={'ids': batch}
                )
                
                # Monitor task
                task_id = result.get('task_id')
                if task_id:
                    # FIXED: Pass continue_on_error=True to keep processing other batches on failure
                    self._wait_for_task(task_id, stage_number, f"Batch {batch_num}", continue_on_error=True)
                
                keywords_processed += len(batch)
                batches_run += 1
                
                # Log progress
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Batch {batch_num} complete"
                )
                
                # Emit per-item trace event for UI progress tracking
                try:
                    self.logger.append_trace(self.account.id, self.site.id, self.run.run_id, {
                        'event': 'stage_item_processed',
                        'run_id': self.run.run_id,
                        'stage': stage_number,
                        'processed': keywords_processed,
                        'total': len(keyword_ids),
                        'batch_num': batch_num,
                        'timestamp': datetime.now().isoformat()
                    })
                except Exception:
                    pass
            except Exception as e:
                # FIXED: Log error but continue processing remaining batches
                error_msg = f"Failed to process batch {batch_num}: {str(e)}"
                logger.error(f"[AutomationService] {error_msg}", exc_info=True)
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                # Continue to next batch
                continue
            
            # ADDED: Within-stage delay (between batches)
            if i + actual_batch_size < len(keyword_ids):  # Not the last batch
                delay = self.config.within_stage_delay
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Waiting {delay} seconds before next batch..."
                )
                time.sleep(delay)
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, "Delay complete, resuming processing"
                )
        
        # Get clusters created count
        clusters_created = Clusters.objects.filter(
            site=self.site,
            created_at__gte=self.run.started_at
        ).count()
        
        # Calculate credits used
        credits_used = self._get_credits_used() - credits_before
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, keywords_processed, time_elapsed, credits_used
        )
        
        # ADDED: Post-stage validation
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Validation: {keywords_processed} keywords processed, {clusters_created} clusters created"
        )
        
        # Save results with timestamps
        from django.utils import timezone
        self.run.stage_1_result = {
            'keywords_processed': keywords_processed,
            'clusters_created': clusters_created,
            'batches_run': batches_run,
            'credits_used': credits_used,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 2
        self.run.total_credits_used += credits_used
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 1 complete: {keywords_processed} keywords → {clusters_created} clusters")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before next stage..."
        )
        time.sleep(delay)
    
    def run_stage_2(self):
        """Stage 2: Clusters → Ideas"""
        stage_number = 2
        stage_name = "Clusters → Ideas (AI)"
        start_time = time.time()
        
        # ADDED: Pre-stage validation - verify Stage 1 completion
        pending_keywords = Keywords.objects.filter(
            site=self.site,
            status='new',
            cluster__isnull=True,
            disabled=False
        ).count()
        
        if pending_keywords > 0:
            error_msg = f"Stage 1 incomplete: {pending_keywords} keywords still pending"
            self.logger.log_stage_error(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, error_msg
            )
            logger.error(f"[AutomationService] {error_msg}")
            # Continue anyway but log warning
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: Proceeding despite {pending_keywords} pending keywords from Stage 1"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Pre-stage validation passed: 0 keywords pending from Stage 1"
            )
        
        # Query clusters with status='new'
        pending_clusters = Clusters.objects.filter(
            site=self.site,
            status='new',
            disabled=False
        )
        
        total_count = pending_clusters.count()
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No clusters to process - skipping stage"
            )
            self.run.stage_2_result = {'clusters_processed': 0, 'ideas_created': 0, 'credits_used': 0}
            self.run.current_stage = 3
            self.run.save()
            return
        
        # Process one at a time
        clusters_processed = 0
        credits_before = self._get_credits_used()
        
        for cluster in pending_clusters:
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({clusters_processed} clusters processed)"
                )
                # Save current progress
                ideas_created = ContentIdeas.objects.filter(
                    site=self.site,
                    created_at__gte=self.run.started_at
                ).count()
                credits_used = self._get_credits_used() - credits_before
                time_elapsed = self._format_time_elapsed(start_time)
                from django.utils import timezone
                self.run.stage_2_result = {
                    'clusters_processed': clusters_processed,
                    'ideas_created': ideas_created,
                    'credits_used': credits_used,
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.total_credits_used += credits_used
                self.run.save()
                return
            
            try:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Generating ideas for cluster: {cluster.name}"
                )
                
                # Call AI function via AIEngine
                engine = AIEngine(account=self.account)
                result = engine.execute(
                    fn=GenerateIdeasFunction(),
                    payload={'ids': [cluster.id]}
                )
                
                # Monitor task
                task_id = result.get('task_id')
                if task_id:
                    # FIXED: Pass continue_on_error=True to keep processing other clusters on failure
                    self._wait_for_task(task_id, stage_number, f"Cluster '{cluster.name}'", continue_on_error=True)
                
                clusters_processed += 1
                
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Cluster '{cluster.name}' complete"
                )
            except Exception as e:
                # FIXED: Log error but continue processing remaining clusters
                error_msg = f"Failed to generate ideas for cluster '{cluster.name}': {str(e)}"
                logger.error(f"[AutomationService] {error_msg}", exc_info=True)
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                # Continue to next cluster
                continue
        
        # Get ideas created count
        ideas_created = ContentIdeas.objects.filter(
            site=self.site,
            created_at__gte=self.run.started_at
        ).count()
        
        # Calculate credits used
        credits_used = self._get_credits_used() - credits_before
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, clusters_processed, time_elapsed, credits_used
        )
        
        # Save results with timestamps
        from django.utils import timezone
        self.run.stage_2_result = {
            'clusters_processed': clusters_processed,
            'ideas_created': ideas_created,
            'credits_used': credits_used,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 3
        self.run.total_credits_used += credits_used
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 2 complete: {clusters_processed} clusters → {ideas_created} ideas")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before next stage..."
        )
        time.sleep(delay)
    
    def run_stage_3(self):
        """Stage 3: Ideas → Tasks (Local Queue)"""
        stage_number = 3
        stage_name = "Ideas → Tasks (Local Queue)"
        start_time = time.time()
        
        # ADDED: Pre-stage validation - verify Stage 2 completion
        pending_clusters = Clusters.objects.filter(
            site=self.site,
            status='new',
            disabled=False
        ).exclude(
            ideas__isnull=False
        ).count()
        
        if pending_clusters > 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: {pending_clusters} clusters from Stage 2 still pending"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Pre-stage validation passed: 0 clusters pending from Stage 2"
            )
        
        # Query pending ideas
        pending_ideas = ContentIdeas.objects.filter(
            site=self.site,
            status='new'
        )
        
        total_count = pending_ideas.count()
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No ideas to process - skipping stage"
            )
            self.run.stage_3_result = {'ideas_processed': 0, 'tasks_created': 0, 'batches_run': 0}
            self.run.current_stage = 4
            self.run.save()
            return
        
        # Process in batches
        batch_size = self.config.stage_3_batch_size
        ideas_processed = 0
        tasks_created = 0
        batches_run = 0
        
        idea_list = list(pending_ideas)
        
        for i in range(0, len(idea_list), batch_size):
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({ideas_processed} ideas processed, {tasks_created} tasks created)"
                )
                # Save current progress
                time_elapsed = self._format_time_elapsed(start_time)
                from django.utils import timezone
                self.run.stage_3_result = {
                    'ideas_processed': ideas_processed,
                    'tasks_created': tasks_created,
                    'batches_run': batches_run,
                    'started_at': self.run.started_at.isoformat(),
                    'completed_at': timezone.now().isoformat(),
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.save()
                return
            
            batch = idea_list[i:i + batch_size]
            batch_num = (i // batch_size) + 1
            total_batches = (len(idea_list) + batch_size - 1) // batch_size
            
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Queueing batch {batch_num}/{total_batches} ({len(batch)} ideas)"
            )
            
            # Create tasks (local operation)
            for idea in batch:
                # Build keywords string
                keywords_str = ''
                if idea.keyword_objects.exists():
                    keywords_str = ', '.join([kw.keyword for kw in idea.keyword_objects.all()])
                elif idea.target_keywords:
                    keywords_str = idea.target_keywords
                
                # Create task
                task = Tasks.objects.create(
                    title=idea.idea_title,
                    description=idea.description or '',
                    cluster=idea.keyword_cluster,
                    content_type=idea.content_type or 'post',
                    content_structure=idea.content_structure or 'article',
                    keywords=keywords_str,
                    status='queued',
                    account=idea.account,
                    site=idea.site,
                    sector=idea.sector,
                    idea=idea,
                )
                
                # Update idea status
                idea.status = 'queued'
                idea.save()
                
                tasks_created += 1
            
            ideas_processed += len(batch)
            batches_run += 1
            
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Batch {batch_num} complete: {len(batch)} tasks created"
            )
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, ideas_processed, time_elapsed, 0  # No credits for local operation
        )
        
        # Save results with timestamps
        from django.utils import timezone
        self.run.stage_3_result = {
            'ideas_processed': ideas_processed,
            'tasks_created': tasks_created,
            'batches_run': batches_run,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 4
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 3 complete: {ideas_processed} ideas → {tasks_created} tasks")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before next stage..."
        )
        time.sleep(delay)
    
    def run_stage_4(self):
        """Stage 4: Tasks → Content"""
        stage_number = 4
        stage_name = "Tasks → Content (AI)"
        start_time = time.time()
        
        # ADDED: Pre-stage validation - verify Stage 3 completion
        pending_ideas = ContentIdeas.objects.filter(
            site=self.site,
            status='new'
        ).count()
        
        if pending_ideas > 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: {pending_ideas} ideas from Stage 3 still pending"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Pre-stage validation passed: 0 ideas pending from Stage 3"
            )
        
        # Query queued tasks (all queued tasks need content generated)
        pending_tasks = Tasks.objects.filter(
            site=self.site,
            status='queued'
        )
        
        total_count = pending_tasks.count()
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No tasks to process - skipping stage"
            )
            self.run.stage_4_result = {'tasks_processed': 0, 'content_created': 0, 'total_words': 0, 'credits_used': 0}
            self.run.current_stage = 5
            self.run.save()
            return
        
        # Process one at a time
        tasks_processed = 0
        credits_before = self._get_credits_used()
        
        # FIXED: Ensure ALL tasks are processed by iterating over queryset list
        task_list = list(pending_tasks)
        total_tasks = len(task_list)
        
        for idx, task in enumerate(task_list, 1):
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({tasks_processed} tasks processed)"
                )
                # Save current progress
                content_created = Content.objects.filter(
                    site=self.site,
                    created_at__gte=self.run.started_at
                ).count()
                credits_used = self._get_credits_used() - credits_before
                time_elapsed = self._format_time_elapsed(start_time)
                from django.utils import timezone
                self.run.stage_4_result = {
                    'tasks_processed': tasks_processed,
                    'content_created': content_created,
                    'credits_used': credits_used,
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.total_credits_used += credits_used
                self.run.save()
                return
            
            try:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Generating content for task {idx}/{total_tasks}: {task.title}"
                )
                
                # Call AI function via AIEngine
                engine = AIEngine(account=self.account)
                result = engine.execute(
                    fn=GenerateContentFunction(),
                    payload={'ids': [task.id]}
                )
                
                # Monitor task
                task_id = result.get('task_id')
                if task_id:
                    # FIXED: Pass continue_on_error=True to keep processing other tasks on failure
                    self._wait_for_task(task_id, stage_number, f"Task '{task.title}'", continue_on_error=True)
                
                tasks_processed += 1
                
                # Log progress
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Task '{task.title}' complete ({tasks_processed}/{total_tasks})"
                )
                
                # Emit per-item trace event for UI progress tracking
                try:
                    self.logger.append_trace(self.account.id, self.site.id, self.run.run_id, {
                        'event': 'stage_item_processed',
                        'run_id': self.run.run_id,
                        'stage': stage_number,
                        'processed': tasks_processed,
                        'total': total_tasks,
                        'item': {'id': task.id, 'title': task.title},
                        'timestamp': datetime.now().isoformat()
                    })
                except Exception:
                    pass
            except Exception as e:
                # FIXED: Log error but continue processing remaining tasks
                error_msg = f"Failed to process task '{task.title}': {str(e)}"
                logger.error(f"[AutomationService] {error_msg}", exc_info=True)
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                # Continue to next task
                continue
            
            # ADDED: Within-stage delay between tasks (if not last task)
            if idx < total_tasks:
                delay = self.config.within_stage_delay
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Waiting {delay} seconds before next task..."
                )
                time.sleep(delay)
        
        # Get content created count and total words
        content_created = Content.objects.filter(
            site=self.site,
            created_at__gte=self.run.started_at
        ).count()
        
        total_words = Content.objects.filter(
            site=self.site,
            created_at__gte=self.run.started_at
        ).aggregate(total=Count('id'))['total'] * 2500  # Estimate
        
        # Calculate credits used
        credits_used = self._get_credits_used() - credits_before
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, tasks_processed, time_elapsed, credits_used
        )
        
        # ADDED: Post-stage validation - verify all tasks processed
        remaining_tasks = Tasks.objects.filter(
            site=self.site,
            status='queued'
        ).count()
        
        if remaining_tasks > 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: {remaining_tasks} tasks still queued after Stage 4"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Post-stage validation passed: 0 tasks remaining"
            )
        
        # Save results with timestamps
        from django.utils import timezone
        self.run.stage_4_result = {
            'tasks_processed': tasks_processed,
            'content_created': content_created,
            'total_words': total_words,
            'credits_used': credits_used,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 5
        self.run.total_credits_used += credits_used
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 4 complete: {tasks_processed} tasks → {content_created} content")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before next stage..."
        )
        time.sleep(delay)
    
    def run_stage_5(self):
        """Stage 5: Content → Image Prompts"""
        stage_number = 5
        stage_name = "Content → Image Prompts (AI)"
        start_time = time.time()
        
        # ADDED: Pre-stage validation - verify Stage 4 completion
        remaining_tasks = Tasks.objects.filter(
            site=self.site,
            status='queued'
        ).count()
        
        if remaining_tasks > 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: {remaining_tasks} tasks from Stage 4 still queued"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Pre-stage validation passed: 0 tasks pending from Stage 4"
            )
        
        # FIXED: Query content without Images records (ensure status='draft')
        content_without_images = Content.objects.filter(
            site=self.site,
            status='draft'  # Explicitly check for draft status
        ).annotate(
            images_count=Count('images')
        ).filter(
            images_count=0
        )
        
        total_count = content_without_images.count()
        
        # ADDED: Enhanced logging
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage 5: Found {total_count} content pieces without images (status='draft', images_count=0)"
        )
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No content to process - skipping stage"
            )
            self.run.stage_5_result = {'content_processed': 0, 'prompts_created': 0, 'credits_used': 0}
            self.run.current_stage = 6
            self.run.save()
            return
        
        # Process one at a time
        content_processed = 0
        credits_before = self._get_credits_used()
        
        content_list = list(content_without_images)
        total_content = len(content_list)
        
        for idx, content in enumerate(content_list, 1):
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({content_processed} content processed)"
                )
                # Save current progress
                prompts_created = Images.objects.filter(
                    site=self.site,
                    created_at__gte=self.run.started_at
                ).count()
                credits_used = self._get_credits_used() - credits_before
                time_elapsed = self._format_time_elapsed(start_time)
                from django.utils import timezone
                self.run.stage_5_result = {
                    'content_processed': content_processed,
                    'prompts_created': prompts_created,
                    'credits_used': credits_used,
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.total_credits_used += credits_used
                self.run.save()
                return
            
            try:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Extracting prompts {idx}/{total_content}: {content.title}"
                )
                
                # Call AI function via AIEngine
                engine = AIEngine(account=self.account)
                result = engine.execute(
                    fn=GenerateImagePromptsFunction(),
                    payload={'ids': [content.id]}
                )
                
                # Monitor task
                task_id = result.get('task_id')
                if task_id:
                    # FIXED: Pass continue_on_error=True to keep processing other content on failure
                    self._wait_for_task(task_id, stage_number, f"Content '{content.title}'", continue_on_error=True)
                
                content_processed += 1
                
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Content '{content.title}' complete ({content_processed}/{total_content})"
                )
            except Exception as e:
                # FIXED: Log error but continue processing remaining content
                error_msg = f"Failed to extract prompts for content '{content.title}': {str(e)}"
                logger.error(f"[AutomationService] {error_msg}", exc_info=True)
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                # Continue to next content
                continue
            
            # ADDED: Within-stage delay between content pieces
            if idx < total_content:
                delay = self.config.within_stage_delay
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Waiting {delay} seconds before next content..."
                )
                time.sleep(delay)
        
        # Get prompts created count
        prompts_created = Images.objects.filter(
            site=self.site,
            status='pending',
            created_at__gte=self.run.started_at
        ).count()
        
        # Calculate credits used
        credits_used = self._get_credits_used() - credits_before
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, content_processed, time_elapsed, credits_used
        )
        
        # Save results with timestamps
        from django.utils import timezone
        self.run.stage_5_result = {
            'content_processed': content_processed,
            'prompts_created': prompts_created,
            'credits_used': credits_used,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 6
        self.run.total_credits_used += credits_used
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 5 complete: {content_processed} content → {prompts_created} prompts")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before next stage..."
        )
        time.sleep(delay)
    
    def run_stage_6(self):
        """Stage 6: Image Prompts → Generated Images"""
        stage_number = 6
        stage_name = "Images (Prompts) → Generated Images (AI)"
        start_time = time.time()
        
        # ADDED: Pre-stage validation - verify Stage 5 completion
        content_without_images = Content.objects.filter(
            site=self.site,
            status='draft'
        ).annotate(
            images_count=Count('images')
        ).filter(
            images_count=0
        ).count()
        
        if content_without_images > 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Warning: {content_without_images} content pieces from Stage 5 still without images"
            )
        else:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "Pre-stage validation passed: All content has image prompts"
            )
        
        # Query pending images
        pending_images = Images.objects.filter(
            site=self.site,
            status='pending'
        )
        
        total_count = pending_images.count()
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        if total_count == 0:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, "No images to process - skipping stage"
            )
            self.run.stage_6_result = {'images_processed': 0, 'images_generated': 0, 'content_moved_to_review': 0, 'credits_used': 0}
            self.run.current_stage = 7
            self.run.save()
            return
        
        # Process one at a time
        images_processed = 0
        credits_before = self._get_credits_used()
        
        image_list = list(pending_images)
        total_images = len(image_list)
        
        for idx, image in enumerate(image_list, 1):
            # Check if automation should stop (paused or cancelled)
            should_stop, reason = self._check_should_stop()
            if should_stop:
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Stage {reason} - saving progress ({images_processed} images processed)"
                )
                # Save current progress
                images_generated = Images.objects.filter(
                    site=self.site,
                    status='completed',
                    created_at__gte=self.run.started_at
                ).count()
                content_moved = Content.objects.filter(
                    site=self.site,
                    status='review',
                    updated_at__gte=self.run.started_at
                ).count()
                credits_used = self._get_credits_used() - credits_before
                time_elapsed = self._format_time_elapsed(start_time)
                from django.utils import timezone
                self.run.stage_6_result = {
                    'images_processed': images_processed,
                    'images_generated': images_generated,
                    'content_moved_to_review': content_moved,
                    'credits_used': credits_used,
                    'time_elapsed': time_elapsed,
                    'partial': True,
                    'stopped_reason': reason
                }
                self.run.total_credits_used += credits_used
                self.run.save()
                return
            
            try:
                content_title = image.content.title if image.content else 'Unknown'
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Generating image {idx}/{total_images}: {image.image_type} for '{content_title}'"
                )
                
                # Call process_image_generation_queue directly (same as Writer/Images page)
                # Queue the task
                if hasattr(process_image_generation_queue, 'delay'):
                    task = process_image_generation_queue.delay(
                        image_ids=[image.id],
                        account_id=self.account.id,
                        content_id=image.content.id if image.content else None
                    )
                    task_id = str(task.id)
                else:
                    # Fallback for testing (synchronous)
                    result = process_image_generation_queue(
                        image_ids=[image.id],
                        account_id=self.account.id,
                        content_id=image.content.id if image.content else None
                    )
                    task_id = None
                
                # Monitor task (if async)
                if task_id:
                    # FIXED: Pass continue_on_error=True to keep processing other images on failure
                    self._wait_for_task(task_id, stage_number, f"Image for '{content_title}'", continue_on_error=True)
                
                images_processed += 1
                
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Image generated for '{content_title}' ({images_processed}/{total_images})"
                )
            except Exception as e:
                # FIXED: Log error but continue processing remaining images
                content_title = image.content.title if image.content else 'Unknown'
                error_msg = f"Failed to generate image for content '{content_title}': {str(e)}"
                logger.error(f"[AutomationService] {error_msg}", exc_info=True)
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                # Continue to next image
                continue
            
            # ADDED: Within-stage delay between images
            if idx < total_images:
                delay = self.config.within_stage_delay
                self.logger.log_stage_progress(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, f"Waiting {delay} seconds before next image..."
                )
                time.sleep(delay)
        
        # Get images generated count
        images_generated = Images.objects.filter(
            site=self.site,
            status='generated',
            updated_at__gte=self.run.started_at
        ).count()
        
        # Count content moved to review (automatic side effect)
        content_moved_to_review = Content.objects.filter(
            site=self.site,
            status='review',
            updated_at__gte=self.run.started_at
        ).count()
        
        # Calculate credits used
        credits_used = self._get_credits_used() - credits_before
        
        # Calculate time elapsed
        time_elapsed = self._format_time_elapsed(start_time)
        
        # Log completion
        self.logger.log_stage_complete(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, images_processed, time_elapsed, credits_used
        )
        
        # Save results with timestamps
        from django.utils import timezone
        stage_6_start = start_time  # Capture stage start time
        self.run.stage_6_result = {
            'images_processed': images_processed,
            'images_generated': images_generated,
            'content_moved_to_review': content_moved_to_review,
            'credits_used': credits_used,
            'started_at': self.run.started_at.isoformat(),
            'completed_at': timezone.now().isoformat(),
            'time_elapsed': time_elapsed
        }
        self.run.current_stage = 7
        self.run.total_credits_used += credits_used
        self.run.save()
        
        logger.info(f"[AutomationService] Stage 6 complete: {images_processed} images generated, {content_moved_to_review} content moved to review")
        
        # ADDED: Between-stage delay
        delay = self.config.between_stage_delay
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Stage complete. Waiting {delay} seconds before final stage..."
        )
        time.sleep(delay)
    
    def run_stage_7(self):
        """Stage 7: Manual Review Gate (Count Only)"""
        stage_number = 7
        stage_name = "Manual Review Gate"
        
        # Query content ready for review
        ready_for_review = Content.objects.filter(
            site=self.site,
            status='review'
        )
        
        total_count = ready_for_review.count()
        content_ids = list(ready_for_review.values_list('id', flat=True))
        
        # Log stage start
        self.logger.log_stage_start(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, stage_name, total_count
        )
        
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Automation complete. {total_count} content pieces ready for review"
        )
        
        if content_ids:
            self.logger.log_stage_progress(
                self.run.run_id, self.account.id, self.site.id,
                stage_number, f"Content IDs ready: {content_ids[:10]}..." if len(content_ids) > 10 else f"Content IDs ready: {content_ids}"
            )
        
        # Save results
        self.run.stage_7_result = {
            'ready_for_review': total_count,
            'content_ids': content_ids
        }
        self.run.status = 'completed'
        self.run.completed_at = datetime.now()
        self.run.save()
        
        # Release lock
        cache.delete(f'automation_lock_{self.site.id}')
        
        logger.info(f"[AutomationService] Stage 7 complete: Automation ended, {total_count} content ready for review")
    
    def pause_automation(self):
        """Pause current automation run"""
        if self.run:
            self.run.status = 'paused'
            self.run.save()
            logger.info(f"[AutomationService] Paused run: {self.run.run_id}")
    
    def resume_automation(self):
        """Resume paused automation run"""
        if self.run and self.run.status == 'paused':
            self.run.status = 'running'
            self.run.save()
            logger.info(f"[AutomationService] Resumed run: {self.run.run_id}")
    
    def estimate_credits(self) -> int:
        """Estimate total credits needed for automation"""
        # Count items
        keywords_count = Keywords.objects.filter(site=self.site, status='new', cluster__isnull=True, disabled=False).count()
        clusters_count = Clusters.objects.filter(site=self.site, status='new').exclude(ideas__isnull=False).count()
        ideas_count = ContentIdeas.objects.filter(site=self.site, status='new').count()
        tasks_count = Tasks.objects.filter(site=self.site, status='queued').count()
        content_count = Content.objects.filter(site=self.site, status='draft').annotate(images_count=Count('images')).filter(images_count=0).count()
        
        # Estimate credits
        clustering_credits = (keywords_count // 5) + 1 if keywords_count > 0 else 0  # 1 credit per 5 keywords
        ideas_credits = clusters_count * 2  # 2 credits per cluster
        content_credits = tasks_count * 5  # Assume 2500 words avg = 5 credits
        prompts_credits = content_count * 2  # Assume 4 prompts per content = 2 credits
        images_credits = content_count * 8  # Assume 4 images * 2 credits avg
        
        total = clustering_credits + ideas_credits + content_credits + prompts_credits + images_credits
        
        logger.info(f"[AutomationService] Estimated credits: {total}")
        return total
    
    # Helper methods
    
    def _wait_for_task(self, task_id: str, stage_number: int, item_name: str, continue_on_error: bool = True):
        """Wait for Celery task to complete
        
        Args:
            task_id: Celery task ID
            stage_number: Current stage number
            item_name: Name of item being processed
            continue_on_error: If True, log error but don't raise exception (default: True)
        """
        result = AsyncResult(task_id)
        
        # FIXED: Wrap result.ready() in try-catch to handle Celery backend errors
        max_retries = 100  # Maximum polling attempts (100 * 3s = 5 minutes)
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                is_ready = result.ready()
                if is_ready:
                    break
            except (TypeError, ValueError, Exception) as e:
                # FIXED: Handle Celery backend errors (e.g., Redis key serialization issues)
                logger.warning(f"[AutomationService] Celery backend error while checking task {task_id}: {e}")
                # Wait a bit longer and try to get result directly
                time.sleep(5)
                try:
                    # Try to get the result status from database/backend
                    if result.state in ['SUCCESS', 'FAILURE']:
                        break
                except Exception:
                    # If still failing, assume task is done and break
                    logger.warning(f"[AutomationService] Cannot determine task state, assuming complete")
                    break
            
            time.sleep(3)  # Poll every 3 seconds
            retry_count += 1
            
            # Check for pause
            self.run.refresh_from_db()
            if self.run.status == 'paused':
                logger.info(f"[AutomationService] Paused during {item_name}")
                # Wait until resumed
                while self.run.status == 'paused':
                    time.sleep(5)
                    self.run.refresh_from_db()
        
        # Check if task failed
        try:
            if result.failed():
                error_msg = f"Task failed for {item_name}"
                self.logger.log_stage_error(
                    self.run.run_id, self.account.id, self.site.id,
                    stage_number, error_msg
                )
                if not continue_on_error:
                    raise Exception(error_msg)
                else:
                    logger.warning(f"[AutomationService] {error_msg} - continuing with next item")
        except (TypeError, ValueError, Exception) as e:
            # FIXED: If we can't determine task status due to backend error, log and continue
            logger.warning(f"[AutomationService] Cannot determine task failure state for {item_name}: {e}")
            if not continue_on_error:
                raise
    
    def _get_credits_used(self) -> int:
        """Get total credits used by this run so far"""
        if not self.run:
            return 0
        
        total = AITaskLog.objects.filter(
            account=self.account,
            created_at__gte=self.run.started_at
        ).aggregate(total=Count('id'))['total'] or 0
        
        return total
    
    def _format_time_elapsed(self, start_time: float) -> str:
        """Format elapsed time"""
        elapsed = time.time() - start_time
        minutes = int(elapsed // 60)
        seconds = int(elapsed % 60)
        return f"{minutes}m {seconds}s"
    
    def get_current_processing_state(self) -> dict:
        """
        Get real-time processing state for current automation run
        Returns detailed info about what's currently being processed
        """
        if not self.run:
            return None
        
        # Allow paused runs to show state (UI needs this to avoid blanking)
        if self.run.status not in ('running', 'paused'):
            return None
        
        stage = self.run.current_stage

        # Get stage-specific data based on current stage
        state = None
        if stage == 1:  # Keywords → Clusters
            state = self._get_stage_1_state()
        elif stage == 2:  # Clusters → Ideas
            state = self._get_stage_2_state()
        elif stage == 3:  # Ideas → Tasks
            state = self._get_stage_3_state()
        elif stage == 4:  # Tasks → Content
            state = self._get_stage_4_state()
        elif stage == 5:  # Content → Image Prompts
            state = self._get_stage_5_state()
        elif stage == 6:  # Image Prompts → Images
            state = self._get_stage_6_state()
        elif stage == 7:  # Manual Review Gate
            state = self._get_stage_7_state()

        # Trace what we return so UI mapping can be diagnosed
        try:
            if state is not None and self.logger:
                # Avoid heavy payloads for very large queues — only include keys we care about
                lightweight = {
                    'stage_number': state.get('stage_number'),
                    'stage_name': state.get('stage_name'),
                    'stage_type': state.get('stage_type'),
                    'total_items': state.get('total_items'),
                    'processed_items': state.get('processed_items'),
                    'percentage': state.get('percentage'),
                    'currently_processing_count': len(state.get('currently_processing') or []),
                    'up_next_count': len(state.get('up_next') or []),
                    'remaining_count': state.get('remaining_count'),
                }
                trace = {
                    'event': 'get_current_processing',
                    'run_id': getattr(self.run, 'run_id', None),
                    'site_id': getattr(self.site, 'id', None),
                    'account_id': getattr(self.account, 'id', None),
                    'payload': lightweight,
                    'timestamp': datetime.now().isoformat(),
                }
                # Best-effort append
                try:
                    self.logger.append_trace(self.account.id, self.site.id, self.run.run_id, trace)
                except Exception:
                    # fallback: try writing to diagnostic
                    try:
                        self.logger._append_to_main_log(self.account.id, self.site.id, self.run.run_id, f"TRACE: {trace}")
                    except Exception:
                        pass
        except Exception:
            pass

        return state
    
    def _get_stage_1_state(self) -> dict:
        """Get processing state for Stage 1: Keywords → Clusters"""
        queue = Keywords.objects.filter(
            site=self.site, status='new'
        ).order_by('id')
        
        processed = self._get_processed_count(1)
        total = queue.count() + processed
        
        return {
            'stage_number': 1,
            'stage_name': 'Keywords → Clusters',
            'stage_type': 'AI',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 3),
            'up_next': self._get_next_items(queue, 2, skip=3),
            'remaining_count': queue.count()
        }
    
    def _get_stage_2_state(self) -> dict:
        """Get processing state for Stage 2: Clusters → Ideas"""
        queue = Clusters.objects.filter(
            site=self.site, status='new', disabled=False
        ).order_by('id')
        
        processed = self._get_processed_count(2)
        total = queue.count() + processed
        
        return {
            'stage_number': 2,
            'stage_name': 'Clusters → Ideas',
            'stage_type': 'AI',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 1),
            'up_next': self._get_next_items(queue, 2, skip=1),
            'remaining_count': queue.count()
        }
    
    def _get_stage_3_state(self) -> dict:
        """Get processing state for Stage 3: Ideas → Tasks"""
        queue = ContentIdeas.objects.filter(
            site=self.site, status='approved'
        ).order_by('id')
        
        processed = self._get_processed_count(3)
        total = queue.count() + processed
        
        return {
            'stage_number': 3,
            'stage_name': 'Ideas → Tasks',
            'stage_type': 'Local',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 1),
            'up_next': self._get_next_items(queue, 2, skip=1),
            'remaining_count': queue.count()
        }
    
    def _get_stage_4_state(self) -> dict:
        """Get processing state for Stage 4: Tasks → Content"""
        queue = Tasks.objects.filter(
            site=self.site, status='ready'
        ).order_by('id')
        
        processed = self._get_processed_count(4)
        total = queue.count() + processed
        
        return {
            'stage_number': 4,
            'stage_name': 'Tasks → Content',
            'stage_type': 'AI',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 1),
            'up_next': self._get_next_items(queue, 2, skip=1),
            'remaining_count': queue.count()
        }
    
    def _get_stage_5_state(self) -> dict:
        """Get processing state for Stage 5: Content → Image Prompts"""
        queue = Content.objects.filter(
            site=self.site,
            status='draft'
        ).annotate(
            images_count=Count('images')
        ).filter(
            images_count=0
        ).order_by('id')
        
        processed = self._get_processed_count(5)
        total = queue.count() + processed
        
        return {
            'stage_number': 5,
            'stage_name': 'Content → Image Prompts',
            'stage_type': 'AI',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 1),
            'up_next': self._get_next_items(queue, 2, skip=1),
            'remaining_count': queue.count()
        }
    
    def _get_stage_6_state(self) -> dict:
        """Get processing state for Stage 6: Image Prompts → Images"""
        queue = Images.objects.filter(
            site=self.site, status='pending'
        ).order_by('id')
        
        processed = self._get_processed_count(6)
        total = queue.count() + processed
        
        return {
            'stage_number': 6,
            'stage_name': 'Image Prompts → Images',
            'stage_type': 'AI',
            'total_items': total,
            'processed_items': processed,
            'percentage': round((processed / total * 100) if total > 0 else 0),
            'currently_processing': self._get_current_items(queue, 1),
            'up_next': self._get_next_items(queue, 2, skip=1),
            'remaining_count': queue.count()
        }
    
    def _get_stage_7_state(self) -> dict:
        """Get processing state for Stage 7: Manual Review Gate"""
        queue = Content.objects.filter(
            site=self.site, status='review'
        ).order_by('id')
        
        total = queue.count()
        
        return {
            'stage_number': 7,
            'stage_name': 'Manual Review Gate',
            'stage_type': 'Manual',
            'total_items': total,
            'processed_items': total,
            'percentage': 100,
            'currently_processing': [],
            'up_next': self._get_current_items(queue, 3),
            'remaining_count': total
        }
    
    def _get_processed_count(self, stage: int) -> int:
        """Get count of items processed in current stage during this run"""
        if not self.run:
            return 0
        
        # Count items that were updated during this run and changed status from pending
        if stage == 1:
            # Keywords that changed status from 'new' during this run
            return Keywords.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at
            ).exclude(status='new').count()
        elif stage == 2:
            # Clusters that changed status from 'new' during this run
            return Clusters.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at
            ).exclude(status='new').count()
        elif stage == 3:
            # Ideas that changed status from 'approved' during this run
            return ContentIdeas.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at
            ).exclude(status='approved').count()
        elif stage == 4:
            # Tasks that changed status from 'ready'/'queued' during this run
            return Tasks.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at
            ).exclude(status__in=['ready', 'queued']).count()
        elif stage == 5:
            # Content processed for image prompts during this run
            return Content.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at,
                images__isnull=False
            ).distinct().count()
        elif stage == 6:
            # Images completed during this run
            return Images.objects.filter(
                site=self.site,
                updated_at__gte=self.run.started_at,
                status='completed'
            ).count()
        
        return 0
    
    def _get_current_items(self, queryset, count: int) -> list:
        """Get currently processing items"""
        items = queryset[:count]
        return [
            {
                'id': item.id,
                'title': self._get_item_title(item),
                'type': queryset.model.__name__.lower()
            }
            for item in items
        ]
    
    def _get_next_items(self, queryset, count: int, skip: int = 0) -> list:
        """Get next items in queue"""
        items = queryset[skip:skip + count]
        return [
            {
                'id': item.id,
                'title': self._get_item_title(item),
                'type': queryset.model.__name__.lower()
            }
            for item in items
        ]
    
    def _get_item_title(self, item) -> str:
        """Extract title from various model types"""
        # Try different title fields based on model type
        if hasattr(item, 'keyword'):
            return item.keyword
        elif hasattr(item, 'cluster_name'):
            return item.cluster_name
        elif hasattr(item, 'idea_title'):
            return item.idea_title
        elif hasattr(item, 'title'):
            return item.title
        elif hasattr(item, 'image_type') and hasattr(item, 'content'):
            content_title = item.content.title if item.content else 'Unknown'
            return f"{item.image_type} for '{content_title}'"
        
        return 'Unknown'
