#!/usr/bin/env python3
"""
Celery app for async batch prediction jobs.
Training jobs continue to use the existing featrix_queue system.
"""

from celery import Celery
from config import config
import json
import traceback
import sys
import pickle
import gc
import logging
from pathlib import Path

# Set up proper logging for Celery
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)-8s] %(name)-45s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Install Featrix exception hook for better error tracking
try:
    from lib.featrix_debug import install_featrix_excepthook
    install_featrix_excepthook()
except Exception:
    pass  # Don't fail if debug module not available

# CRITICAL: Set up Python path BEFORE any featrix imports to avoid circular dependencies
lib_path = Path(__file__).parent / "lib"
current_path = Path(__file__).parent

if str(lib_path.resolve()) not in sys.path:
    sys.path.insert(0, str(lib_path.resolve()))
if str(current_path.resolve()) not in sys.path:
    sys.path.insert(0, str(current_path.resolve()))

logger.info(f"🔧 Celery module loading - Python paths set up")
logger.info(f"  lib_path: {lib_path.resolve()}")
logger.info(f"  current_path: {current_path.resolve()}")

# Import NON-CUDA modules at module level only
# CRITICAL: Do NOT import torch here - causes CUDA fork issues
try:
    import numpy as np
    import pandas as pd
    logger.info(f"✅ Non-CUDA modules imported at module level")
except ImportError as e:
    logger.warning(f"⚠️ Core modules not available: {e}")

# Will check torch availability inside tasks after fork
TORCH_AVAILABLE = False

try:
    from featrix_queue import load_session
    FEATRIX_QUEUE_AVAILABLE = True
    logger.info(f"✅ featrix_queue imported at module level")
except ImportError as e:
    FEATRIX_QUEUE_AVAILABLE = False
    logger.warning(f"⚠️ featrix_queue not available: {e}")

# CRITICAL: Do NOT import featrix modules at module level
# They contain torch imports which cause CUDA fork issues in Celery
# Will import them inside tasks after the fork happens
logger.info(f"🔧 Skipping featrix imports at module level to avoid CUDA fork issues")
logger.info(f"🔧 Will import torch and featrix modules inside task after fork")

# Initialize Celery app
app = Celery('sphere_predictions')

# Configure Celery to use Redis as broker and result backend
app.conf.update(
    broker_url='redis://localhost:6379/1',
    result_backend='redis://localhost:6379/1',
    task_serializer='json',
    accept_content=['json'],
    result_serializer='json',
    timezone='America/New_York',
    enable_utc=True,
    result_expires=3600,  # Results expire after 1 hour
    task_routes={
        'celery_app.predict_batch': {'queue': 'predictions'},
    },
    worker_prefetch_multiplier=1,  # One task at a time for GPU jobs
    worker_max_tasks_per_child=10,  # Restart worker after 10 tasks to prevent memory leaks
    task_acks_late=True,
    worker_disable_rate_limits=True,
    # Logging configuration
    worker_log_format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    worker_task_log_format='%(asctime)s - %(name)s - %(levelname)s - [%(task_name)s(%(task_id)s)] %(message)s',
    worker_hijack_root_logger=False,  # Let our logging config work
)

@app.task(bind=True)
def predict_batch(self, session_id: str, records: list, prediction_options: dict = None):
    """
    Celery task for batch predictions.
    
    Args:
        session_id: Session ID with trained single predictor
        records: List of record dictionaries to predict
        prediction_options: Optional parameters (batch_size, etc.)
    
    Returns:
        dict with predictions or error info
    """
    import pickle
    from pathlib import Path
    
    logger.info(f"🚀 CELERY TASK STARTED - Session: {session_id}, Records: {len(records)}")
    
    try:
        # Update task state to indicate processing has started
        self.update_state(
            state='PROGRESS',
            meta={
                'current': 0,
                'total': len(records),
                'status': 'Loading model...'
            }
        )
        
        # Paths already set up at module level - no need to repeat
        logger.info(f"🔧 Python paths already configured at module level")
        
        # Load the session and single predictor
        if not FEATRIX_QUEUE_AVAILABLE:
            raise ImportError("featrix_queue module not available")
            
        logger.info(f"🔧 Loading session {session_id}...")
        try:
            session = load_session(session_id)
            logger.info(f"✅ Session loaded successfully")
        except Exception as session_error:
            logger.error(f"❌ SESSION LOAD FAILED: {session_error}")
            logger.error(f"❌ Session traceback: {traceback.format_exc()}")
            raise
        
        # Get the single predictor path (backwards compatible)
        single_predictors = session.get("single_predictors")
        single_predictor = session.get("single_predictor")
        
        # Set prediction options (need this before accessing it below)
        options = prediction_options or {}
        
        # Helper function to generate predictor ID
        def generate_predictor_id(predictor_path: str) -> str:
            import hashlib
            import os
            filename = os.path.basename(predictor_path) if predictor_path else 'unknown'
            # Remove .pickle extension from user-facing ID
            if filename.endswith('.pickle'):
                filename = filename[:-7]
            path_hash = hashlib.md5(predictor_path.encode('utf-8')).hexdigest()[:8]
            return f"{filename}_{path_hash}"
        
        # NEW: Check if predictor_id or target_column specified in options
        requested_predictor_id = options.get('predictor_id')
        requested_target_column = options.get('target_column')
        logger.info(f"🔍 CELERY PREDICT_BATCH - predictor_id='{requested_predictor_id}', target_column='{requested_target_column}'")
        
        predictor_path = None
        if single_predictors and isinstance(single_predictors, list):
            # NEW: If multiple predictors exist, REQUIRE target_column or predictor_id
            if len(single_predictors) > 1 and not requested_predictor_id and not requested_target_column:
                # Get available targets for error message
                available_targets = []
                available_ids = []
                for pred_path in single_predictors:
                    # Get predictor_id
                    if Path(pred_path).exists():
                        available_ids.append(generate_predictor_id(pred_path))
                    
                    # Get target_column from metadata
                    metadata_path = str(Path(pred_path).parent / "model_metadata.json")
                    if Path(metadata_path).exists():
                        try:
                            import json
                            with open(metadata_path, 'r') as f:
                                metadata = json.load(f)
                                target_col = metadata.get('target_column')
                                if target_col:
                                    available_targets.append(target_col)
                        except Exception:
                            pass
                
                error_msg = f"Multiple predictors found ({len(single_predictors)}) - must specify 'target_column' or 'predictor_id'. Available targets: {available_targets}, Available IDs: {available_ids}"
                logger.error(f"❌ CELERY PREDICT_BATCH - {error_msg}")
                raise ValueError(error_msg)
            
            # PRIORITY 1: predictor_id (most precise)
            if requested_predictor_id:
                logger.info(f"🔍 CELERY PREDICT_BATCH - Looking for predictor with predictor_id='{requested_predictor_id}'")
                predictor_found = False
                available_predictor_ids = []
                
                for pred_path in single_predictors:
                    if Path(pred_path).exists():
                        current_predictor_id = generate_predictor_id(pred_path)
                        available_predictor_ids.append(current_predictor_id)
                        
                        if current_predictor_id == requested_predictor_id:
                            predictor_path = pred_path
                            predictor_found = True
                            logger.info(f"✅ CELERY PREDICT_BATCH - Found matching predictor by ID: {pred_path}")
                            break
                
                if not predictor_found:
                    raise ValueError(f"No predictor found for predictor_id '{requested_predictor_id}'. Available IDs: {available_predictor_ids}")
            
            # PRIORITY 2: target_column (semantic match)
            elif requested_target_column:
                logger.info(f"🔍 CELERY PREDICT_BATCH - Looking for predictor with target_column='{requested_target_column}'")
                predictor_found = False
                
                for pred_path in single_predictors:
                    # Skip None or invalid paths
                    if not pred_path:
                        logger.warning(f"⚠️  CELERY PREDICT_BATCH - Skipping None predictor path in list")
                        continue
                    
                    # Load predictor metadata to check target_column
                    metadata_path = str(Path(pred_path).parent / "model_metadata.json")
                    if Path(metadata_path).exists():
                        try:
                            import json
                            with open(metadata_path, 'r') as f:
                                metadata = json.load(f)
                                target_col = metadata.get('target_column')
                                
                                logger.info(f"🔍 CELERY PREDICT_BATCH - Checking predictor {pred_path}: target_column='{target_col}'")
                                
                                if target_col == requested_target_column:
                                    predictor_path = pred_path
                                    predictor_found = True
                                    logger.info(f"✅ CELERY PREDICT_BATCH - Found matching predictor: {pred_path}")
                                    break
                        except Exception as metadata_error:
                            logger.warning(f"⚠️  CELERY PREDICT_BATCH - Could not read metadata for {pred_path}: {metadata_error}")
                            continue
                
                if not predictor_found:
                    raise ValueError(f"No predictor found for target column '{requested_target_column}'")
            
            # PRIORITY 3: Single predictor fallback (backwards compatibility)
            else:
                # Only allow fallback if there's exactly one predictor
                if len(single_predictors) == 1:
                    predictor_path = single_predictors[0]
                    logger.info(f"🔍 CELERY PREDICT_BATCH - Single predictor session, using: {predictor_path}")
                else:
                    # This should never be reached due to check above, but keep for safety
                    raise ValueError(f"Multiple predictors found ({len(single_predictors)}) but no target_column or predictor_id specified")
        elif single_predictor:
            predictor_path = single_predictor
        else:
            raise ValueError("No single predictor found in session")
        
        if not predictor_path or not Path(predictor_path).exists():
            raise ValueError(f"Single predictor not found at: {predictor_path}")
        
        # Update progress
        self.update_state(
            state='PROGRESS',
            meta={
                'current': 0,
                'total': len(records),
                'status': 'Loading single predictor model...'
            }
        )
        
        # NEW APPROACH: Use subprocess to avoid all import/fork/CUDA issues
        logger.info(f"🚀 Using subprocess approach to avoid fork/CUDA issues")
        logger.info(f"📦 Predictor path: {predictor_path}")
        
        # Create temporary files for input/output
        import tempfile
        import uuid
        import subprocess
        import json
        import redis
        import time
        
        # Initialize Redis for progress monitoring
        try:
            redis_client = redis.Redis(host='localhost', port=6379, db=1, decode_responses=True)
            redis_available = True
            logger.info(f"✅ Redis connected for progress monitoring")
        except Exception as redis_error:
            logger.warning(f"⚠️ Redis not available for progress monitoring: {redis_error}")
            redis_available = False
        
        temp_dir = Path("/tmp") / f"celery_prediction_{uuid.uuid4().hex[:8]}"
        temp_dir.mkdir(exist_ok=True)
        
        input_file = temp_dir / "input.json"
        output_file = temp_dir / "output.json"
        
        logger.info(f"📁 Created temp directory: {temp_dir}")
        
        # Set prediction options
        options = prediction_options or {}
        
        try:
            # Write input data to file
            with open(input_file, 'w') as f:
                json.dump({
                    'session_id': session_id,
                    'predictor_path': predictor_path,
                    'records': records,
                    'batch_size': options.get('batch_size', 256)
                }, f)
            logger.info(f"✅ Written {len(records)} records to {input_file}")
            
            # Run subprocess prediction
            script_path = Path(__file__).parent / "standalone_prediction.py"
            
            logger.info(f"🔧 Running subprocess: python {script_path} {input_file} {output_file}")
            
            # Use subprocess to call the standalone prediction script
            import sys
            
            # Get task ID for Redis progress tracking
            task_id = self.request.id
            redis_key = f"prediction_progress:{task_id}"
            
            # Start subprocess asynchronously - write output to log file to avoid pipe deadlock
            logger.info(f"🚀 Starting subprocess with Redis progress tracking...")
            
            # Create log files in temp directory
            subprocess_log = temp_dir / "subprocess.log"
            subprocess_err = temp_dir / "subprocess_error.log"
            
            with open(subprocess_log, 'w') as stdout_file, open(subprocess_err, 'w') as stderr_file:
                process = subprocess.Popen(
                    [sys.executable, str(script_path), str(input_file), str(output_file), task_id],
                    stdout=stdout_file,
                    stderr=stderr_file,
                    text=True
                )
                
                # Poll for subprocess completion while monitoring Redis progress
                start_time = time.time()
                last_progress_percentage = 0
                
                while process.poll() is None:
                    # Check for timeout (1 hour)
                    if time.time() - start_time > 3600:
                        process.terminate()
                        process.wait()
                        raise Exception("Prediction subprocess timed out after 1 hour")
                    
                    # Check Redis for progress updates
                    if redis_available:
                        try:
                            progress_data = redis_client.get(redis_key)
                            if progress_data:
                                progress_info = json.loads(progress_data)
                                current = progress_info.get('current', 0)
                                total = progress_info.get('total', 100)
                                status = progress_info.get('status', 'Processing...')
                                percentage = progress_info.get('percentage', 0)
                                
                                # Only update Celery state if progress percentage has changed
                                if percentage != last_progress_percentage:
                                    self.update_state(
                                        state='PROGRESS',
                                        meta={
                                            'current': current,
                                            'total': total,
                                            'status': status,
                                            'percentage': percentage,
                                            'timestamp': progress_info.get('timestamp')
                                        }
                                    )
                                    last_progress_percentage = percentage
                                    logger.info(f"📊 Relayed progress: {percentage}% - {status}")
                        except Exception as redis_err:
                            logger.warning(f"Failed to read progress from Redis: {redis_err}")
                    
                    # Sleep briefly to avoid busy polling
                    time.sleep(2)
            
            # Check final return code
            if process.returncode != 0:
                # Read error logs
                with open(subprocess_err, 'r') as f:
                    stderr_content = f.read()
                logger.error(f"❌ Subprocess failed with return code {process.returncode}")
                logger.error(f"❌ STDERR: {stderr_content}")
                raise Exception(f"Prediction subprocess failed: {stderr_content}")
            
            logger.info(f"✅ Subprocess completed successfully")
            
            # Log subprocess output for debugging
            if subprocess_log.exists():
                with open(subprocess_log, 'r') as f:
                    stdout_content = f.read()
                    if stdout_content:
                        logger.info(f"📊 Subprocess output logged to {subprocess_log}")
            
            # Read results from output file
            if not output_file.exists():
                raise Exception(f"Output file not created: {output_file}")
            
            with open(output_file, 'r') as f:
                result_data = json.load(f)
            
            logger.info(f"✅ Results loaded: {result_data.get('successful_predictions', 0)} successful predictions")
            
            # Handle both JSON Tables and legacy formats
            output_format = result_data.get('format', 'legacy')
            logger.info(f"📊 Output format: {output_format}")
            
            if output_format == 'json_tables' and 'results_table' in result_data:
                # Convert JSON Tables to legacy format for backward compatibility
                # But KEEP the results_table too for consistent API format
                try:
                    from jsontables import JSONTablesDecoder
                    
                    results_table = result_data['results_table']
                    records_with_predictions = JSONTablesDecoder.to_records(results_table)
                    
                    # Convert to legacy prediction format for backward compatibility
                    predictions = []
                    for i, record in enumerate(records_with_predictions):
                        # Extract prediction data from the enhanced record
                        prediction_obj = {
                            'row_index': record.get('row_index', i),
                            'prediction_id': None,  # Not stored in JSON Tables format
                            'error': record.get('prediction_error'),
                            'metadata': {}
                        }
                        
                        # Extract the actual prediction probabilities
                        if record.get('prediction_error') is None:
                            # Look for pred_* columns
                            prediction_probs = {}
                            for key, value in record.items():
                                if key.startswith('pred_') and key != 'pred_':
                                    class_name = key[5:]  # Remove 'pred_' prefix
                                    prediction_probs[class_name] = value
                            
                            if prediction_probs:
                                prediction_obj['prediction'] = prediction_probs
                            elif 'prediction' in record:
                                prediction_obj['prediction'] = record['prediction']
                            else:
                                prediction_obj['prediction'] = None
                                prediction_obj['error'] = "No prediction data found"
                            
                            # Add guardrails if available
                            if 'guardrails' in record:
                                prediction_obj['metadata']['guardrails'] = record['guardrails']
                        else:
                            prediction_obj['prediction'] = None
                        
                        predictions.append(prediction_obj)
                    
                    # Add legacy format alongside JSON Tables format
                    result_data['predictions'] = predictions
                    # Keep results_table as-is for consistent API format
                    logger.info(f"✅ Added legacy predictions format alongside JSON Tables: {len(predictions)} predictions")
                    
                except ImportError:
                    logger.warning(f"⚠️ JSON Tables not available for conversion, keeping raw format")
                except Exception as convert_error:
                    logger.warning(f"⚠️ Failed to convert JSON Tables format: {convert_error}")
                    # Keep the original format but log the issue
            
            elif output_format == 'legacy' and 'predictions' in result_data:
                logger.info(f"✅ Using legacy prediction format")
            else:
                logger.warning(f"⚠️ Unknown output format: {output_format}")
            
            # DON'T call update_state(state='SUCCESS') - Celery does this automatically
            # when the task returns. Calling it manually overwrites the return value!
            # Just update progress one final time with PROGRESS state
            self.update_state(
                state='PROGRESS',  # Use PROGRESS, not SUCCESS!
                meta={
                    'current': len(records),
                    'total': len(records),
                    'status': f"Completed: {result_data.get('successful_predictions', 0)} successful, {result_data.get('failed_predictions', 0)} failed",
                    'output_format': output_format,
                    'percentage': 100
                }
            )
            
            # Log what we're returning for debugging
            logger.info(f"📦 Returning result_data keys: {list(result_data.keys())}")
            logger.info(f"📦 results_table type: {type(result_data.get('results_table'))}")
            logger.info(f"📦 predictions count: {len(result_data.get('predictions', []))}")
            
            # Return the full result - Celery will automatically set state to SUCCESS
            return result_data
            
        except subprocess.TimeoutExpired:
            logger.error(f"❌ Subprocess timed out after 1 hour")
            raise Exception("Prediction subprocess timed out")
            
        except Exception as subprocess_error:
            logger.error(f"❌ Subprocess error: {subprocess_error}")
            logger.error(f"❌ Subprocess traceback: {traceback.format_exc()}")
            raise Exception(f"Prediction subprocess failed: {subprocess_error}")
            
        finally:
            # Cleanup temp files
            try:
                if input_file.exists():
                    input_file.unlink()
                if output_file.exists():
                    output_file.unlink()
                temp_dir.rmdir()
                logger.info(f"✅ Cleaned up temp directory: {temp_dir}")
            except Exception as cleanup_err:
                logger.warning(f"⚠️ Failed to cleanup temp files: {cleanup_err}")

        
    except Exception as e:
        # Log the full error
        error_msg = str(e)
        error_traceback = traceback.format_exc()
        
        # Log to worker logs for debugging
        logger.error(f"CELERY WORKER ERROR: {error_msg}")
        logger.error(f"CELERY WORKER TRACEBACK:\n{error_traceback}")
        
        # Return error state instead of re-raising
        return {
            'success': False,
            'error': error_msg,
            'traceback': error_traceback,
            'session_id': session_id,
            'total_records': len(records) if 'records' in locals() else 0,
            'exc_type': type(e).__name__,
            'exc_message': str(e)
        }
        
    finally:
        # Minimal cleanup since subprocess handles GPU/model cleanup
        try:
            import gc
            gc.collect()
            logger.info(f"✅ Python garbage collection completed")
        except Exception as cleanup_error:
            logger.warning(f"Cleanup failed: {cleanup_error}")

if __name__ == '__main__':
    app.start() 