import contextlib
import fcntl
import json
import logging
import math
import os
import pickle
import shutil
import signal
import socket
import stat
import sys
import time
import traceback
import ast
import requests
import sqlite3

import pandas as pd

# CRITICAL: Set up Python path BEFORE any featrix imports
# This allows inline imports of featrix.neural modules to work correctly
from pathlib import Path
lib_path = Path(__file__).parent / "lib"
if str(lib_path.resolve()) not in sys.path:
    sys.path.insert(0, str(lib_path.resolve()))

# from sessions import does_session_exist
from copy import deepcopy
from datetime import datetime, timedelta
from enum import Enum
from shutil import copyfile
from uuid import uuid4
from zoneinfo import ZoneInfo

# Job recovery imports
import glob
import re

from lib.utils import validate_user_metadata

from slack import get_hostname
from version import get_version
import boto3
from botocore.exceptions import ClientError
import hashlib

# from config import QUEUE_DIR, OUTPUT_DIR, DATA_DIR
from config import config
from slack import send_slack_message
from utils import convert_from_iso, convert_to_iso

from typing import Dict, Any, List

# Version tracking
try:
    from version import get_version
except ImportError:
    # Fallback if version module not available
    def get_version():
        class FakeVersion:
            def dict(self):
                return {
                    "semantic_version": "unknown",
                    "git_hash": None,
                    "git_date": None
                }
        return FakeVersion()

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

os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"] = "0"

import torch

def _log_gpu_memory_featrix_queue(context: str = ""):
    """Quick GPU memory logging for tracing memory usage in featrix_queue."""
    try:
        if not torch.cuda.is_available():
            return
        allocated = torch.cuda.memory_allocated() / (1024**3)  # GB
        reserved = torch.cuda.memory_reserved() / (1024**3)  # GB
        max_allocated = torch.cuda.max_memory_allocated() / (1024**3)  # GB
        logger.info(f"📊 GPU MEMORY [featrix_queue: {context}]: Allocated={allocated:.3f} GB, Reserved={reserved:.3f} GB, Peak={max_allocated:.3f} GB")
    except Exception as e:
        logger.debug(f"Could not log GPU memory: {e}")

class JobStatus(Enum):
    READY = "ready"
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"
    CANCELLED = "cancelled"
    PAUSED = "paused"


def clear_directory(dir_path: Path, absent_ok: bool = False, remove_if_present: bool = False):
    """Delete all sessions and queues."""

    if not dir_path.exists():
        if not absent_ok:
            raise Exception(f"The directory at {dir_path.resolve()} does not exist")
        else:
            # If the directory doesn't exist and we're OK with it, we're done.
            return

    for content_path in dir_path.iterdir():
        if content_path.is_dir():
            clear_directory(content_path, remove_if_present=True)
        else:
            content_path.unlink()

    if remove_if_present: 
        dir_path.rmdir()

    logger.info(f"Directory {dir_path.resolve()} cleared")


def ensure_directories_exist():
    """Ensure all necessary featrix directories exist."""
    directories_to_create = [
        config.queue_dir,
        config.output_dir, 
        config.data_dir,
        config.session_dir,
        config.session_private_dir
    ]
    
    for directory in directories_to_create:
        if not directory.exists():
            logger.info(f"Creating directory: {directory}")
            directory.mkdir(parents=True, exist_ok=True)


def initialize_standard_queues():
    """Initialize all standard queue directories that workers expect."""
    ensure_directories_exist()
    
    standard_queues = [
        "create_structured_data",
        "train_es", 
        "train_knn",
        "build_es_projections",
        "train_single_predictor"
    ]
    
    for queue_name in standard_queues:
        create_queue(queue_name)


def create_queue(queue_name: str):
    """
    This is idempotent - if the queue already exists, it will not be created again, or overwritten.
    """
    
    # Ensure parent directories exist first
    ensure_directories_exist()

    queue_path = config.queue_dir / queue_name

    # If the queue exsits, just log this and return. 
    # Creating a queue that already exists is just a noop.
    if queue_path.exists():
        logger.debug(f"Queue {queue_name} already exists")
        return

    queue_path.mkdir(exist_ok=True, parents=True)

    logger.info(f"New queue created at {queue_path}")


def iterate_over_jobs_in_queue(queue_name: str):
    queue_path = config.queue_dir / queue_name
    
    # Automatically create the queue directory if it doesn't exist
    if not queue_path.exists():
        logger.info(f"Queue {queue_name} does not exist, creating it automatically...")
        create_queue(queue_name)
        # Return early since a newly created queue will be empty
        return

    for job_path in queue_path.iterdir():
        # Skip non-job files.
        if job_path.suffix != ".job":
            continue

        # Just the name, without the extension
        job_id = job_path.stem

        yield load_job(queue_name, job_id)


def print_job(job: dict):
    serialized_job = serialize_job(job)

    logger.info(f"Job ID: {job.get('job_id')}")
    logger.info(json.dumps(serialized_job, indent=4))


def list_jobs_in_queue(queue_name: str):

    jobs = list(iterate_over_jobs_in_queue(queue_name))

    # Sort jobs from oldest to newest
    jobs = sorted(jobs, key=lambda x: x.get("created_at"), reverse=False)

    for job in jobs:
        print_job(job)
        logger.info("")  # Empty line for readability


def get_ready_jobs(queue_name: str, order: str = "oldest_first"):
    
    if order not in ['oldest_first', 'newest_first']:
        raise ValueError(f"Unsupported order: {order}")

    ready_jobs = []

    for job in iterate_over_jobs_in_queue(queue_name):
        if job.get("status") == JobStatus.READY:
            ready_jobs.append(job)

    # reverse=False means sorting in ascending order, which means the oldest ready job will be first.
    # Handle None values by treating them as epoch time (very old)
    epoch = datetime(1970, 1, 1, tzinfo=ZoneInfo("America/New_York"))
    reverse = order == 'newest_first'
    ready_jobs = sorted(ready_jobs, key=lambda x: x.get("created_at") or epoch, reverse=reverse)

    return ready_jobs


def get_version_info() -> dict:
    """Get version information for tracking in jobs and recovery info."""
    try:
        version_info = get_version()
        version_dict = version_info.dict()
        return {
            "version": version_dict.get("semantic_version", "unknown"),
            "git_hash": version_dict.get("git_hash"),
            "git_date": version_dict.get("git_date"),
        }
    except Exception as e:
        logger.warning(f"Failed to get version info: {e}")
        return {
            "version": "unknown",
            "git_hash": None,
            "git_date": None,
        }


def get_job_position_in_queue(queue_name: str, job_id: str):
    job = load_job(queue_name, job_id)
    job_id = job.get("job_id")
    job_status = job.get("status")

    # If the job has already been completed, is currently running, or has failed,
    # we return None to indicate that the job is not in the queue.
    if job_status != JobStatus.READY:
        return None
    
    # Get the list of ready jobs in the queue, starting with the oldest,
    # i.e. the job at position 0 is the next one to be worked on.
    ready_jobs = get_ready_jobs(queue_name, order="oldest_first")
    job_ids = [job.get("job_id") for job in ready_jobs]
    
    try:
        position = job_ids.index(job_id)
        return position
    except ValueError:
        # If not in the queue at all, something's gone wrong - if the job is not ready,
        # we should never get here
        raise ValueError(f"Something went wrong - job {job_id} status is {job_status}, but it's not in the queue")


def add_job_to_queue(queue_name: str, session_id: str, job_type: str, job_spec: dict, job_id: str = None):

    # Ensure all necessary directories exist before adding jobs
    ensure_directories_exist()

    # Defensive check: ensure session_id is provided
    if not session_id:
        logger.error(f"❌ add_job_to_queue called with empty session_id for job_type={job_type}, queue_name={queue_name}")
        raise ValueError(f"session_id is required but was empty or None for job_type={job_type}")

    # Session ID validation removed - we trust the session ID generation logic
    # No need to validate format as it causes more problems than it solves

    created_at = datetime.now(tz=ZoneInfo("America/New_York"))

    queue_path = config.queue_dir / queue_name
    
    # Automatically create the queue directory if it doesn't exist
    if not queue_path.exists():
        logger.info(f"Queue {queue_name} does not exist, creating it automatically...")
        create_queue(queue_name)

    # If job_id is not provided, generate a new one.
    # New format: {uuid}-{date}-{time} (job_type goes in path, not ID)
    # This makes job IDs shorter and paths more organized
    if job_id is None:
        unique_suffix = str(uuid4())[:6]
        job_timestamp = created_at.strftime('%Y%m%d-%H%M%S')
        job_id = f"{unique_suffix}-{job_timestamp}"

    session_exists = does_session_exist(session_id)
    if not session_exists:
        raise FileNotFoundError(f"Session {session_id} does not exist")

    # Get version info for this job
    version_info = get_version_info()

    job_content = dict(
        created_at=created_at,
        started_at=None,
        finished_at=None,
        session_id=session_id,
        type=job_type,
        status=JobStatus.READY,
        job_id=job_id,
        spec=job_spec,
        progress=None,
        version=version_info,
    )

    logger.debug(f"🔍 add_job_to_queue: Creating job {job_id} with session_id={session_id}, job_type={job_type}, queue_name={queue_name}")
    save_job(queue_name, job_id, job_content, exist_ok=False)
    logger.debug(f"✅ add_job_to_queue: Job {job_id} saved with session_id={session_id}")
    
    logger.info(f"Job {job_id} added to queue {queue_name}")

    return job_id


def get_job_output_path(job_id: str, session_id: str = None, job_type: str = None) -> Path:
    """
    Get the output directory path for a job.
    New structure: {output_dir}/{session_id}/{job_type}/{job_id}/
    
    If session_id or job_type are not provided, attempts to load from job file.
    Falls back to old structure for backward compatibility if job file doesn't exist.
    """
    # If we have all required info, use new structure
    if session_id and job_type:
        return config.output_dir / session_id / job_type / job_id
    
    # Try to load job to get session_id and job_type
    try:
        # Try to find job in any queue (job_id format might help us find it)
        # First, try to parse job_id to determine queue
        if '_' in job_id:
            # Old format: {job_type}_{timestamp}_{uuid}
            parts = job_id.split('_')
            if len(parts) >= 3:
                queue_name = parts[0]
                try:
                    job = load_job(queue_name, job_id)
                    session_id = job.get("session_id")
                    job_type = job.get("type")
                    if session_id and job_type:
                        return config.output_dir / session_id / job_type / job_id
                except (FileNotFoundError, KeyError):
                    pass
        else:
            # New format: {uuid}-{date}-{time}
            # Need to search all queues - this is expensive but rare
            for queue_name in ["train_es", "train_knn", "build_es_projections", "train_single_predictor", "cpu_data_tasks"]:
                try:
                    job = load_job(queue_name, job_id)
                    session_id = job.get("session_id")
                    job_type = job.get("type")
                    if session_id and job_type:
                        return config.output_dir / session_id / job_type / job_id
                except (FileNotFoundError, KeyError):
                    continue
    except Exception:
        pass
    
    # Fallback to old structure for backward compatibility
    return config.output_dir / job_id


def load_job(queue_name: str, job_id: str):
    queue_path = config.queue_dir / queue_name
    job_path = queue_path / f"{job_id}.job"
    lock_path = queue_path / f"{job_id}.lock"

    if not job_path.exists():
        raise FileNotFoundError(f"Job {job_id} does not exist in queue {queue_name}")

    # Use file locking to prevent reading while another process is writing
    lock_file = open(lock_path, 'w')
    fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH)  # Shared lock for reading
    
    try:
        job_dict = json.loads(job_path.read_text())

        job_dict['created_at'] = convert_from_iso(job_dict['created_at'])
        job_dict['started_at'] = convert_from_iso(job_dict['started_at'])
        job_dict['finished_at'] = convert_from_iso(job_dict['finished_at'])

        # Handle job status - convert to enum, with fallback for invalid values
        status_str = job_dict.get('status', 'ready')
        if isinstance(status_str, JobStatus):
            # Already an enum
            job_dict['status'] = status_str
        else:
            # Convert string to enum, handling case-insensitive and invalid values
            status_str_lower = str(status_str).lower()
            try:
                # Try exact match first
                job_dict['status'] = JobStatus(status_str)
            except ValueError:
                try:
                    # Try case-insensitive match
                    job_dict['status'] = JobStatus(status_str_lower)
                except ValueError:
                    # Invalid status - log warning and use READY as fallback
                    import logging
                    logger = logging.getLogger(__name__)
                    logger.warning(f"⚠️  Invalid job status '{status_str}' in job {job_id}, defaulting to READY")
                    job_dict['status'] = JobStatus.READY

        return job_dict
    finally:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)  # Release lock
        lock_file.close()


def serialize_job(job: dict):

    job = deepcopy(job)

    # Convert timestamp to string (handle missing fields for placeholder/deleted jobs)
    if "created_at" in job:
        job["created_at"] = convert_to_iso(job["created_at"])
    if "started_at" in job:
        job["started_at"] = convert_to_iso(job["started_at"])
    if "finished_at" in job:
        job["finished_at"] = convert_to_iso(job["finished_at"])
    
    # Convert status enum to string (handle string status for placeholder jobs)
    if "status" in job:
        if hasattr(job["status"], "value"):
            job["status"] = job["status"].value
        # else: already a string (e.g., "deleted")

    # Convert any numpy types to native Python types for JSON serialization
    serialized_job = clean_numpy_values(job)
    return serialized_job


def save_job(queue_name: str, job_id: str, job_content: dict, exist_ok=True):
    
    queue_path = config.queue_dir / queue_name
    final_job_path = queue_path / f"{job_id}.job"
    lock_path = queue_path / f"{job_id}.lock"
    
    if final_job_path.exists() and not exist_ok:
        raise FileExistsError(f"Job {job_id} already exists in queue {queue_name}")
    
    serialized_job = serialize_job(job_content)

    # Use file locking to prevent concurrent writes
    lock_file = open(lock_path, 'w')
    fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)  # Exclusive lock
    
    try:
        # We're assuming job descriptions can be long, and we don't want to kick
        # off a job before it is fully written to disk. Renames are atomic, so we
        # first create a temp file, and then rename it to the final name.
        temp_job_filename = job_id + ".tmp"
        temp_job_path = queue_path / temp_job_filename

        temp_job_path.write_text(json.dumps(serialized_job, indent=4))
        
        # CRITICAL FIX: Ensure the file is writable so status updates can modify it
        # This fixes the issue where job status updates fail silently due to read-only permissions
        temp_job_path.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)  # 644 permissions
        
        temp_job_path.rename(final_job_path)
        
        # Also ensure the final file has correct permissions after rename
        final_job_path.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)  # 644 permissions

        logger.debug(f"Job {job_id} written to queue {queue_name}")
    except Exception as e:
        # Clean up temp file if anything goes wrong
        temp_job_path = queue_path / f"{job_id}.tmp"
        if temp_job_path.exists():
            try:
                temp_job_path.unlink()
            except:
                pass
        raise
    finally:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)  # Release lock
        lock_file.close()


def update_job_status(queue_name: str, job_id: str, status: JobStatus):
    
    queue_path = config.queue_dir / queue_name
    job_path = queue_path / f"{job_id}.job"

    now = datetime.now(tz=ZoneInfo("America/New_York"))

    if not job_path.exists():
        raise FileNotFoundError(f"Job {job_id} does not exist in queue {queue_name}")
    
    # Load job with file locking to prevent race conditions
    # The lock is released after load_job() returns, but we'll acquire it again in save_job()
    job_content = load_job(queue_name, job_id)
    job_content["status"] = status

    # Update relevant timestamps
    if status == JobStatus.RUNNING:
        job_content["started_at"] = now
    elif status in [JobStatus.DONE, JobStatus.FAILED]:
        job_content["finished_at"] = now

    try:
        save_job(queue_name, job_id, job_content, exist_ok=True)
        logger.info(f"Job {job_id} updated with status {status.value}")
    except PermissionError as e:
        logger.error(f"PERMISSION ERROR updating job {job_id}: {e}")
        logger.error(f"Job file path: {job_path}")
        logger.error(f"Job file permissions: {oct(job_path.stat().st_mode)}")
        
        # Try to fix permissions and retry
        try:
            job_path.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)  # 644 permissions
            logger.info(f"Fixed permissions for job file {job_path}")
            
            # Retry the save
            save_job(queue_name, job_id, job_content, exist_ok=True)
            logger.info(f"Job {job_id} updated with status {status.value} (after permission fix)")
        except Exception as retry_error:
            logger.error(f"Failed to fix permissions and retry for job {job_id}: {retry_error}")
            raise e  # Re-raise the original permission error
    except Exception as e:
        logger.error(f"ERROR updating job {job_id} status: {e}")
        logger.error(f"Job file path: {job_path}")
        logger.error(f"Job file exists: {job_path.exists()}")
        if job_path.exists():
            logger.error(f"Job file permissions: {oct(job_path.stat().st_mode)}")
        raise

    # update_session_status
    # actually, better not to put it here, so they're not coupled


def requeue_job(queue_name: str, job_id: str):
    now = datetime.now(tz=ZoneInfo("America/New_York"))
    
    job = load_job(queue_name, job_id)

    job_status = job.get("status")
    if job_status not in [JobStatus.FAILED, JobStatus.DONE]:
        raise ValueError(f"Job {job_id} is not in a terminal state. Current state: {job_status}")

    job_spec = job.get("spec")
    if job_spec is None:
        raise ValueError(f"Job {job_id} does not have a spec")
    
    job_type = job.get("type")

    # Jobs are immutable, so requeuing a job means creating a new job with the same spec.
    # We'll add a timestamp to the job name to avoid collisions.
    new_job_id = job.get("job_id") + "_requeue_" + now.strftime('%Y%m%d-%H%M%S')

    new_job_id = add_job_to_queue(
        queue_name=queue_name, 
        job_type=job_type, 
        job_spec=job_spec,
        job_id=new_job_id,
        session_id=job.get('session_id')
    )
    return new_job_id


def run_echo_job(job_spec: dict):
    logger.info("echo job")
    time.sleep(5)
    logger.info("echo job done")


def run_generate_movie_frame(job_spec: dict, job_id: str):
    """
    Generate a movie frame on CPU (non-blocking, queued task).
    
    job_spec:
        session_id: str
        epoch: int
        checkpoint_path: str - path to model checkpoint
        data_snapshot_path: str - path to sampled data JSON
        output_dir: str - where to save projection
    """
    from lib.featrix.neural.movie_frame_task import generate_movie_frame_on_cpu
    
    session_id = job_spec['session_id']
    epoch = job_spec['epoch']
    checkpoint_path = job_spec['checkpoint_path']
    data_snapshot_path = job_spec['data_snapshot_path']
    output_dir = job_spec['output_dir']
    
    logger.info(f"🎬 [cpu_data_tasks] Generating movie frame for session {session_id}, epoch {epoch}")
    
    try:
        output_file = generate_movie_frame_on_cpu(
            checkpoint_path=checkpoint_path,
            data_snapshot_path=data_snapshot_path,
            epoch=epoch,
            output_dir=output_dir,
            session_id=session_id
        )
        
        if output_file:
            logger.info(f"✅ Movie frame generated: {output_file}")
        else:
            logger.error(f"❌ Movie frame generation returned None")
            
    except Exception as e:
        logger.error(f"❌ Movie frame generation failed: {e}")
        traceback.print_exc()
        raise


def run_create_structured_data(job_spec: dict, job_id: str, data_file: Path, session_id: str = None):
    from lib.structureddata import CSVtoDB
    
    # CRITICAL: Force CPU for sentence transformer in create_structured_data to avoid GPU OOM
    # This job runs on CPU queue and should not use GPU
    import os
    os.environ['FEATRIX_FORCE_CPU_SENTENCE_MODEL'] = '1'
    logger.info("🔧 Forcing CPU for sentence transformer in create_structured_data job")
    
    # CRITICAL: Store job_id in a local variable immediately to prevent shadowing
    parent_job_id = job_id
    if not parent_job_id:
        logger.error(f"❌ CRITICAL: run_create_structured_data called with job_id=None!")
        logger.error(f"   This will cause JSON ES training to fail!")
        raise ValueError("job_id is required for run_create_structured_data")
    
    logger.info(f"🚀 Starting structured data creation for job {parent_job_id}")
    logger.info(f"   Session ID: {session_id}")
    logger.info(f"   Data file (raw): {data_file}")
    if data_file:
        data_file_path = Path(data_file)
        if data_file_path.exists():
            abs_path = data_file_path.resolve()
            logger.info(f"   Data file (absolute path): {abs_path}")
            logger.info(f"   Data file exists: ✅")
            file_size = data_file_path.stat().st_size
            logger.info(f"   Data file size: {file_size / 1024 / 1024:.2f} MB ({file_size:,} bytes)")
        else:
            logger.warning(f"   Data file (absolute path): {data_file_path.resolve()}")
            logger.warning(f"   Data file exists: ❌")
    else:
        logger.warning(f"   Data file: None (will be handled by job spec)")
    
    # Extract column overrides from job spec
    column_overrides = job_spec.get("column_overrides", {})
    if column_overrides:
        logger.info(f"🔧 Column overrides specified: {len(column_overrides)} columns")
        logger.debug(f"   Override details: {column_overrides}")
    else:
        logger.info("ℹ️  No column overrides specified")
    
    # Check if this is an embedding space job with S3 datasets
    s3_training_path = job_spec.get("s3_training_dataset")
    s3_validation_path = job_spec.get("s3_validation_dataset")
    name = job_spec.get("name", "embedding_space")
    
    if s3_training_path and s3_validation_path:
        logger.info(f"Downloading S3 data for embedding space: {name}")
        
        import boto3
        from urllib.parse import urlparse
        
        # Check if paths are local files (not S3)
        if not s3_training_path.startswith('s3://') or not s3_validation_path.startswith('s3://'):
            # Local files - use them directly
            training_local_path = Path(s3_training_path)
            validation_local_path = Path(s3_validation_path)
            
            if not training_local_path.exists():
                raise ValueError(f"Training file not found: {training_local_path}")
            if not validation_local_path.exists():
                raise ValueError(f"Validation file not found: {validation_local_path}")
            
            logger.info(f"Using local training file: {training_local_path}")
            logger.info(f"Using local validation file: {validation_local_path}")
            
            # Use the training data as the data_file for processing
            data_file = training_local_path
        else:
            # S3 paths - download them
            # Parse S3 URLs
            training_parsed = urlparse(s3_training_path)
            validation_parsed = urlparse(s3_validation_path)
            
            if training_parsed.scheme != 's3' or validation_parsed.scheme != 's3':
                raise ValueError("S3 URLs must start with 's3://'")
            
            training_bucket = training_parsed.netloc
            training_key = training_parsed.path.lstrip('/')
            validation_bucket = validation_parsed.netloc
            validation_key = validation_parsed.path.lstrip('/')
            
            # Set up local file paths
            training_local_path = Path(f"{name}_training_data.csv")
            validation_local_path = Path(f"{name}_validation_data.csv")
            
            try:
                # Initialize S3 client
                s3_client = boto3.client('s3')
                
                # Download training dataset
                logger.info(f"Downloading training dataset from s3://{training_bucket}/{training_key}")
                s3_client.download_file(training_bucket, training_key, str(training_local_path))
                logger.info(f"Training dataset downloaded to {training_local_path}")
                
                # Download validation dataset
                logger.info(f"Downloading validation dataset from s3://{validation_bucket}/{validation_key}")
                s3_client.download_file(validation_bucket, validation_key, str(validation_local_path))
                logger.info(f"Validation dataset downloaded to {validation_local_path}")
                
                # Use the downloaded training data as the data_file for processing
                data_file = training_local_path
                
            except Exception as e:
                logger.error(f"Error downloading S3 data: {e}")
                raise
            
    # Check if data_file is an S3 URL and download it first
    elif isinstance(data_file, str) and data_file.startswith('s3://'):
        logger.info(f"Downloading S3 data: {data_file}")
        
        import boto3
        from urllib.parse import urlparse
        
        # Parse S3 URL
        parsed = urlparse(data_file)
        if parsed.scheme != 's3':
            raise ValueError("S3 URL must start with 's3://'")
        
        bucket = parsed.netloc
        key = parsed.path.lstrip('/')
        
        # Set up local file path - preserve original extension
        local_filename = Path(key).name
        # Keep original extension if it's a supported format, otherwise default to .csv
        if not local_filename.lower().endswith(('.csv', '.parquet', '.json', '.jsonl')):
            local_filename += '.csv'
        data_file = Path(local_filename)
        
        try:
            # Initialize S3 client and download
            s3_client = boto3.client('s3')
            logger.info(f"Downloading from s3://{bucket}/{key} to {data_file}")
            s3_client.download_file(bucket, key, str(data_file))
            logger.info(f"S3 download completed: {data_file}")
            
        except Exception as e:
            logger.error(f"Error downloading S3 data: {e}")
            raise
    
    # Proceed with normal CSV to DB processing
    # Extract JSON analysis configuration from job spec (if provided)
    json_analysis_config = job_spec.get("json_analysis_config", {})
    
    logger.info(f"📦 Creating CSVtoDB instance...")
    logger.info(f"   Column overrides: {len(column_overrides)} columns")
    logger.info(f"   JSON analysis config: {json_analysis_config}")
    start_time = time.time()
    
    c2d = CSVtoDB(
        column_overrides=column_overrides,
        json_analysis_config=json_analysis_config
    )
    logger.info(f"✅ CSVtoDB instance created")
    
    logger.info(f"🔄 Starting CSV to SQLite conversion...")
    logger.info(f"   Input file: {data_file}")
    logger.info(f"   File size: {Path(data_file).stat().st_size / 1024 / 1024:.2f} MB" if Path(data_file).exists() else "   File not found!")
    
    csv_start_time = time.time()
    c2d.csv_to_sqlite(csv_path=data_file)
    csv_elapsed = time.time() - csv_start_time
    logger.info(f"✅ CSV to SQLite conversion completed in {csv_elapsed:.1f} seconds")
    logger.info(f"   SQLite DB path: {c2d.db_path}")
    if Path(c2d.db_path).exists():
        db_size = Path(c2d.db_path).stat().st_size
        logger.info(f"   SQLite DB size: {db_size / 1024 / 1024:.2f} MB")
    
    # SENTENCE TRANSFORMER EMBEDDING GENERATION
    # Always generate embeddings and store in SQLite for PCA initialization
    # Optionally also export to JSONL if requested
    generate_embeddings = job_spec.get("generate_embeddings", False)
    embeddings_output_path = None
    
    # DISABLED: PCA embedding generation - moved to structureddata.py if needed
    # This was causing sentence transformer to load unnecessarily in featrix_queue
    logger.info(f"ℹ️  PCA embedding generation disabled - will be handled in structureddata.py if needed")
    embeddings_output_path = None
    
    # SCHEMA DUMPING AND METADATA COLLECTION
    logger.info(f"📋 Generating final schema and transformation metadata for job {parent_job_id}...")
    schema_start_time = time.time()
    
    try:
        # Get the input data file to inspect final schema
        from featrix.neural.input_data_file import FeatrixInputDataFile
        from featrix.neural.input_data_set import FeatrixInputDataSet
        
        input_data_file = FeatrixInputDataFile(c2d.db_path)
        df = input_data_file.df
        
        # Create a dataset to get column detection and encoder info
        dataset = FeatrixInputDataSet(
            df=df,
            ignore_cols=[],
            limit_rows=None,
            encoder_overrides=c2d.column_overrides,  # NEW: Use updated overrides (includes auto-detected)
            project_row_meta_data=None,
        )
        
        # Build comprehensive schema metadata
        schema_metadata = {
            "job_id": parent_job_id,
            "original_file": str(data_file),
            "sqlite_db": str(c2d.db_path),
            "total_rows": len(df),
            "total_columns": 0,  # Will be updated after filtering internal columns
            "creation_timestamp": datetime.now().isoformat(),
            
            # Column information
            "columns": {},
            
            # Transformation metadata
            "json_transformations": getattr(c2d, 'json_transformations', {}),
            "added_columns_stats": getattr(c2d, 'added_columns_stats', []),
            "json_parseable_columns": getattr(c2d, 'json_parseable_columns', []),  # Columns with JSON-parseable dicts or lists of dicts
            "column_overrides": c2d.column_overrides,  # NEW: Use updated overrides (includes auto-detected)
            
            # Encoder metadata for .encode() compatibility
            "encoder_metadata": {},
            
            # Original vs final schema comparison
            "schema_changes": {
                "original_columns": [],  # Will be populated if we can read original
                "final_columns": list(df.columns),
                "columns_added": len(getattr(c2d, 'added_columns_stats', [])),
                "json_columns_expanded": len(getattr(c2d, 'json_transformations', {}))
            },
            
            # Child ES sessions (for JSON columns with separate schemas)
            # Will be populated after child ES sessions are created
            "child_es_sessions": {}  # {col_name: child_session_id} - tracks which columns have separate ES schemas
        }
        
        # Analyze each column and build encoder metadata
        logger.info(f"🔍 Analyzing {len(df.columns)} columns for encoder metadata...")
        
        for col_name in df.columns:
            # Wrap entire column analysis in try-except to prevent one bad column from failing the whole schema generation
            try:
                col_data = df[col_name]
                
                # Try to get sample values, handling non-UTF-8 data gracefully
                sample_values = []
                try:
                    sample_data = col_data.dropna().head(5)
                    for val in sample_data:
                        try:
                            # Try to convert to string
                            str_val = str(val)
                            # Check if it's valid UTF-8
                            str_val.encode('utf-8')
                            sample_values.append(str_val)
                        except (UnicodeDecodeError, UnicodeEncodeError):
                            # If conversion fails, use a placeholder
                            sample_values.append(f"<non-UTF-8 data: {type(val).__name__}>")
                except Exception as e:
                    logger.debug(f"⚠️  Could not extract sample values for column '{col_name}': {e}")
                    sample_values = ["<unable to convert>"]
                
                # Basic column info
                col_info = {
                    "dtype": str(col_data.dtype),
                    "null_count": col_data.isnull().sum(),
                    "unique_count": col_data.nunique(),
                    "sample_values": sample_values,
                }
                
                # Detect encoder type using dataset's logic
                try:
                    # Use the dataset's column detection logic
                    col_codecs = dataset.column_codecs()
                    if col_name in col_codecs:
                        codec = col_codecs[col_name]
                        codec_type = type(codec).__name__
                        
                        col_info["encoder_type"] = codec_type
                        col_info["encoder_class"] = codec_type
                        
                        # Store encoder-specific metadata
                        if hasattr(codec, 'string_cache') and codec.string_cache:
                            col_info["string_cache_size"] = len(codec.string_cache.cache_dict)
                            col_info["vocabulary"] = list(codec.string_cache.cache_dict.keys())[:100]  # First 100 for preview
                        
                        if hasattr(codec, 'numeric_stats'):
                            col_info["numeric_stats"] = codec.numeric_stats
                            
                        if hasattr(codec, 'set_delimiter'):
                            col_info["set_delimiter"] = codec.set_delimiter
                            
                    else:
                        col_info["encoder_type"] = "unknown"
                        col_info["encoder_class"] = "unknown"
                        
                except Exception as e:
                    logger.warning(f"Could not determine encoder type for column {col_name}: {e}")
                    col_info["encoder_type"] = "error"
                    col_info["encoder_error"] = str(e)
                
                # Add column info to metadata
                schema_metadata["columns"][col_name] = col_info
                
                # Store encoder metadata for .encode() usage
                schema_metadata["encoder_metadata"][col_name] = {
                    "encoder_type": col_info.get("encoder_type"),
                    "encoder_class": col_info.get("encoder_class"),
                    "column_override": c2d.column_overrides.get(col_name),  # NEW: Use updated overrides
                    "transformation_source": None  # Will be set below for JSON-derived columns
                }
                    
            except Exception as col_error:
                # Catch any errors during column analysis and continue with other columns
                logger.warning(f"⚠️  Error analyzing column '{col_name}': {col_error}")
                logger.debug(f"   Traceback: {traceback.format_exc()}")
                # Add minimal column info so schema generation can continue
                try:
                    schema_metadata["columns"][col_name] = {
                        "dtype": str(df[col_name].dtype),
                        "error": str(col_error),
                        "encoder_type": "error"
                    }
                    schema_metadata["encoder_metadata"][col_name] = {
                        "encoder_type": "error",
                        "error": str(col_error)
                    }
                except Exception:
                    # If even minimal info fails, skip this column entirely
                    logger.error(f"❌ Could not add any info for column '{col_name}', skipping")
        
        # Link JSON transformation metadata to encoder metadata
        # FIXME: we might need to call the trained ES's to get the encodings???
        # for original_col, transformation in getattr(c2d, 'json_transformations', {}).items():
        #     for expanded_col in transformation.get('expanded_columns', []):
        #         if expanded_col in schema_metadata["encoder_metadata"]:
        #             transformation_type = transformation.get("transformation_type")
                    
        #             # Handle list conversion columns specially
        #             if expanded_col.endswith('.list'):
        #                 json_key_path = expanded_col.replace(f"{original_col}.dict.", "").replace(".list", "")
        #                 transformation_source = {
        #                     "original_column": original_col,
        #                     "transformation_type": transformation_type,
        #                     "json_key_path": json_key_path,
        #                     "list_format": True,  # NEW: Flag indicating this is in key=value list format
        #                     "delimiter": "|"      # NEW: Delimiter used for list values
        #                 }
        #             else:
        #                 json_key_path = expanded_col.replace(f"{original_col}.dict.", "")
        #                 transformation_source = {
        #                     "original_column": original_col,
        #                     "transformation_type": transformation_type,
        #                     "json_key_path": json_key_path,
        #                     "list_format": False
        #                 }
                    
        #             schema_metadata["encoder_metadata"][expanded_col]["transformation_source"] = transformation_source
        
        # Filter out internal __featrix* columns from user-visible lists
        # These are internal metadata columns that should not appear in schema documentation
        from lib.featrix.neural.utils import filter_internal_columns
        
        # Filter the columns list in schema_changes
        final_columns_all = list(df.columns)
        final_columns_user_visible = filter_internal_columns(final_columns_all)
        
        # Update schema metadata to only show user-visible columns
        schema_metadata["schema_changes"]["final_columns"] = final_columns_user_visible
        schema_metadata["schema_changes"]["final_columns_including_internal"] = final_columns_all
        schema_metadata["total_columns"] = len(final_columns_user_visible)
        schema_metadata["total_columns_including_internal"] = len(final_columns_all)
        
        # Remove internal columns from the columns dict
        for internal_col in [col for col in df.columns if col.startswith('__featrix')]:
            if internal_col in schema_metadata["columns"]:
                del schema_metadata["columns"][internal_col]
            if internal_col in schema_metadata["encoder_metadata"]:
                del schema_metadata["encoder_metadata"][internal_col]
        
        logger.info(f"🔒 Filtered {len(final_columns_all) - len(final_columns_user_visible)} internal columns from schema metadata")
        
        # Try to get original column info if we can read the source file
        try:
            original_df = pd.read_csv(data_file, nrows=0)  # Just get column names
            original_columns_all = list(original_df.columns)
            original_columns_user_visible = filter_internal_columns(original_columns_all)
            schema_metadata["schema_changes"]["original_columns"] = original_columns_user_visible
            schema_metadata["schema_changes"]["original_columns_including_internal"] = original_columns_all
            logger.info(f"📋 Original file had {len(original_columns_user_visible)} columns, final has {len(final_columns_user_visible)} columns (user-visible)")
        except Exception as e:
            logger.warning(f"Could not read original file for schema comparison: {e}")
            schema_metadata["schema_changes"]["original_columns"] = ["could_not_read_original"]
        
        # Save schema metadata to file
        schema_file = Path("final_schema.json")
        with open(schema_file, 'w') as f:
            json.dump(schema_metadata, f, indent=2, default=str)
        
        logger.info(f"📋 Final schema saved to {schema_file.resolve()}")
        logger.info(f"✅ Schema metadata generated:")
        logger.info(f"   • Total columns: {schema_metadata['total_columns']}")
        logger.info(f"   • JSON transformations: {len(schema_metadata['json_transformations'])}")
        json_parseable = schema_metadata.get('json_parseable_columns', [])
        if json_parseable:
            logger.info(f"   • JSON-parseable columns: {len(json_parseable)}")
            for col_info in json_parseable:
                types = []
                if col_info.get('has_dicts'):
                    types.append('dicts')
                if col_info.get('has_list_of_dicts'):
                    types.append('list-of-dicts')
                logger.info(f"     - {col_info['column']}: {', '.join(types)}")
        logger.info(f"   • Column overrides: {len(schema_metadata['column_overrides'])} (includes auto-detected)")
        logger.info(f"   • Encoder types detected: {len(set(info.get('encoder_type', 'unknown') for info in schema_metadata['columns'].values()))}")
        
        # Queue ES training for JSON columns
        # TEMPORARILY DISABLED: Skip JSON columns to save time
        json_parseable = schema_metadata.get('json_parseable_columns', [])
        if json_parseable:
            logger.warning(f"⚠️  Skipping {len(json_parseable)} JSON columns - JSON columns are temporarily disabled: {json_parseable}")
            json_parseable = []  # Clear to prevent processing
        if json_parseable:
            logger.info(f"🚀 Queueing embedding space training jobs for {len(json_parseable)} JSON columns...")
            
            # Get current host for ES training targeting
            current_host = None
            try:
                hostname = socket.gethostname()
                short_name = hostname.split('.')[0].lower()
                if short_name in ['taco', 'burrito', 'churro']:
                    current_host = short_name
                    logger.info(f"🎯 Current host detected: {current_host}")
            except Exception as e:
                logger.warning(f"⚠️ Could not detect current host: {e}")
            
            # First pass: queue all ES training jobs and collect session IDs
            json_column_sessions = {}  # Map column_name -> session_id
            processed_columns = set()  # Track which columns we've already processed to prevent duplicates
            
            # CRITICAL: Ensure we have a session_id to pass as parent_session_id
            if not session_id:
                logger.error(f"❌ run_create_structured_data: session_id is None! Cannot create ES training jobs under parent session directory.")
                logger.error(f"   This will cause ES training jobs to be created in separate session directories.")
                # Try to get session_id from the job
                try:
                    job = load_job("cpu_data_tasks", parent_job_id)
                    session_id = job.get("session_id")
                    if session_id:
                        logger.info(f"✅ Retrieved session_id={session_id} from job file")
                    else:
                        logger.error(f"❌ Job file also missing session_id!")
                except Exception as e:
                    logger.error(f"❌ Failed to load job to get session_id: {e}")
            
            # CRITICAL: Skip JSON column ES training if we don't have a valid session_id
            # This prevents jobs from going to wrong directories
            if not session_id:
                logger.error(f"❌ SKIPPING JSON column ES training - session_id is still None after all attempts!")
                logger.error(f"   This job will complete without ES training for JSON columns")
                logger.error(f"   JSON columns will need to be handled manually")
            else:
                logger.info(f"✅ Have valid session_id={session_id} for JSON column ES training")
            
            for col_info in json_parseable:
                col_name = col_info.get('column')
                if not col_name or col_name not in df.columns:
                    continue
                
                # CRITICAL: Skip if we've already processed this column (prevent duplicates)
                if col_name in processed_columns:
                    logger.warning(f"⚠️ Column '{col_name}' already processed, skipping duplicate")
                    continue
                processed_columns.add(col_name)
                
                logger.info(f"📋 Queueing ES training job for JSON column '{col_name}'...")
                
                # Extract all JSON values from the column
                json_values = df[col_name].dropna().tolist()
                
                if len(json_values) < 10:
                    logger.warning(f"⚠️ Column '{col_name}' has too few values ({len(json_values)}) for ES training, skipping")
                    continue
                
                # CRITICAL: Only queue ES training if we have a valid session_id
                if not session_id:
                    logger.error(f"❌ SKIPPING ES training for column '{col_name}' - session_id is None!")
                    logger.error(f"   Cannot create ES training jobs without parent session_id")
                    continue
                
                # Queue ES training job for this column (don't wait)
                # Pass parent session_id so ES training jobs are created under the parent session's output directory
                logger.info(f"🔍 run_create_structured_data: Calling queue_es_training_for_json_column with parent_session_id={session_id} for column '{col_name}'")
                try:
                    es_training_session_id = queue_es_training_for_json_column(
                        col_name, json_values, current_host, 
                        json_column_sessions,  # Pass map of other JSON column sessions for dependency detection
                        parent_es_name=name,  # Pass parent ES name for naming child ES
                        parent_session_id=session_id  # Pass parent session ID for output directory
                    )
                except ValueError as e:
                    if "parent_session_id is required" in str(e):
                        logger.error(f"❌ FAILED to create ES training for column '{col_name}': {e}")
                        logger.error(f"   Skipping this column - ES training will not be available")
                        continue
                    else:
                        raise
                except Exception as e:
                    # CRITICAL: Catch all exceptions to prevent one failed column from stopping processing of others
                    logger.error(f"❌ FAILED to create ES training for column '{col_name}': {e}")
                    logger.error(f"   Error type: {type(e).__name__}")
                    logger.error(f"   Skipping this column - ES training will not be available")
                    logger.error(traceback.format_exc())
                    continue
                logger.debug(f"🔍 run_create_structured_data: queue_es_training_for_json_column returned es_training_session_id={es_training_session_id}")
                
                if es_training_session_id:
                    json_column_sessions[col_name] = es_training_session_id
                    # Get job IDs for logging
                    try:
                        es_session = load_session(es_training_session_id)
                        job_plan = es_session.get("job_plan", [])
                        job_ids = []
                        for job_desc in job_plan:
                            job_id = job_desc.get("job_id")
                            job_type = job_desc.get("job_type")
                            if job_id:
                                job_ids.append(f"{job_type}:{job_id}")
                        job_ids_str = f" (jobs: {', '.join(job_ids)})" if job_ids else ""
                    except Exception as e:
                        logger.debug(f"Could not load session for job IDs: {e}")
                        job_ids_str = ""
                    logger.info(f"✅ Queued ES training job for column '{col_name}' (session: {es_training_session_id}){job_ids_str}")
                else:
                    logger.warning(f"⚠️ Failed to queue ES training job for column '{col_name}', will continue without it")
            
            # Store child ES session references in schema_metadata for tracking multiple schemas
            if json_column_sessions:
                schema_metadata["child_es_sessions"] = json_column_sessions  # {col_name: child_session_id}
                logger.info(f"📋 Stored {len(json_column_sessions)} child ES session references in schema_metadata")
        
    except Exception as e:
        logger.error(f"❌ Error generating schema metadata: {e}")
        logger.error(f"Traceback: {traceback.format_exc()}")
        # Continue without schema metadata if there's an error
        schema_metadata = {
            "error": f"Failed to generate schema metadata: {str(e)}",
            "job_id": job_id,
            "creation_timestamp": datetime.now().isoformat()
        }
    
    schema_elapsed = time.time() - schema_start_time
    logger.info(f"✅ Schema metadata generation completed in {schema_elapsed:.1f} seconds")
    
    logger.info(f"🔒 Closing CSVtoDB connection...")
    c2d.close()
    logger.info(f"✅ CSVtoDB connection closed")
    
    logger.info(f"📁 Getting output files...")
    output_files = c2d.get_output_files()
    logger.info(f"✅ Output files retrieved:")
    logger.info(f"   SQLite DB: {output_files.get('sqlite_db', 'N/A')}")
    logger.info(f"   Strings cache: {output_files.get('strings_cache', 'N/A')}")

    total_elapsed = time.time() - start_time
    logger.info(f"🎉 Structured data creation job finished successfully!")
    logger.info(f"   Total time: {total_elapsed:.1f} seconds ({total_elapsed / 60:.1f} minutes)")
    logger.info(f"   SQLite DB: {output_files['sqlite_db']}")
    
    # Include comprehensive metadata in session output
    result = {
        "sqlite_db": str(Path(output_files['sqlite_db']).resolve()),
        "column_overrides": c2d.column_overrides,  # NEW: Use updated overrides (includes auto-detected)
        "schema_metadata": schema_metadata,  # NEW: Complete schema and transformation metadata
        "final_schema_file": str(Path("final_schema.json").resolve()),  # NEW: Path to schema file
    }
    
    # Only include strings_cache if the file actually exists (may not exist if there are no strings)
    strings_cache_path_str = output_files.get('strings_cache', '')
    if strings_cache_path_str:
        try:
            strings_cache_path = Path(strings_cache_path_str)
            # Try to resolve to absolute path, but handle case where working directory doesn't exist
            try:
                if strings_cache_path.is_absolute():
                    abs_path = strings_cache_path
                else:
                    # Relative path - try to resolve, but fallback to just using the path as-is if getcwd() fails
                    abs_path = strings_cache_path.resolve()
            except (FileNotFoundError, OSError):
                # Working directory doesn't exist - use the path as-is (it should already be absolute from get_output_files)
                abs_path = strings_cache_path
            
            if abs_path.exists():
                result["strings_cache"] = str(abs_path)
                logger.info(f"✅ Strings cache included: {result['strings_cache']}")
            else:
                logger.info(f"ℹ️  Strings cache path provided but file doesn't exist: {abs_path}")
                result["strings_cache"] = None
        except Exception as e:
            logger.warning(f"⚠️  Error processing strings_cache path '{strings_cache_path_str}': {e}")
            result["strings_cache"] = None
    else:
        logger.info(f"ℹ️  No strings cache file (no strings in data)")
        result["strings_cache"] = None
    
    # Add embeddings path if generated
    if embeddings_output_path:
        result["embeddings_file"] = str(Path(embeddings_output_path).resolve())
        logger.info(f"🔮 Embeddings file included in session: {result['embeddings_file']}")
    
    # NEW: Log the updated column overrides if any were auto-detected
    if len(c2d.column_overrides) > len(column_overrides):
        auto_detected_count = len(c2d.column_overrides) - len(column_overrides)
        logger.info(f"🔧 Added {auto_detected_count} auto-detected encoders to column overrides")
        logger.info(f"   Final column overrides: {len(c2d.column_overrides)} total")
    
    # Include JSON transformations if present
    json_transformations = output_files.get('json_transformations', {})
    if json_transformations:
        result["json_transformations"] = json_transformations
        logger.info(f"🔧 Stored JSON transformations for {len(json_transformations)} columns in session")
    
    return result


def run_train_es_job(job_spec: dict, job_id: str, data_file: Path, strings_cache: str):

    from lib.es_training import LightTrainingArgs, train_es

    logger.info(f"Starting embedding space training for job {job_id}")
    
    # Get session early to check for dependencies
    job = load_job("train_es", job_id)
    session_id = job.get("session_id")
    if session_id:
        logger.info(f"   Session ID: {session_id}")
    else:
        logger.error(f"   ❌ CRITICAL: train_es job {job_id} has NO session_id!")
    session = load_session(session_id)
    
    # Check if this is a parent session that needs to wait for child ES sessions to complete
    # Only check if this is NOT a child ES training session (i.e., no parent_session_id)
    parent_session_id = session.get("parent_session_id")
    if not parent_session_id:
        # This is a parent session - check for child ES sessions from create_structured_data
        schema_metadata = session.get("schema_metadata", {})
        child_es_sessions = schema_metadata.get("child_es_sessions", {})
        
        if child_es_sessions:
            logger.info(f"🔗 Parent session has {len(child_es_sessions)} child ES sessions to wait for before training...")
            logger.info(f"   Child sessions: {list(child_es_sessions.keys())}")
            
            # Wait for all child sessions' train_es jobs to complete
            max_wait = 3600 * 4  # 4 hours max
            check_interval = 30  # Check every 30 seconds
            start_time = time.time()
            
            while time.time() - start_time < max_wait:
                all_complete = True
                incomplete_columns = []
                
                for col_name, child_session_id in child_es_sessions.items():
                    try:
                        child_session = load_session(child_session_id)
                        job_plan = child_session.get("job_plan", [])
                        
                        # Find the train_es job in the child session's job plan
                        train_es_job_id = None
                        for job_desc in job_plan:
                            if job_desc.get("job_type") == "train_es":
                                train_es_job_id = job_desc.get("job_id")
                                break
                        
                        if not train_es_job_id:
                            # Job not queued yet or not found
                            all_complete = False
                            incomplete_columns.append(f"{col_name} (no train_es job)")
                            continue
                        
                        # Check if the train_es job is DONE
                        train_es_job = load_job("train_es", train_es_job_id)
                        job_status = train_es_job.get("status")
                        
                        if job_status != JobStatus.DONE:
                            all_complete = False
                            incomplete_columns.append(f"{col_name} (status: {job_status})")
                    except Exception as e:
                        logger.debug(f"Failed to check child session {child_session_id} for column '{col_name}': {e}")
                        all_complete = False
                        incomplete_columns.append(f"{col_name} (error)")
                
                if all_complete:
                    logger.info(f"✅ All {len(child_es_sessions)} child ES sessions' train_es jobs are complete!")
                    for col_name, child_session_id in child_es_sessions.items():
                        logger.info(f"   ✅ Column '{col_name}': child session {child_session_id} train_es complete")
                    break
                
                elapsed_minutes = (time.time() - start_time) / 60
                logger.info(f"⏳ Waiting for {len(incomplete_columns)} child ES sessions... ({elapsed_minutes:.1f} minutes elapsed)")
                logger.info(f"   Incomplete: {incomplete_columns}")
                time.sleep(check_interval)
            else:
                # Timeout
                logger.error(f"❌ Timeout waiting for child ES sessions to complete (waited {max_wait/60:.1f} minutes)")
                raise TimeoutError(f"Child ES sessions not complete: {incomplete_columns}")
    
    # Check for required ES dependencies before starting training
    # Dependencies are stored as {col_name: session_id} mapping
    required_child_es_mapping = session.get("required_child_es_mapping", {})
    if required_child_es_mapping:
        logger.info(f"🔗 Checking {len(required_child_es_mapping)} required ES dependencies before training...")
        logger.info(f"   Required columns: {list(required_child_es_mapping.keys())}")
        
        # Wait for all required ES's to be available locally
        max_wait = 3600 * 4  # 4 hours max
        check_interval = 30  # Check every 30 seconds
        start_time = time.time()
        
        while time.time() - start_time < max_wait:
            all_ready = True
            missing_columns = []
            
            for col_name, dep_session_id in required_child_es_mapping.items():
                try:
                    dep_session = load_session(dep_session_id)
                    dep_es_path = dep_session.get("embedding_space")
                    
                    if not dep_es_path or not Path(dep_es_path).exists():
                        all_ready = False
                        missing_columns.append(col_name)
                except Exception as e:
                    logger.debug(f"Failed to load session {dep_session_id} for column '{col_name}': {e}")
                    all_ready = False
                    missing_columns.append(col_name)
            
            if all_ready:
                logger.info(f"✅ All {len(required_child_es_mapping)} required ES dependencies are ready!")
                for col_name, dep_session_id in required_child_es_mapping.items():
                    dep_session = load_session(dep_session_id)
                    dep_es_path = dep_session.get("embedding_space")
                    logger.info(f"   ✅ Column '{col_name}': {dep_es_path}")
                break
            
            elapsed_minutes = (time.time() - start_time) / 60
            logger.info(f"⏳ Waiting for {len(missing_columns)} required ES's... ({elapsed_minutes:.1f} minutes elapsed)")
            logger.info(f"   Missing columns: {missing_columns}")
            time.sleep(check_interval)
        else:
            # Timeout
            logger.error(f"❌ Timeout waiting for required ES dependencies (waited {max_wait/60:.1f} minutes)")
            raise TimeoutError(f"Required ES dependencies not available for columns: {missing_columns}")
    
    args = LightTrainingArgs()
    args.epochs = job_spec.get("epochs")
    # Get row_limit from job_spec, fall back to config.json if not specified
    args.row_limit = job_spec.get("row_limit")
    if args.row_limit is None:
        from lib.sphere_config import get_row_limit
        args.row_limit = get_row_limit()
        logger.info(f"📊 Using row_limit from config.json: {args.row_limit}")
    args.is_production = job_spec.get("is_production")
    args.input_file = str(data_file.resolve())
    args.ignore_cols = job_spec.get("ignore_cols")
    args.learning_rate = job_spec.get("learning_rate")
    args.batch_size = job_spec.get("batch_size")
    args.job_id = job_id
    args.job_queue = "train_es"
    args.string_cache = strings_cache
    
    # Use the session's own session_id (child ES sessions are top-level)
    args.session_id = session_id
    
    # Get the name from the session if available
    name = session.get('name')
    if name:
        args.name = name
        logger.info(f"🏷️  Embedding space name: {name}")
    else:
        logger.info("ℹ️  No name specified for embedding space")
    
    # Get JSON transformations
    json_transformations = session.get('json_transformations', {})
    args.json_transformations = json_transformations
    
    if json_transformations:
        logger.info(f"🔧 Found JSON transformations for {len(json_transformations)} columns:")
        for orig_col, transform_info in json_transformations.items():
            logger.info(f"   • {orig_col} → {len(transform_info.get('expanded_columns', []))} expanded columns")
    else:
        logger.info("ℹ️  No JSON transformations found in session")
    
    # Get column overrides
    column_overrides = session.get('column_overrides', {})
    args.column_overrides = column_overrides
    
    # Get user_metadata from job_spec or session
    user_metadata = job_spec.get("user_metadata") or session.get('user_metadata')
    if user_metadata:
        # Validate size
        try:
            validated_metadata = validate_user_metadata(user_metadata, max_size_kb=32)
            # Set user_metadata if the field exists (handles old server versions)
            try:
                args.user_metadata = validated_metadata
                logger.info(f"✅ User metadata validated: {len(str(validated_metadata))} chars")
            except (ValueError, AttributeError) as e:
                # Field doesn't exist in this version of LightTrainingArgs - skip it
                logger.debug(f"⚠️  user_metadata field not available in LightTrainingArgs: {e}")
        except ValueError as e:
            logger.error(f"❌ Invalid user_metadata: {e}")
            raise
    
    if column_overrides:
        logger.info(f"🔧 Found column overrides for {len(column_overrides)} columns:")
        for col_name, override_type in column_overrides.items():
            logger.info(f"   • {col_name} → {override_type}")
    else:
        logger.info("ℹ️  No column overrides found in session")
    
    # DropoutScheduler parameters
    args.enable_dropout_scheduler = job_spec.get("enable_dropout_scheduler", True)
    args.dropout_schedule_type = job_spec.get("dropout_schedule_type", "linear_decay")
    args.initial_dropout = job_spec.get("initial_dropout", 0.5)
    args.final_dropout = job_spec.get("final_dropout", 0.1)
    
    # Movie frame and WeightWatcher intervals (default to 5 for better performance)
    args.movie_frame_interval = job_spec.get("movie_frame_interval", 3)  # Changed default from 5 to 3
    args.weightwatcher_save_every = job_spec.get("weightwatcher_save_every", 5)
    
    # Pass resume information if present (from job recovery or extension)
    # Note: For backward compatibility, we only pass resume info via job_spec
    # The training code checks both args.resume_from_epoch AND args.job_spec.get('resume_from_epoch')
    if hasattr(args, 'resume_from_epoch'):
        args.resume_from_epoch = job_spec.get("resume_from_epoch")
    if hasattr(args, 'job_spec'):
        args.job_spec = job_spec  # Pass entire spec for additional recovery info
    
    resume_epoch = job_spec.get("resume_from_epoch")
    resume_from_es_path = job_spec.get("resume_from_embedding_space")
    
    # Check if this is a fine-tuning job (has parent_embedding_space_path)
    parent_es_path_str = job_spec.get("parent_embedding_space_path")
    is_fine_tune = job_spec.get("fine_tune", False) and parent_es_path_str
    
    if is_fine_tune:
        # Resolve parent ES path, checking published location first
        # For parent ES, we need to get the session_id from the parent session
        # For now, try to resolve it directly
        parent_es_path = Path(parent_es_path_str)
        if not parent_es_path.exists():
            # Try to resolve if it's from a published session
            # Extract session_id from path if possible, or try common locations
            published_candidate = Path("/sphere/published") / parent_es_path.name.replace("embedding_space.pickle", "").replace("embedded_space.pickle", "").strip("/")
            if published_candidate.exists():
                # Try to find the ES in published location
                for possible_es in [published_candidate / "output" / "embedding_space.pickle",
                                   published_candidate / "output" / "embedded_space.pickle",
                                   published_candidate / "embedding_space.pickle"]:
                    if possible_es.exists():
                        parent_es_path = possible_es
                        break
        
        if not parent_es_path.exists():
            raise FileNotFoundError(f"Parent embedding space not found at: {parent_es_path} (original: {parent_es_path_str})")
        # Pass parent ES path to training args
        args.parent_embedding_space_path = str(parent_es_path.resolve())
        logger.info(f"   ✅ Parent embedding space path set for fine-tuning: {parent_es_path}")
    
    # If resuming and session already has an embedding space (from extension/clone), use it
    if resume_epoch is not None and session.get("embedding_space"):
        # Resolve path, checking published location first
        existing_es_path_str = session.get("embedding_space")
        existing_es_path = resolve_session_path(session_id, existing_es_path_str)
        
        if existing_es_path.exists():
            logger.info(f"🔄 Job {job_id} will resume training from epoch {resume_epoch}")
            logger.info(f"   Using existing embedding space: {existing_es_path}")
            # Copy the existing ES to working directory so training can use it
            from shutil import copyfile
            copyfile(existing_es_path, Path("embedded_space.pickle"))
            logger.info(f"   ✅ Copied existing ES to working directory for resume")
        else:
            logger.warning(f"⚠️  Resume requested but ES path doesn't exist: {existing_es_path}")
    # Also check resume_from_embedding_space if specified
    if resume_from_es_path:
        # Resolve path, checking published location first
        # Try to extract session_id from the path or use current session_id
        resolved_es_path = Path(resume_from_es_path)
        if not resolved_es_path.exists():
            # Try to resolve using current session_id (assuming it's from the same session)
            try:
                resolved_es_path = resolve_session_path(session_id, resume_from_es_path)
            except Exception as e:
                logger.warning(f"Could not resolve resume_from_embedding_space path: {e}")
        
        if resolved_es_path.exists():
            logger.info(f"🔄 Job {job_id} will resume from embedding space: {resolved_es_path}")
            # Copy to working directory
            from shutil import copyfile
            copyfile(resolved_es_path, Path("embedded_space.pickle"))
            logger.info(f"   ✅ Copied ES from {resolved_es_path} to working directory for resume")
        else:
            logger.warning(f"⚠️  resume_from_embedding_space specified but path doesn't exist: {resume_from_es_path}")
    
    elif resume_epoch is not None:
        logger.info(f"🔄 Job {job_id} will resume training from epoch {resume_epoch}")

    train_es(args)

    logger.info(f"Embedding space training job finished successfully - epochs: {args.epochs}")
    return {
        "embedding_space": str(Path("embedded_space.pickle").resolve()),
    }


def run_train_knn_job(model_path: Path, sqlite_db_path: Path, strings_cache, job_id: str = None):
    from lib.knn_training import train_knn

    logger.info(f"Starting KNN training job")
    logger.info(f"   Embedding space path: {model_path}")
    
    # Check if the embedding space file exists, try alternative locations if not
    model_path_obj = Path(model_path)
    if not model_path_obj.exists():
        logger.warning(f"⚠️  Embedding space not found at {model_path}, trying alternative locations...")
        
        # Try best_model_package/best_model.pickle (new format)
        best_model_path = model_path_obj.parent / "best_model_package" / "best_model.pickle"
        if best_model_path.exists():
            logger.info(f"✅ Found embedding space in best_model_package: {best_model_path}")
            model_path = best_model_path
        else:
            # Try embedded_space.pickle in the same directory (alternative naming)
            alt_path = model_path_obj.parent / "embedded_space.pickle"
            if alt_path.exists():
                logger.info(f"✅ Found embedding space with alternative name: {alt_path}")
                model_path = alt_path
            else:
                # Try parent directory (in case path is wrong)
                parent_embedded = model_path_obj.parent.parent / "embedded_space.pickle"
                if parent_embedded.exists():
                    logger.info(f"✅ Found embedding space in parent directory: {parent_embedded}")
                    model_path = parent_embedded
                else:
                    # Last resort: try to find any .pickle file in the job directory
                    job_dir = model_path_obj.parent
                    pickle_files = list(job_dir.glob("*.pickle"))
                    if pickle_files:
                        logger.warning(f"⚠️  Using first pickle file found in job directory: {pickle_files[0]}")
                        model_path = pickle_files[0]
                    else:
                        raise FileNotFoundError(f"Embedding space not found at {model_path} and no alternatives found in {job_dir}")
    
    output = train_knn(es_path=model_path, sqlite_db_path=sqlite_db_path, job_id=job_id)

    logger.info(f"KNN training job finished successfully - vector DB: {output['vector_db_path']}")
    return {
        # 'sqlite_db': str(Path(output['sqlite_db']).resolve()),
        'vector_db': str(Path(output['vector_db_path']).resolve()),
    }


def run_build_es_projections_job(model_path, sqlite_db, strings_cache):

    from lib.es_projections import build_es_projections

    logger.info(f"Starting ES projections job")
    logger.info(f"   Embedding space path: {model_path}")
    
    # Check if the embedding space file exists, try alternative locations if not
    model_path_obj = Path(model_path)
    if not model_path_obj.exists():
        logger.warning(f"⚠️  Embedding space not found at {model_path}, trying alternative locations...")
        
        # Try best_model_package/best_model.pickle (new format)
        best_model_path = model_path_obj.parent / "best_model_package" / "best_model.pickle"
        if best_model_path.exists():
            logger.info(f"✅ Found embedding space in best_model_package: {best_model_path}")
            model_path = str(best_model_path)
        else:
            # Try parent directory (in case path is wrong)
            parent_embedded = model_path_obj.parent.parent / "embedded_space.pickle"
            if parent_embedded.exists():
                logger.info(f"✅ Found embedding space in parent directory: {parent_embedded}")
                model_path = str(parent_embedded)
            else:
                # Last resort: try to find any .pickle file in the job directory
                job_dir = model_path_obj.parent
                pickle_files = list(job_dir.glob("*.pickle"))
                if pickle_files:
                    logger.warning(f"⚠️  Using first pickle file found in job directory: {pickle_files[0]}")
                    model_path = str(pickle_files[0])
                else:
                    raise FileNotFoundError(f"Embedding space not found at {model_path} and no alternatives found in {job_dir}")
    
    output = build_es_projections(model_path, sqlite_db)

    logger.info(f"ES projections job finished successfully - projections: {output['projections']}")
    
    # Handle case where preview_png is None (e.g., when all embeddings are NaN)
    result = {
        "projections": str(Path(output['projections']).resolve()),
    }
    
    if output.get('preview_png') is not None:
        result["preview_png"] = str(Path(output['preview_png']).resolve())
    else:
        logger.warning("No preview PNG generated - likely due to invalid embeddings")
        result["preview_png"] = None
    
    return result


def run_train_single_predictor_job(job_spec: dict, job_id: str, data_file: Path, strings_cache: str):
    from lib.single_predictor_training import LightSinglePredictorArgs, train_single_predictor
    
    logger.info(f"Starting single predictor training for job {job_id}")
    logger.info(f"Target column: {job_spec.get('target_column')}, type: {job_spec.get('target_column_type')}")
    
    # Get the embedding space path from the session
    job = load_job("train_single_predictor", job_id)
    session_id = job.get("session_id")
    session = load_session(session_id)
    embedding_space_path = session.get("embedding_space")
    
    if not embedding_space_path or not Path(embedding_space_path).exists():
        raise FileNotFoundError(f"Embedding space not found for session {session_id}. Path: {embedding_space_path}")
    
    logger.info(f"Using embedding space: {embedding_space_path}")
    
    # Check if a custom training file was provided
    custom_training_file = job_spec.get("custom_training_file")
    if custom_training_file:
        training_data_file = Path(custom_training_file)
        if not training_data_file.exists():
            raise FileNotFoundError(f"Custom training file not found: {custom_training_file}")
        logger.info(f"✅ Using custom training file: {custom_training_file}")
    else:
        training_data_file = data_file
        logger.info(f"✅ Using session data file: {data_file}")
    
    # Get user_metadata from job_spec or session
    user_metadata = job_spec.get("user_metadata") or session.get('user_metadata')
    if user_metadata:
        # Validate size
        try:
            validated_user_metadata = validate_user_metadata(user_metadata, max_size_kb=32)
            logger.info(f"✅ User metadata validated for predictor: {len(str(validated_user_metadata))} chars")
        except ValueError as e:
            logger.error(f"❌ Invalid user_metadata: {e}")
            raise
    else:
        validated_user_metadata = None
    
    # Build the arguments dict first to ensure required fields are present
    args_dict = {
        "n_epochs": job_spec.get("n_epochs", 0),
        "batch_size": job_spec.get("batch_size", 0),
        "target_column": job_spec.get("target_column"),  # Required
        "target_column_type": job_spec.get("target_column_type"),  # Required: "set" or "scalar"
        "fine_tune": job_spec.get("fine_tune", True),
        "learning_rate": job_spec.get("learning_rate", 0.0001),
        "input_file": str(training_data_file.resolve()),
        "job_id": job_id,
        "job_queue": "train_single_predictor",
        "string_cache": strings_cache,
        "embedding_space_path": str(embedding_space_path),
        "positive_label": job_spec.get("positive_label"),  # Positive label for binary classification
        "class_imbalance": job_spec.get("class_imbalance"),  # Expected class ratios/counts from real world
        "name": job_spec.get("name"),  # Name for the single predictor
        "session_id": session_id,  # Pass session_id for webhook callbacks
        "webhooks": job_spec.get("webhooks"),  # Pass webhook configuration
        "user_metadata": validated_user_metadata,  # Pass user metadata for identification
    }
    
    # Validate required parameters before creating the object
    if not args_dict["target_column"]:
        raise ValueError("target_column is required in job spec")
    if not args_dict["target_column_type"]:
        raise ValueError("target_column_type is required in job spec")
    
    # Create the args object with all required fields
    args = LightSinglePredictorArgs(**args_dict)
    
    # Run the training
    train_single_predictor(args)
    
    logger.info(f"Single predictor training job finished successfully - target: {args_dict['target_column']}, epochs: {args_dict['n_epochs']}")
    return {
        "single_predictor": str(Path("single_predictor.pickle").resolve()),
        "training_metrics": str(Path("training_metrics.json").resolve()),
    }


def run_train_single_predictor_more_job(job_spec: dict, job_id: str, data_file: Path, strings_cache: str):
    """
    Continue training an existing single predictor for more epochs.
    Loads the existing predictor and resumes training.
    """
    from lib.single_predictor_training import LightSinglePredictorArgs, train_single_predictor, load_single_predictor
    
    logger.info(f"🔄 Starting single predictor continuation training for job {job_id}")
    
    # Get the existing predictor path from job spec
    resume_from_predictor = job_spec.get("resume_from_predictor")
    if not resume_from_predictor or not Path(resume_from_predictor).exists():
        raise FileNotFoundError(f"Existing predictor not found at: {resume_from_predictor}")
    
    logger.info(f"📂 Will resume training from existing predictor: {resume_from_predictor}")
    
    # Get session info
    job = load_job("train_single_predictor", job_id)  # Use same queue name
    session_id = job.get("session_id")
    session = load_session(session_id)
    embedding_space_path = session.get("embedding_space")
    
    if not embedding_space_path or not Path(embedding_space_path).exists():
        raise FileNotFoundError(f"Embedding space not found for session {session_id}. Path: {embedding_space_path}")
    
    # Use the same data file that was used for original training
    sqlite_db_path = session.get('sqlite_db')
    if sqlite_db_path and Path(sqlite_db_path).exists():
        training_data_file = Path(sqlite_db_path)
        logger.info(f"✅ Using processed SQLite database: {sqlite_db_path}")
    else:
        training_data_file = data_file
        logger.info(f"✅ Using session data file: {data_file}")
    
    # Get training parameters from job spec or use existing predictor's values
    target_column = job_spec.get("target_column")
    target_column_type = job_spec.get("target_column_type")
    additional_epochs = job_spec.get("n_epochs", 50)
    
    # Use existing batch_size and learning_rate if not specified
    batch_size = job_spec.get("batch_size", 0)  # 0 = use existing
    learning_rate = job_spec.get("learning_rate")  # None = use existing
    
    # Build args dict for continuation
    args_dict = {
        "n_epochs": additional_epochs,
        "batch_size": batch_size,
        "target_column": target_column,
        "target_column_type": target_column_type,
        "fine_tune": True,  # Continue fine-tuning
        "learning_rate": learning_rate,  # None = use existing from predictor
        "input_file": str(training_data_file.resolve()),
        "job_id": job_id,
        "job_queue": "train_single_predictor",  # Use same queue
        "string_cache": strings_cache,
        "embedding_space_path": str(embedding_space_path),
        "session_id": session_id,
        "resume_from_predictor": resume_from_predictor,  # Key: path to existing predictor to load
    }
    
    # Create the args object
    args = LightSinglePredictorArgs(**args_dict)
    
    # Run the continuation training
    train_single_predictor(args)
    
    logger.info(f"✅ Single predictor continuation training finished - additional {additional_epochs} epochs")
    return {
        "single_predictor": str(Path("single_predictor.pickle").resolve()),
        "training_metrics": str(Path("training_metrics.json").resolve()),
    }


@contextlib.contextmanager
def redirect_output_manager(stdout_path=None, stderr_path=None):
    # Save the original stdout and stderr
    original_stdout = sys.stdout
    original_stderr = sys.stderr
    
    # Get hostname for logging
    HOSTNAME = socket.gethostname()
    
    # Open ONLY stdout file - redirect stderr to stdout as well
    # Use buffering=1 to effectively flush every print right away.
    # This helps diagnose issues where the process quits unexpectedly
    # because every log will be inlucded, even the ones just before the crash.
    stdout_file = open(stdout_path, 'w', buffering=1) if stdout_path else None

    # When redirecting stdout/stderr, we must update the logging config
    # so that the root logger knows to put logs into the new stdout stream.
    # This should work with most logging configs.
    # HOWEVER, if custom loggers are used that do NOT propagate to the 
    # root logger, there may be a problem.

    # Save the original logging handlers for root logger.
    original_stdout_handler = None
    original_stderr_handler = None
    root_logger = logging.getLogger()

    try:
        # Redirect BOTH stdout and stderr to the same stdout file
        if stdout_file:
            sys.stdout = stdout_file
            sys.stderr = stdout_file  # ← KEY CHANGE: stderr goes to stdout

        # Remove existing handlers from the root logger
        root_logger.handlers = []

        # Create single logging handler for stdout with timestamps AND HOSTNAME
        if stdout_file:
            stdout_handler = logging.StreamHandler(stdout_file)
            stdout_handler.setLevel(logging.INFO)  # INFO level - no DEBUG spam
            # Add timestamp formatter WITH HOSTNAME and thread ID
            formatter = logging.Formatter(
                f'%(asctime)s [{HOSTNAME}] [%(levelname)-8s] [%(threadName)s] %(name)-45s: %(message)s',
                datefmt='%Y-%m-%d %H:%M:%S'
            )
            stdout_handler.setFormatter(formatter)
            root_logger.addHandler(stdout_handler)
            root_logger.setLevel(logging.INFO)  # INFO level - suppress DEBUG noise

        yield  # Run the code inside the context
    finally:
        # Restore the original stdout and stderr
        sys.stdout = original_stdout
        sys.stderr = original_stderr
        
        # Close the file if it was opened
        if stdout_file:
            stdout_file.close()

        # Remove the new handlers
        root_logger.handlers = []

        # Optionally: re-add any original handlers if needed
        if original_stdout_handler:
            root_logger.addHandler(original_stdout_handler)
        if original_stderr_handler:
            root_logger.addHandler(original_stderr_handler)


def run_job(queue_name: str, job_id: str, auto_step: bool = True, redirect_output=True):

    # Clear GPU cache before starting a new job to free up any lingering memory
    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            logger.info(f"🧹 Cleared GPU cache before starting job {job_id}")
    except Exception as e:
        logger.debug(f"Could not clear GPU cache: {e}")

    # Print version info at start of every job
    try:
        from version import print_version_banner
        print_version_banner(f"Featrix Sphere Job Worker ({queue_name})")
    except Exception as e:
        print(f"⚠️ Could not load version info: {e}")

    # Ensure all necessary directories exist before starting
    ensure_directories_exist()

    # Load job to get session_id and job_type for new directory structure
    job = load_job(queue_name, job_id)
    session_id = job.get("session_id")
    job_type = job.get("type")
    job_spec = job.get("spec", {})
    
    # Use session_id for output directory (child ES sessions are top-level)
    output_session_id = session_id
    
    # Session ID validation removed - we trust the session ID generation logic

    # Create a working directory for the job
    # New structure: {output_dir}/{session_id}/{job_type}/{job_id}/
    # Use output_session_id (parent for child ES jobs, otherwise session_id)
    if output_session_id and job_type:
        work_dir = config.output_dir / output_session_id / job_type / job_id
        logger.info(f"📁 Using new directory structure: {work_dir}")
        logger.info(f"   Session ID: {session_id}")
        logger.info(f"   Job Type: {job_type}")
        logger.info(f"   Job ID: {job_id}")
    else:
        # Fallback to old structure for backward compatibility
        work_dir = config.output_dir / job_id
        if not session_id:
            logger.error(f"❌ Job {job_id} missing session_id - using OLD directory structure: {work_dir}")
            logger.error(f"   This should NOT happen in production! Job will be created at root of featrix_output")
        if not job_type:
            logger.error(f"❌ Job {job_id} missing job_type - using OLD directory structure: {work_dir}")
            logger.error(f"   This should NOT happen in production! Job will be created at root of featrix_output")
    work_dir.mkdir(exist_ok=True, parents=True)
    
    # Create README.txt documenting supported flags
    readme_path = work_dir / "README.txt"
    if not readme_path.exists():
        readme_content = """Job Control Flags
==================

This directory supports the following control flags:

ABORT
-----
Purpose: Stop a running job immediately and prevent automatic restart.
Usage: Create a file named "ABORT" in this directory.
Effect:
  - Job exits immediately when detected
  - Job is marked as FAILED
  - Job will NOT restart automatically
  - Recovery is blocked permanently

RESTART
-------
Purpose: Restart a job from scratch on next system startup.
Usage: Create a file named "RESTART" in this directory.
Effect (on worker startup):
  - Job is reset to READY status
  - All progress and recovery info is cleared
  - Job runs from the beginning
  - RESTART flag is renamed to started.RESTART.<date>
  - Jobs are processed newest first (by file modification time)

FINISH
------
Purpose: Gracefully stop training after current epoch completes.
Usage: Create a file named "FINISH" in this directory.
Effect:
  - Training completes current epoch
  - Model is saved
  - Job exits gracefully
  - FINISHED flag is created when done

NO_STOP
-------
Purpose: Disable early stopping for this job.
Usage: Create a file named "NO_STOP" in this directory.
Effect:
  - Early stopping mechanisms are disabled
  - Job will run for full number of epochs

PUBLISH
-------
Purpose: Save embedding_space.pickle during training for single predictor training.
Usage: Create a file named "PUBLISH" in this directory.
Effect:
  - Embedding space is saved as embedding_space.pickle in the job directory
  - Saved at the start of each epoch when flag is present
  - Can be used to train a single predictor before training completes
  - Flag can remain present - it will save once per epoch

PAUSE
-----
Purpose: Pause a running job gracefully (saves checkpoint and marks as PAUSED).
Usage: Create a file named "PAUSE" in this directory.
Effect:
  - Training completes current epoch/batch
  - Checkpoint is saved (if save_state_after_every_epoch is enabled)
  - Job is marked as PAUSED
  - Job exits gracefully (can be resumed later)
  - To resume: Remove PAUSE file and set job status to READY

Creating Flags
--------------
You can create flags using ffsh:

  ffsh jobs abort <job_path>
  ffsh jobs finish <job_path>
  ffsh jobs restart <job_path>
  ffsh jobs pause <job_path>
  ffsh jobs resume <job_path>

Where <job_path> can be:
  - <session_id>/<job_type>/<job_id> (full path)
  - <job_id> (just the job ID - will search all queues)

Or manually:
  touch <job_directory>/ABORT
  touch <job_directory>/RESTART
  touch <job_directory>/FINISH
  touch <job_directory>/PAUSE
  touch <job_directory>/NO_STOP
  touch <job_directory>/PUBLISH
"""
        readme_path.write_text(readme_content)
        logger.debug(f"Created README.txt in job directory: {readme_path}")
    
    # Verify directory exists and is accessible before chdir
    if not work_dir.exists():
        raise FileNotFoundError(f"Failed to create working directory: {work_dir}")
    if not work_dir.is_dir():
        raise NotADirectoryError(f"Working directory path is not a directory: {work_dir}")
    
    # Change to working directory and verify it worked
    try:
        os.chdir(work_dir)
        # Verify we're actually in the directory (in case it was deleted between chdir and getcwd)
        actual_cwd = os.getcwd()
        if Path(actual_cwd) != work_dir.resolve():
            raise RuntimeError(f"chdir failed: expected {work_dir.resolve()}, got {actual_cwd}")
    except FileNotFoundError:
        # Directory was deleted after we created it - recreate it
        logger.warning(f"⚠️  Working directory was deleted, recreating: {work_dir}")
        work_dir.mkdir(exist_ok=True, parents=True)
        os.chdir(work_dir)
        # Verify again
        actual_cwd = os.getcwd()
        if Path(actual_cwd) != work_dir.resolve():
            raise RuntimeError(f"chdir failed after recreate: expected {work_dir.resolve()}, got {actual_cwd}")

    if redirect_output:
        logs_dir = work_dir / "logs"
        logs_dir.mkdir(exist_ok=True, parents=True)
        stdout_path = logs_dir / "stdout.log"
        
        # Rotate old log file if it exists (for job restarts/resumes)
        if stdout_path.exists():
            # Find the next available number
            n = 1
            while (logs_dir / f"stdout.log.{n}").exists():
                n += 1
            rotated_path = logs_dir / f"stdout.log.{n}"
            stdout_path.rename(rotated_path)
            logger.info(f"📋 Rotated old log: stdout.log → stdout.log.{n}")
        
        # NO LONGER CREATING STDERR - everything goes to stdout
    else:
        stdout_path = None

    with redirect_output_manager(stdout_path, None):

         # Do the work
        session = None  # Initialize to None in case load_session fails
        job_type = None  # Initialize to None in case job loading fails
        try:
            # Job already loaded above, just get session
            # Use session_id from job (child session for child ES jobs), not output_session_id
            session = load_session(session_id)

            # Get job type early - needed for conditional logic below
            job_type = job.get("type")

            job_spec = job.get("spec")
            if job_spec is None:
                raise ValueError("Job spec is missing")

            # 🔒 DEPENDENCY CHECK: Predictor jobs require completed ES training
            if queue_name == "train_single_predictor":
                embedding_space_path = session.get("embedding_space")
                
                if not embedding_space_path or not Path(embedding_space_path).exists():
                    logger.warning(f"⚠️ Embedding space not ready yet: {embedding_space_path}")
                    logger.warning(f"⏳ Waiting for ES training to complete...")
                    
                    # Wait up to 4 hours for ES training to complete
                    max_wait = 14400  # 4 hours
                    wait_interval = 60  # Check every 60 seconds
                    start_time = time.time()
                    
                    while time.time() - start_time < max_wait:
                        time.sleep(wait_interval)
                        
                        # Reload session to get updated embedding_space path
                        session = load_session(job.get("session_id"))
                        embedding_space_path = session.get("embedding_space")
                        
                        if embedding_space_path and Path(embedding_space_path).exists():
                            logger.info(f"✅ Embedding space now available: {embedding_space_path}")
                            break
                        
                        elapsed_minutes = (time.time() - start_time) / 60
                        logger.info(f"⏳ Still waiting for ES training... ({elapsed_minutes:.1f} minutes elapsed)")
                    else:
                        # Timeout - fail the job
                        error_msg = f"Timeout waiting {max_wait/3600:.1f} hours for ES training to complete"
                        logger.error(f"❌ {error_msg}")
                        raise TimeoutError(error_msg)

            # Execute the task
            session_id = job.get("session_id")
            logger.info(f"running job {job_id}")
            if session_id:
                logger.info(f"   Session ID: {session_id}")
            else:
                logger.error(f"   ❌ WARNING: Job {job_id} has NO session_id!")
            logger.debug(f"Job spec: {json.dumps(job.get('spec'), indent=4)}")

            # Mark the job as running so it doesn't get picked up by another worker
            update_job_status(queue_name=queue_name, job_id=job_id, status=JobStatus.RUNNING)       
            
            input_dir = work_dir / "input_data"
            input_dir.mkdir(exist_ok=True, parents=True)

            # Handle different input data scenarios
            input_data = session.get("input_data")
            session_type = session.get("session_type")
            
            if input_data and input_data.startswith('s3://'):
                # For S3 URLs, the job will handle downloading, so we pass the S3 URL as data_file
                data_file = input_data
            elif session_type == "embedding_space" and job_type == "create_structured_data":
                # For embedding space sessions, the first job will download S3 data
                # The job spec contains the S3 paths needed
                data_file = None  # Will be handled by job spec S3 downloads
            else:
                # For local files, handle both absolute paths and relative paths
                input_path = Path(input_data)
                
                # Check if it's an absolute path that exists
                if input_path.is_absolute() and input_path.exists():
                    # Use absolute path as-is
                    original_data_file = input_path
                    logger.info(f"Using absolute input data path: {original_data_file}")
                elif input_path.is_absolute() and not input_path.exists():
                    # Absolute path doesn't exist - this is the problem
                    error_msg = f"Input data file not found at absolute path: {input_data}. The file must exist on this compute node."
                    logger.error(error_msg)
                    raise FileNotFoundError(error_msg)
                else:
                    # Relative path - resolve relative to config.data_dir
                    original_data_file = config.data_dir / input_data
                    logger.info(f"Resolving relative input data path: {input_data} -> {original_data_file}")
                
                # Check if the input data file exists before trying to copy it
                if not original_data_file.exists():
                    error_msg = f"Input data file not found: {original_data_file}. Make sure the file exists in the data directory or provide an absolute path that exists on this compute node."
                    logger.error(error_msg)
                    raise FileNotFoundError(error_msg)
                
                data_file = input_dir / original_data_file.name
                logger.info(f"Copying input data from {original_data_file} to {data_file}")
                copyfile(original_data_file, data_file)

            # Copy master string cache to job working directory for faster startup
            strings_cache = session.get("strings_cache")
            logger.info(f"🔍 CACHE DEBUG: strings_cache from session = {strings_cache}")
            logger.info(f"🔍 CACHE DEBUG: job_type = {job_type}")
            
            if strings_cache and Path(strings_cache).exists() and job_type != "create_structured_data":
                # Copy master cache to local working directory
                master_cache_path = Path(strings_cache)
                local_cache_path = Path("strings.sqlite3")
                
                logger.info(f"🔍 CACHE DEBUG: master_cache_path = {master_cache_path}")
                logger.info(f"🔍 CACHE DEBUG: master_cache_path.exists() = {master_cache_path.exists()}")
                logger.info(f"🔍 CACHE DEBUG: local_cache_path = {local_cache_path.resolve()}")
                logger.info(f"🔍 CACHE DEBUG: current working directory = {os.getcwd()}")
                
                try:
                    copyfile(master_cache_path, local_cache_path)
                    logger.info(f"📦 Copied master string cache from {master_cache_path} to {local_cache_path}")
                    logger.info(f"   Cache size: {master_cache_path.stat().st_size / 1024:.1f} KB")
                    
                    # Verify the copy worked
                    if local_cache_path.exists():
                        logger.info(f"✅ CACHE DEBUG: Local cache copy verified at {local_cache_path.resolve()}")
                        logger.info(f"   Local cache size: {local_cache_path.stat().st_size / 1024:.1f} KB")
                    else:
                        logger.error(f"❌ CACHE DEBUG: Local cache copy failed - file does not exist")
                    
                    # Update strings_cache to point to local copy
                    strings_cache = str(local_cache_path.resolve())
                    logger.info(f"🔍 CACHE DEBUG: Updated strings_cache path to: {strings_cache}")
                except Exception as e:
                    logger.warning(f"Failed to copy master string cache: {e}")
                    logger.warning(f"🔍 CACHE DEBUG: Exception details: {type(e).__name__}: {str(e)}")
                    logger.info("Proceeding with empty cache - will rebuild as needed")
                    strings_cache = None
            elif job_type == "create_structured_data":
                # For create_structured_data, start with no cache (will create the master)
                logger.info(f"🔍 CACHE DEBUG: create_structured_data job - starting with no cache")
                strings_cache = None
            elif not strings_cache:
                logger.warning(f"🔍 CACHE DEBUG: No strings_cache in session")
                strings_cache = None
            elif strings_cache and not Path(strings_cache).exists():
                logger.warning(f"🔍 CACHE DEBUG: strings_cache path does not exist: {strings_cache}")
                strings_cache = None
            else:
                logger.info(f"🔍 CACHE DEBUG: Unexpected cache state - strings_cache={strings_cache}, job_type={job_type}")
                strings_cache = None

            # TODO: handle ctrl-c exits.

            # Do the work
            if job_type == "echo":
                run_echo_job(job_spec)

            elif job_type == "create_structured_data":
                output = run_create_structured_data(job_spec=job_spec, job_id=job_id, data_file=data_file, session_id=session_id)
                session = {**session, **output}

            elif job_type == "generate_movie_frame":
                # Movie frame generation on CPU - non-blocking
                run_generate_movie_frame(job_spec=job_spec, job_id=job_id)

            elif job_type == "train_es":
                # CRITICAL FIX: Use processed SQLite database instead of original CSV
                sqlite_db_path = session.get('sqlite_db')
                if sqlite_db_path and Path(sqlite_db_path).exists():
                    logger.info(f"🔍 train_es using processed SQLite database: {sqlite_db_path}")
                    data_file_for_training = Path(sqlite_db_path)
                else:
                    logger.warning(f"⚠️  SQLite database not found in session, falling back to original data file")
                    data_file_for_training = data_file
                
                output = run_train_es_job(
                    job_spec=job_spec, 
                    job_id=job_id, 
                    data_file=data_file_for_training,
                    strings_cache=strings_cache,  # Use local copy
                )
                session = {**session, **output}

            elif job_type == "train_knn":
                # Skip train_knn for JSON ES sessions - they don't need KNN training
                # JSON ES sessions are child sessions used only for encoding JSON columns
                if session_id and session_id.startswith("json-es-"):
                    logger.info(f"⏭️  Skipping train_knn for JSON ES session {session_id} (not needed for child ES)")
                    logger.info(f"   JSON ES sessions are used only for encoding, not for similarity search")
                    # Mark as complete without running
                    output = {
                        'vector_db': None,  # No vector DB for JSON ES sessions
                    }
                else:
                    output = run_train_knn_job(
                        model_path=Path(session['embedding_space']), 
                        sqlite_db_path=Path(session['sqlite_db']),
                        strings_cache=strings_cache,  # Use local copy
                        job_id=job_id,
                    )
                session = {**session, **output}

            elif job_type == "build_es_projections":
                output = run_build_es_projections_job(
                    model_path=Path(session['embedding_space']), 
                    sqlite_db=Path(session['sqlite_db']),
                    strings_cache=strings_cache  # Use local copy
                )
                session = {**session, **output}

            elif job_type == "train_single_predictor":
                # Single predictor training requires a pre-trained embedding space
                embedding_space_path = session.get("embedding_space")
                if not embedding_space_path:
                    raise ValueError("No embedding space found for single predictor training")
                
                # Copy embedding space to working directory
                embedding_space_src = Path(embedding_space_path)
                embedding_space_dst = Path("embedded_space.pickle")
                copyfile(embedding_space_src, embedding_space_dst)
                
                # CRITICAL FIX: Use processed SQLite database instead of original CSV
                sqlite_db_path = session.get('sqlite_db')
                if sqlite_db_path and Path(sqlite_db_path).exists():
                    logger.info(f"🔍 train_single_predictor using processed SQLite database: {sqlite_db_path}")
                    data_file_for_training = Path(sqlite_db_path)
                else:
                    logger.warning(f"⚠️  SQLite database not found in session, falling back to original data file")
                    data_file_for_training = data_file
                
                output = run_train_single_predictor_job(
                    job_spec=job_spec, 
                    job_id=job_id, 
                    data_file=data_file_for_training,
                    strings_cache=strings_cache  # Use local copy
                )
                
                # For multiple single predictors, we need to track which predictor this is
                # and save to the correct position in the arrays
                job_plan = session.get("job_plan", [])
                predictor_index = None
                
                # Find the job description that matches this job_id
                for job_desc in job_plan:
                    if job_desc.get("job_id") == job_id and job_desc.get("job_type") == "train_single_predictor":
                        predictor_index = job_desc.get("predictor_index")
                        break
                
                if predictor_index is not None:
                    # Multiple predictors - save to correct array position
                    single_predictors = session.get("single_predictors", [])
                    single_predictors = single_predictors or []
                    training_metrics = session.get("training_metrics", [])
                    training_metrics = training_metrics or []
                    if type(training_metrics) == str:
                        try:
                            training_metrics = json.loads(training_metrics)
                        except Exception as e:
                            logger.error(f"Failed to load training metrics: {e}... training_metrics={training_metrics}")
                    if (type(training_metrics) != list):
                        logger.error(f"Training metrics is not a list, converting to empty list from {training_metrics}...!")
                        training_metrics = []
                    
                    while len(single_predictors) <= predictor_index:
                        single_predictors.append(None)
                    while len(training_metrics) <= predictor_index:
                        try:
                            training_metrics.append(None)
                        except Exception as ex:
                            logger.error(f"Failed to append None to training metrics: {ex}... training_metrics={training_metrics}")
                            training_metrics = [None]

                    # Save artifacts to correct positions
                    sp_path = output.get("single_predictor")
                    tm_path = output.get("training_metrics")
                    
                    single_predictors[predictor_index] = sp_path
                    training_metrics[predictor_index] = tm_path
                    
                    session["single_predictors"] = single_predictors
                    session["training_metrics"] = training_metrics
                    
                    logger.info(f"Saved single predictor {predictor_index} artifacts:")
                    logger.info(f"  Single predictor: {sp_path}")
                    logger.info(f"  Training metrics: {tm_path}")
                    logger.info(f"  File exists check: SP={Path(sp_path).exists() if sp_path else False}, TM={Path(tm_path).exists() if tm_path else False}")
                    
                    # Verify training metrics file was actually created
                    if tm_path and not Path(tm_path).exists():
                        logger.error(f"❌ CRITICAL: Training metrics file not found at expected path: {tm_path}")
                        logger.error(f"   This will cause 404 errors when trying to access training metrics!")
                        # List files in the working directory to debug
                        try:
                            cwd_files = list(Path(".").glob("*"))
                            logger.error(f"   Files in working directory: {cwd_files}")
                        except Exception as e:
                            logger.error(f"   Could not list working directory: {e}")
                    else:
                        logger.info(f"✅ Training metrics file verified at: {tm_path}")
                else:
                    # Backward compatibility - single predictor (old style)
                    sp_path = output.get("single_predictor")
                    tm_path = output.get("training_metrics")
                    
                    session = {**session, **output}
                    
                    # Also update single_predictors[0] if it exists but is null/empty
                    if "single_predictors" in session and isinstance(session.get("single_predictors"), list):
                        single_predictors = session.get("single_predictors", [])
                        if len(single_predictors) == 0:
                            single_predictors = [sp_path]
                        elif len(single_predictors) == 1 and (single_predictors[0] is None or not single_predictors[0]):
                            single_predictors[0] = sp_path
                        session["single_predictors"] = single_predictors
                        logger.info(f"Updated single_predictors[0] to: {sp_path}")
                    
                    logger.info(f"Backward compatibility: single predictor artifacts:")
                    logger.info(f"  Single predictor: {sp_path}")
                    logger.info(f"  Training metrics: {tm_path}")
                    logger.info(f"  File exists check: SP={Path(sp_path).exists() if sp_path else False}, TM={Path(tm_path).exists() if tm_path else False}")
                    
                    # Verify training metrics file for legacy format too
                    if tm_path and not Path(tm_path).exists():
                        logger.error(f"❌ CRITICAL: Legacy training metrics file not found at: {tm_path}")
                        logger.error(f"   This will cause 404 errors when trying to access training metrics!")
                    else:
                        logger.info(f"✅ Legacy training metrics file verified at: {tm_path}")
                
            elif job_type == "train_single_predictor_more":
                # Continue training an existing single predictor
                embedding_space_path = session.get("embedding_space")
                if not embedding_space_path:
                    raise ValueError("No embedding space found for predictor continuation")
                
                # Copy embedding space to working directory
                embedding_space_src = Path(embedding_space_path)
                embedding_space_dst = Path("embedded_space.pickle")
                copyfile(embedding_space_src, embedding_space_dst)
                
                # Use processed SQLite database
                sqlite_db_path = session.get('sqlite_db')
                if sqlite_db_path and Path(sqlite_db_path).exists():
                    logger.info(f"🔍 train_single_predictor_more using processed SQLite database: {sqlite_db_path}")
                    data_file_for_training = Path(sqlite_db_path)
                else:
                    logger.warning(f"⚠️  SQLite database not found in session, falling back to original data file")
                    data_file_for_training = data_file
                
                # Use the same queue as train_single_predictor for continuation
                # The job runner will handle loading the existing predictor
                output = run_train_single_predictor_more_job(
                    job_spec=job_spec,
                    job_id=job_id,
                    data_file=data_file_for_training,
                    strings_cache=strings_cache
                )
                
                # Update the same predictor (don't create a new one)
                predictor_index = job_spec.get("predictor_index")
                if predictor_index is not None:
                    single_predictors = session.get("single_predictors", [])
                    training_metrics = session.get("training_metrics", [])
                    
                    if isinstance(training_metrics, str):
                        try:
                            training_metrics = json.loads(training_metrics)
                        except:
                            training_metrics = []
                    if not isinstance(training_metrics, list):
                        training_metrics = []
                    
                    # Ensure lists are long enough
                    while len(single_predictors) <= predictor_index:
                        single_predictors.append(None)
                    while len(training_metrics) <= predictor_index:
                        training_metrics.append(None)
                    
                    # Update existing predictor
                    sp_path = output.get("single_predictor")
                    tm_path = output.get("training_metrics")
                    single_predictors[predictor_index] = sp_path
                    training_metrics[predictor_index] = tm_path
                    
                    session["single_predictors"] = single_predictors
                    session["training_metrics"] = training_metrics
                    
                    logger.info(f"✅ Updated predictor {predictor_index} with continuation training")
                else:
                    # Backward compatibility
                    session = {**session, **output}

            # test jobs
            elif job_type == "echo_task_1":
                pass
            elif job_type == "echo_task_2":
                pass
            else:
                raise ValueError(f"Unsupported job type {job_type}")

            update_job_status(queue_name=queue_name, job_id=job_id, status=JobStatus.DONE)
            save_session(session_id=session.get("session_id"), session_doc=session, exist_ok=True)

            # CRITICAL: For create_structured_data jobs, call step_session immediately to queue next job
            # This ensures the next job in the plan is queued even if the worker doesn't call step_session
            if job_type == "create_structured_data":
                current_session_id = session.get("session_id")
                if current_session_id:
                    logger.info(f"🚀 create_structured_data job {job_id} completed, calling step_session for session {current_session_id}")
                    try:
                        time.sleep(0.1)  # Small delay to ensure locks are released
                        step_session(session_id=current_session_id)
                        logger.info(f"✅ step_session completed successfully for session {current_session_id}")
                    except Exception as step_error:
                        logger.error(f"❌ CRITICAL: Failed to step session {current_session_id} after create_structured_data job {job_id}: {step_error}")
                        logger.error(f"   This means the next job in the plan was NOT queued!")
                        logger.error(f"   Traceback: {traceback.format_exc()}")
                        # Don't fail the job - stepping is a separate concern
                        logger.error(f"🔴🔴🔴 SESSION STEP FAILED - MANUAL INTERVENTION NEEDED 🔴🔴🔴")
                else:
                    logger.error(f"❌ CRITICAL: create_structured_data job {job_id} has no session_id in session, cannot step session!")
                    logger.error(f"   Session keys: {list(session.keys())}")
                    logger.error(f"   This means the next job in the plan was NOT queued!")

            # NOTE: For other job types, we do NOT call step_session() from within job execution to avoid deadlocks.
            # step_session() will be called externally by workers or a separate process after
            # the job completes. This prevents lock contention between job execution and session stepping.
            # The session status is already updated to reflect completion, so external processes
            # can detect when to step the session.
            
            try:
                if torch.cuda.is_available():
                    torch.cuda.reset_peak_memory_stats()
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
            except Exception as e:
                logger.warning(f"Error resetting CUDA stats: {e}")

        except Exception as e:
            logger.error(f"Error executing job {job_id}:")
            logger.error(f"Exception type: {type(e).__name__}")
            logger.error(f"Exception message: {str(e)}")
            logger.error(f"Full traceback:\n{traceback.format_exc()}")

            # Check if this was an ABORT exit (sys.exit(2))
            is_abort = isinstance(e, SystemExit) and getattr(e, 'code', None) == 2
            if is_abort:
                logger.warning(f"🚫 Job {job_id} aborted with sys.exit(2) - worker will restart after cleanup")

            # Post crash to Slack
            try:
                from slack import send_slack_message
                
                slack_msg = f"💥 **Job Crashed**\n"
                slack_msg += f"Job ID: {job_id}\n"
                slack_msg += f"Queue: {queue_name}\n"
                slack_msg += f"Job Type: {job_type or 'unknown'}\n"
                slack_msg += f"Session ID: {session.get('session_id', session_id) if session else session_id}\n"
                slack_msg += f"Exception: {type(e).__name__}: {str(e)}\n"
                slack_msg += f"```\n{traceback.format_exc()}\n```"
                
                send_slack_message(slack_msg)
                logger.info("✅ Slack notification sent for job crash")
            except Exception as slack_error:
                logger.warning(f"Failed to send Slack notification for crash: {slack_error}")

            try:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
            except Exception as cuda_e:
                logger.warning(f"Error with CUDA cleanup: {cuda_e}")

            update_job_status(queue_name=queue_name, job_id=job_id, status=JobStatus.FAILED)
            
            # Only mark session as failed for core jobs, not for single predictor failures
            # Only if we successfully loaded the session
            if session is not None:
                if job_type and job_type != "train_single_predictor":
                    session["status"] = SessionStatus.FAILED
                else:
                    logger.warning(f"Single predictor job {job_id} failed, but keeping session active for other predictors")
                    
                save_session(session_id=session.get("session_id"), session_doc=session, exist_ok=True)
            else:
                logger.error(f"⚠️ Cannot update session status - session {session_id} could not be loaded")

            # Return exit code 2 to signal worker should restart
            if is_abort:
                logger.warning(f"🔄 Returning exit code 2 to trigger worker restart")
                return 2

            return e

        logger.info(f"Job {job_id} finished")

        return None


def print_job_dir(queue_name: str, job_id: str):
    from asciitree import LeftAligned

    from tree import build_tree

    # Use helper function to get job output path (handles new structure)
    job_path = get_job_output_path(job_id)

    tree = build_tree(job_path)

    tree_builder = LeftAligned()
    tree_string = tree_builder(tree)
    
    logger.info(tree_string)





##############################################################################
# 
# SESSIONS
# 
##############################################################################


class SessionStatus(Enum):
    READY = "ready"
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"
    CANCELLED = "cancelled"


def does_session_exist(session_id: str):
    """Check if a session exists by session_id."""
    session_path = config.session_dir / f"{session_id}.session"
    return session_path.is_file()

def create_echo_session(session_id: str, created_at: datetime, input_filename: str | None = None):

    input_filename = input_filename or "test.csv"

    session_doc = dict(
        created_at=created_at,
        session_type="echo",
        session_id=session_id,
        status=SessionStatus.READY,
        job_plan=[
            {
                "job_type": "echo_task_1",
                "spec": {
                    "spec_a": "abc",
                    "spec_b": "def",
                }, 
                "job_id": None,
            },
            {
                "job_type": "echo_task_2",
                "spec": {
                    "spec_1": "123",
                    "spec_2": "456",
                },
                "job_id": None,
            },
        ],
        input_data=input_filename,
        # Path to the file containing the output artifact
        task_1_output=None,
        task_2_output=None,
        # Any other stuff shared across jobs should be listed
        # as a property of the session.
    )

    return session_doc


def create_sphere_session(session_id: str, created_at: datetime, input_filename: str | None = None, name: str = None, single_predictors: list = None, epochs: int = None, column_overrides: dict = None, string_list_delimiter: str = "|", movie_frame_interval: int = 3, weightwatcher_save_every: int = 5, important_columns_for_visualization: list = None, user_metadata: dict = None):

    # Allow sphere sessions without input_filename for testing/API usage
    if not input_filename:
        logger.info(f"Creating sphere session {session_id} without input file (API/testing mode)")
        input_filename = "api_created_session.csv"  # Placeholder filename
    single_predictors = single_predictors or []
    important_columns_for_visualization = important_columns_for_visualization or []

    # Import here to avoid circular dependencies
    from lib.sphere_config import get_row_limit
    
    # Get row limit from config
    default_row_limit = get_row_limit()

    # Build job plan with basic sphere jobs
    job_plan = [
        dict(
            job_type="create_structured_data",
            spec={
                "column_overrides": column_overrides or {},
                "string_list_delimiter": string_list_delimiter,
            },
            job_id=None,
        ),
        dict(
            job_type= "train_es",
            spec={
                # Use provided epochs or default to auto-calculate
                "epochs": epochs if epochs is not None else 0,
                # Row limit from config (customizable via /sphere/app/config.json)
                "row_limit": default_row_limit,
                "is_production": True,
                "ignore_cols": [],
                "learning_rate": 0.0003,
                "batch_size": 0,  # Auto-calculate based on dataset size
                # DropoutScheduler parameters
                "enable_dropout_scheduler": True,
                "dropout_schedule_type": "linear_decay",
                "initial_dropout": 0.5,
                "final_dropout": 0.1,
                # Movie frame and WeightWatcher intervals (configurable via API)
                "movie_frame_interval": movie_frame_interval,
                "weightwatcher_save_every": weightwatcher_save_every,
                # VISUALIZATION ONLY: Prefer rows with non-null values in these columns for epoch projections
                # ⚠️ IMPORTANT: Has NO EFFECT on model training, only affects visualization sampling
                "important_columns_for_visualization": important_columns_for_visualization,
                # User metadata for identification
                "user_metadata": user_metadata,
            },
            job_id=None,
        ),
        dict(
            job_type="train_knn",
            spec={},
            job_id=None,
        ),
        dict(
            job_type="build_es_projections",
            spec={},
            job_id=None,
        ), 
    ]
    
    # Add single predictor training jobs to the plan
    single_predictors_artifacts = []
    training_metrics_artifacts = []
    
    for i, predictor_spec in enumerate(single_predictors):
        # Convert pydantic model to dict if needed
        if hasattr(predictor_spec, 'dict'):
            spec_dict = predictor_spec.dict()
        else:
            spec_dict = predictor_spec
        
        # Add user_metadata to predictor spec if not already present
        if user_metadata and 'user_metadata' not in spec_dict:
            spec_dict['user_metadata'] = user_metadata
            
        # Add job to plan
        job_plan.append(dict(
            job_type="train_single_predictor",
            spec=spec_dict,
            job_id=None,
            predictor_index=i,  # Track which predictor this is
        ))
        
        # Initialize artifact paths (will be filled when jobs complete)
        single_predictors_artifacts.append(None)
        training_metrics_artifacts.append(None)

    session_doc = dict(
        created_at=created_at,
        session_type="sphere",
        session_id=session_id,
        name=name,  # Optional name for identification and metadata
        status=SessionStatus.READY,
        job_plan=job_plan,
        # Paths to artifacts that are shared across jobs.
        input_data=input_filename,
        embedding_space=None,
        sqlite_db=None,
        vector_db=None,
        projections=None,
        preview_png=None,
        strings_cache=None,
        # Support multiple single predictors
        single_predictors=single_predictors_artifacts,
        training_metrics=training_metrics_artifacts,
        # User metadata for identification
        user_metadata=user_metadata,
    )

    return session_doc


def create_cloned_session(session_id: str, created_at: datetime, embedding_space_path: str, strings_cache_path: str | None = None,
                         source_session_id: str = None, source_compute: str = None, validation_loss: float | None = None,
                         epoch: int | None = None, training_metadata: dict | None = None, name: str | None = None) -> dict:
    """
    Create a ready session from a cloned embedding space.
    
    Args:
        session_id: New session ID (should have prefix extracted from source if applicable)
        created_at: Creation timestamp
        embedding_space_path: Path to the cloned embedding space pickle
        strings_cache_path: Path to the cloned strings cache (optional)
        source_session_id: Original session ID that was cloned
        source_compute: Compute node where the original session was located
        validation_loss: Current validation loss from training
        epoch: Current epoch from training
        training_metadata: Full training metadata dictionary
        name: Optional name for the session
        
    Returns:
        Session document dictionary
    """
    session_doc = dict(
        created_at=created_at,
        session_type="sphere",
        session_id=session_id,
        name=name,
        status=SessionStatus.READY,  # Ready - no job_plan needed
        job_plan=[],  # Empty - session is ready to use
        input_data=None,  # No input data for cloned sessions
        embedding_space=embedding_space_path,
        sqlite_db=None,
        vector_db=None,
        projections=None,
        preview_png=None,
        strings_cache=strings_cache_path,
        single_predictors=[],
        training_metrics=[],
        clone_metadata={
            "source_session_id": source_session_id,
            "source_compute": source_compute,
            "validation_loss": validation_loss,
            "epoch": epoch,
            "training_metadata": training_metadata,
            "cloned_at": created_at.isoformat() if created_at else None,
        }
    )
    
    return session_doc


def create_predictor_session(session_id: str, created_at: datetime, input_filename: str | None = None, target_spec: dict = None):
    """Create a session that trains an embedding space and then a single predictor."""
    
    if not input_filename:
        raise ValueError("input_filename is required for predictor sessions")
    
    if target_spec is None:
        raise ValueError("target_spec is required for predictor session")
    
    # Validate target_spec
    if "target_column" not in target_spec:
        raise ValueError("target_column is required in target_spec")
    if "target_column_type" not in target_spec:
        raise ValueError("target_column_type is required in target_spec")

    session_doc = dict(
        created_at=created_at,
        session_type="predictor",
        session_id=session_id,
        status=SessionStatus.READY,
        job_plan=[
            dict(
                job_type="create_structured_data",
                spec={},
                job_id=None,
            ),
            dict(
                job_type="train_es",
                spec={
                    "epochs": 0,
                    "row_limit": 1_000_000,  # 1M default, use 10K for testing/samples
                    "is_production": True,
                    "ignore_cols": [],
                    "learning_rate": 0.0003,
                    "batch_size": 0,  # Auto-calculate based on dataset size
                    # DropoutScheduler parameters
                    "enable_dropout_scheduler": True,
                    "dropout_schedule_type": "linear_decay",
                    "initial_dropout": 0.5,
                    "final_dropout": 0.1,
                    # Movie frame and WeightWatcher intervals (set to 5 by default for better performance)
                    "movie_frame_interval": 5,
                    "weightwatcher_save_every": 5,
                },
                job_id=None,
            ),
            dict(
                job_type="train_single_predictor",
                spec=target_spec,
                job_id=None,
            ),
        ],
        # Paths to artifacts that are shared across jobs.
        input_data=input_filename,
        embedding_space=None,
        sqlite_db=None,
        strings_cache=None,
        single_predictor=None,
        training_metrics=None,
    )

    return session_doc


def create_foundation_model_session(session_id: str, created_at: datetime, foundation_model_id: str, target_spec: dict, input_filename: str | None = None, name: str | None = None):
    """
    Create a session that trains a predictor on a foundation model (existing embedding space).
    
    This function REUSES the foundation model's embedding space - it does NOT create a new one.
    The embedding space and strings cache are copied/referenced from the foundation model session.
    
    Args:
        session_id: New session ID for the predictor training
        created_at: Creation timestamp
        foundation_model_id: Session ID of the foundation model (must have a trained embedding space)
        target_spec: Dict with target_column, target_column_type, and other training parameters
        input_filename: Optional input data file (if None, will use foundation model's data)
        name: Optional name for the session
        
    Returns:
        Session document dictionary
    """    
    logger.info(f"🏗️  Creating foundation model session {session_id} based on foundation {foundation_model_id}")
    logger.info(f"   This session will REUSE the foundation's embedding space - NO new ES will be created")
    
    # Load the foundation model session to get its embedding space
    foundation_session = load_session(foundation_model_id)
    embedding_space_path = foundation_session.get("embedding_space")
    strings_cache_path = foundation_session.get("strings_cache")
    
    # Track embedding space metadata for the predictor session
    embedding_space_md5 = foundation_session.get("embedding_space_md5")
    embedding_space_ctime = foundation_session.get("embedding_space_ctime")
    
    # Resolve to absolute path and validate
    if embedding_space_path:
        embedding_space_path = str(Path(embedding_space_path).resolve())
        if not Path(embedding_space_path).exists():
            # Try to find the embedding space in common locations
            logger.warning(f"⚠️  Embedding space not found at expected path: {embedding_space_path}")
            logger.info(f"🔍 Starting comprehensive search for embedding space file for foundation_model_id: {foundation_model_id}")
            logger.info(f"   Expected path from session: {embedding_space_path}")
            
            # Common locations to search
            search_paths = []
            search_attempts = []
            
            # 1. Try relative to session directory
            logger.info(f"📍 Search Strategy 1: Checking session directory relative paths")
            session_dir = config.session_dir
            if session_dir:
                logger.info(f"   Session directory: {session_dir}")
                candidate1 = session_dir.parent / "featrix_output" / foundation_model_id / "embedded_space.pickle"
                candidate2 = session_dir.parent / "featrix_output" / foundation_model_id / "embedding_space.pickle"
                logger.info(f"   Candidate 1: {candidate1} (exists: {candidate1.exists()})")
                logger.info(f"   Candidate 2: {candidate2} (exists: {candidate2.exists()})")
                search_paths.append(candidate1)
                search_paths.append(candidate2)
                search_attempts.append(("session_dir_relative", [candidate1, candidate2]))
            else:
                logger.warning(f"   ⚠️  Session directory not configured")
            
            # 2. Try in featrix_output directories (always in /sphere/app/)
            logger.info(f"📍 Search Strategy 2: Searching /sphere/app/featrix_output")
            featrix_output_dir = Path("/sphere/app/featrix_output")
            logger.info(f"   Checking if /sphere/app/featrix_output exists: {featrix_output_dir.exists()}")
            if featrix_output_dir.exists():
                logger.info(f"   ✅ /sphere/app/featrix_output exists, scanning for foundation_model_id: {foundation_model_id}")
                
                # Look for any subdirectory with the foundation_model_id
                matching_subdirs = []
                all_subdirs = []
                try:
                    for subdir in featrix_output_dir.iterdir():
                        all_subdirs.append(subdir.name)
                        if subdir.is_dir() and foundation_model_id in str(subdir):
                            matching_subdirs.append(subdir)
                            logger.info(f"   ✅ Found matching subdirectory: {subdir}")
                            search_paths.append(subdir / "embedded_space.pickle")
                            search_paths.append(subdir / "embedding_space.pickle")
                            
                            # Also check train_es subdirectories
                            train_es_dir = subdir / "train_es"
                            logger.info(f"      Checking train_es subdirectory: {train_es_dir} (exists: {train_es_dir.exists()})")
                            if train_es_dir.exists():
                                train_es_subdirs = []
                                for es_subdir in train_es_dir.iterdir():
                                    if es_subdir.is_dir():
                                        train_es_subdirs.append(es_subdir.name)
                                        es_candidate1 = es_subdir / "embedded_space.pickle"
                                        es_candidate2 = es_subdir / "embedding_space.pickle"
                                        logger.info(f"         Checking: {es_candidate1} (exists: {es_candidate1.exists()})")
                                        logger.info(f"         Checking: {es_candidate2} (exists: {es_candidate2.exists()})")
                                        search_paths.append(es_candidate1)
                                        search_paths.append(es_candidate2)
                                        
                                        # Also check for best_model_package
                                        best_model_pkg = es_subdir / "best_model_package" / "best_model.pickle"
                                        logger.info(f"         Checking best_model_package: {best_model_pkg} (exists: {best_model_pkg.exists()})")
                                        if best_model_pkg.exists():
                                            search_paths.append(best_model_pkg)
                                logger.info(f"      Found {len(train_es_subdirs)} train_es subdirectories: {train_es_subdirs}")
                        else:
                            if subdir.is_dir():
                                logger.debug(f"      Skipping subdirectory (no match): {subdir.name}")
                except Exception as e:
                    logger.error(f"   ❌ Error iterating featrix_output_dir: {e}")
                    import traceback
                    logger.error(traceback.format_exc())
                
                logger.info(f"   Total subdirectories in featrix_output: {len(all_subdirs)}")
                logger.info(f"   Matching subdirectories found: {len(matching_subdirs)}")
                if all_subdirs:
                    logger.info(f"   Sample subdirectories: {all_subdirs[:10]}")
                
                # Also do a broader search: look for any train_es subdirectories
                logger.info(f"📍 Search Strategy 2b: Recursive search for train_es directories")
                train_es_dirs_found = []
                try:
                    for train_es_dir in featrix_output_dir.rglob("train_es"):
                        if train_es_dir.is_dir():
                            train_es_dirs_found.append(train_es_dir)
                            logger.info(f"   Found train_es directory: {train_es_dir}")
                            for es_subdir in train_es_dir.iterdir():
                                if es_subdir.is_dir():
                                    # Check if this might be the right one by checking parent directories
                                    path_str = str(es_subdir)
                                    parent_str = str(es_subdir.parent)
                                    matches = foundation_model_id in path_str or foundation_model_id in parent_str
                                    logger.info(f"      Checking subdir: {es_subdir.name} (matches foundation_model_id: {matches})")
                                    
                                    # Check for embedding space files
                                    for es_file in ["embedded_space.pickle", "embedding_space.pickle"]:
                                        es_path = es_subdir / es_file
                                        if es_path.exists():
                                            logger.info(f"         ✅ Found {es_file} at: {es_path}")
                                            if matches:
                                                logger.info(f"         ✅ Path matches foundation_model_id, adding to search")
                                                search_paths.append(es_path)
                                            else:
                                                logger.info(f"         ⚠️  Path doesn't match foundation_model_id, skipping")
                                    
                                    # Also check best_model_package
                                    best_model_pkg = es_subdir / "best_model_package" / "best_model.pickle"
                                    if best_model_pkg.exists():
                                        logger.info(f"         ✅ Found best_model.pickle at: {best_model_pkg}")
                                        if matches:
                                            logger.info(f"         ✅ Path matches foundation_model_id, adding to search")
                                            search_paths.append(best_model_pkg)
                except Exception as e:
                    logger.error(f"   ❌ Error in recursive train_es search: {e}")
                    import traceback
                    logger.error(traceback.format_exc())
                
                logger.info(f"   Total train_es directories found: {len(train_es_dirs_found)}")
            else:
                logger.warning(f"   ❌ /sphere/app/featrix_output does not exist!")
            
            # 3. Try the original path with /sphere/app/ as base
            logger.info(f"📍 Search Strategy 3: Reconstructing path with /sphere/app/ base")
            original_path = Path(embedding_space_path)
            logger.info(f"   Original path: {original_path}")
            logger.info(f"   Original path parts: {original_path.parts}")
            if original_path.name in ["embedded_space.pickle", "embedding_space.pickle"]:
                base_path = Path("/sphere/app")
                logger.info(f"   Base path: {base_path} (exists: {base_path.exists()})")
                if base_path.exists():
                    relative_parts = original_path.parts
                    # Find where "featrix_output" appears
                    try:
                        featrix_idx = [i for i, part in enumerate(relative_parts) if part == "featrix_output"][0]
                        logger.info(f"   Found 'featrix_output' at index: {featrix_idx}")
                        reconstructed = base_path / Path(*relative_parts[featrix_idx:])
                        logger.info(f"   Reconstructed path: {reconstructed} (exists: {reconstructed.exists()})")
                        search_paths.append(reconstructed)
                    except (IndexError, ValueError) as e:
                        logger.warning(f"   ⚠️  Could not find 'featrix_output' in path parts: {e}")
            else:
                logger.info(f"   Original path name '{original_path.name}' is not an embedding space file, skipping reconstruction")
            
            # Search all candidate paths
            logger.info(f"📍 Final Search: Checking {len(search_paths)} candidate paths")
            found_path = None
            for idx, candidate_path in enumerate(search_paths):
                exists = candidate_path.exists()
                is_file = candidate_path.is_file() if exists else False
                logger.info(f"   [{idx+1}/{len(search_paths)}] {candidate_path}")
                logger.info(f"       exists: {exists}, is_file: {is_file}")
                if exists and is_file:
                    found_path = str(candidate_path.resolve())
                    logger.info(f"   ✅✅✅ FOUND EMBEDDING SPACE AT: {found_path}")
                    break
                elif exists:
                    logger.warning(f"       ⚠️  Path exists but is not a file (might be a directory)")
            
            if found_path:
                embedding_space_path = found_path
                logger.info(f"✅✅✅ Successfully found and using embedding space: {embedding_space_path}")
                
                # Calculate MD5 and ctime for found path                
                es_path = Path(embedding_space_path)
                if es_path.exists():
                    # Calculate MD5
                    md5_hash = hashlib.md5()
                    with open(es_path, 'rb') as f:
                        for chunk in iter(lambda: f.read(4096), b""):
                            md5_hash.update(chunk)
                    embedding_space_md5 = md5_hash.hexdigest()
                    
                    # Get creation time
                    ctime = os.path.getctime(es_path)
                    embedding_space_ctime = datetime.fromtimestamp(ctime).isoformat()
                    
                    logger.info(f"   Embedding space file info:")
                    logger.info(f"      Path: {embedding_space_path}")
                    logger.info(f"      MD5: {embedding_space_md5}")
                    logger.info(f"      Created: {embedding_space_ctime}")
            else:
                # Provide helpful error message with all search details
                logger.error(f"❌❌❌ EMBEDDING SPACE SEARCH FAILED")
                logger.error(f"   Foundation model ID: {foundation_model_id}")
                logger.error(f"   Expected path from session: {embedding_space_path}")
                logger.error(f"   Total candidate paths checked: {len(search_paths)}")
                logger.error(f"   /sphere/app/featrix_output exists: {Path('/sphere/app/featrix_output').exists()}")
                if Path('/sphere/app/featrix_output').exists():
                    try:
                        subdirs = list(Path('/sphere/app/featrix_output').iterdir())
                        logger.error(f"   Subdirectories in /sphere/app/featrix_output: {[d.name for d in subdirs if d.is_dir()][:20]}")
                    except Exception as e:
                        logger.error(f"   Error listing subdirectories: {e}")
                
                error_msg = (
                    f"Foundation model {foundation_model_id} embedding space file does not exist on this compute node.\n"
                    f"Expected path: {embedding_space_path}\n"
                    f"Searched {len(search_paths)} candidate locations in /sphere/app/featrix_output.\n"
                    f"The foundation model may be on a different compute node. "
                    f"Please clone the embedding space to this node first using clone_embedding_space, "
                    f"or ensure the embedding space file exists at the expected location."
                )
                logger.error(f"❌ {error_msg}")
                raise ValueError(error_msg)
        else:
            logger.info(f"✅ Found foundation embedding space: {embedding_space_path}")
            
            # Get MD5 and ctime if not already set
            if embedding_space_md5 is None:
                es_path = Path(embedding_space_path)
                if es_path.exists():
                    # Calculate MD5
                    md5_hash = hashlib.md5()
                    with open(es_path, 'rb') as f:
                        for chunk in iter(lambda: f.read(4096), b""):
                            md5_hash.update(chunk)
                    embedding_space_md5 = md5_hash.hexdigest()
                    
                    # Get creation time
                    ctime = os.path.getctime(es_path)
                    embedding_space_ctime = datetime.fromtimestamp(ctime).isoformat()
                    
                    logger.info(f"   Embedding space file info:")
                    logger.info(f"      Path: {embedding_space_path}")
                    logger.info(f"      MD5: {embedding_space_md5}")
                    logger.info(f"      Created: {embedding_space_ctime}")
    
    # If embedding_space_path is still None, try to search for it
    if not embedding_space_path:
        logger.warning(f"⚠️  Foundation model {foundation_model_id} does not have embedding_space path set in session")
        logger.info(f"🔍 Searching for embedding space file for foundation_model_id: {foundation_model_id}")
        
        # Search in common locations
        search_paths = []
        
        # Search patterns - include best_model_package and .pth files
        es_file_patterns = [
            "embedded_space.pickle",
            "embedding_space.pickle", 
            "best_model.pickle",
            "best_model_package/best_model.pickle",
            "best_model_package/embedded_space.pickle",
            "best_model_package/embedding_space.pickle"
        ]
        
        def search_in_train_es(base_dir: Path, search_paths: list):
            """Search for embedding space files within train_es directories only (recursive within train_es is OK)."""
            if not base_dir.exists():
                return
            
            # Find all train_es directories (recursive search)
            for train_es_dir in base_dir.rglob("train_es"):
                # Only search within train_es, not outside
                if not train_es_dir.is_dir():
                    continue
                
                logger.info(f"   Found train_es directory: {train_es_dir}")
                
                # Search for pickle files within train_es (recursive within train_es is OK)
                for es_file in es_file_patterns:
                    candidate = train_es_dir / es_file
                    logger.info(f"      Checking: {candidate} (exists: {candidate.exists()})")
                    if candidate.exists():
                        search_paths.append(candidate)
                        logger.info(f"      ✅ Found: {candidate}")
                
                # Also check for any .pth files within train_es directories (recursive)
                for pth_file in train_es_dir.rglob("*.pth"):
                    logger.info(f"      Found .pth file: {pth_file}")
                    # Prefer best_model.pth or model files
                    if "best" in pth_file.name.lower() or "model" in pth_file.name.lower():
                        search_paths.append(pth_file)
                        logger.info(f"      ✅ Added .pth file: {pth_file}")
                
                # Recursive search for pickle files within train_es only
                for pickle_file in train_es_dir.rglob("*.pickle"):
                    if pickle_file.name in ["embedded_space.pickle", "embedding_space.pickle", "best_model.pickle"]:
                        logger.info(f"      Found pickle file in train_es: {pickle_file}")
                        if pickle_file not in search_paths:
                            search_paths.append(pickle_file)
                            logger.info(f"      ✅ Added: {pickle_file}")
        
        # 1. Try in featrix_output directories (only within train_es)
        featrix_output_dir = Path("/sphere/app/featrix_output")
        if featrix_output_dir.exists():
            session_es_dir = featrix_output_dir / foundation_model_id
            logger.info(f"   Checking featrix_output: {session_es_dir} (exists: {session_es_dir.exists()})")
            if session_es_dir.exists():
                search_in_train_es(session_es_dir, search_paths)
        
        # 2. Try in published directory (only within train_es)
        published_dir = Path("/sphere/published") / foundation_model_id
        logger.info(f"   Checking published: {published_dir} (exists: {published_dir.exists()})")
        if published_dir.exists():
            search_in_train_es(published_dir, search_paths)
        
        # 3. Check session directory structure (only within train_es)
        session_dir = config.session_dir
        if session_dir:
            session_path = Path(session_dir) / foundation_model_id
            logger.info(f"   Checking session_dir: {session_path} (exists: {session_path.exists()})")
            if session_path.exists():
                search_in_train_es(session_path, search_paths)
        
        logger.info(f"   Total search paths found: {len(search_paths)}")
        for idx, path in enumerate(search_paths):
            logger.info(f"      [{idx+1}] {path}")
        
        # Use first found path
        if search_paths:
            embedding_space_path = str(search_paths[0].resolve())
            logger.info(f"✅ Found embedding space by searching: {embedding_space_path}")
            
            # Calculate MD5 hash and get ctime
            es_path = Path(embedding_space_path)
            if es_path.exists():
                # Calculate MD5
                md5_hash = hashlib.md5()
                with open(es_path, 'rb') as f:
                    for chunk in iter(lambda: f.read(4096), b""):
                        md5_hash.update(chunk)
                md5_hex = md5_hash.hexdigest()
                
                # Get creation time
                ctime = os.path.getctime(es_path)
                ctime_dt = datetime.fromtimestamp(ctime)
                
                logger.info(f"   Embedding space file info:")
                logger.info(f"      Path: {embedding_space_path}")
                logger.info(f"      MD5: {md5_hex}")
                logger.info(f"      Created: {ctime_dt.isoformat()}")
                
                # Store for predictor session metadata
                embedding_space_md5 = md5_hex
                embedding_space_ctime = ctime_dt.isoformat()
                
                # Update foundation session with this info
                foundation_session["embedding_space"] = embedding_space_path
                foundation_session["embedding_space_md5"] = md5_hex
                foundation_session["embedding_space_ctime"] = ctime_dt.isoformat()
                save_session(foundation_model_id, foundation_session, exist_ok=True)
                logger.info(f"   ✅ Updated foundation session with embedding space path, MD5, and ctime")
        else:
            logger.error(f"❌ Could not find embedding space file after searching:")
            logger.error(f"   Checked featrix_output: {featrix_output_dir / foundation_model_id}")
            logger.error(f"   Checked published: {Path('/sphere/published') / foundation_model_id}")
            if session_dir:
                logger.error(f"   Checked session_dir: {Path(session_dir) / foundation_model_id}")
            raise ValueError(f"Foundation model {foundation_model_id} does not have an embedding_space path set and could not be found by searching")
    
    if strings_cache_path:
        strings_cache_path = str(Path(strings_cache_path).resolve())
        if not Path(strings_cache_path).exists():
            logger.warning(f"⚠️  Foundation strings cache not found: {strings_cache_path} (will continue without it)")
            strings_cache_path = None
        else:
            logger.info(f"✅ Found foundation strings cache: {strings_cache_path}")
    else:
        logger.info(f"ℹ️  Foundation model has no strings cache (this is OK)")
    
    # Validate target_spec
    if "target_column" not in target_spec:
        raise ValueError("target_column is required in target_spec")
    if "target_column_type" not in target_spec:
        raise ValueError("target_column_type is required in target_spec")
    
    # Use foundation model's input data if not provided
    if input_filename is None:
        input_filename = foundation_session.get("input_data")
        logger.info(f"📁 Using foundation model's input data: {input_filename}")
    else:
        # Clean up input_filename: if it's an absolute path, extract just the filename
        # The file should already be in config.data_dir or uploaded separately
        input_path = Path(input_filename)
        
        if input_path.is_absolute():
            # Extract just the filename - assume it's been uploaded to data_dir
            input_filename = input_path.name
            logger.info(f"📁 Extracted filename from absolute path: {input_filename}")
            logger.info(f"   Original path was: {input_path}")
            logger.info(f"   Will look for file in: {config.data_dir / input_filename}")
        else:
            logger.info(f"📁 Using provided input data (relative path): {input_filename}")
    
    # Create job plan: ONLY create_structured_data and train_single_predictor
    # CRITICAL: train_es, train_knn, and build_es_projections are NOT in the plan
    # They will be skipped automatically by step_session because embedding_space already exists
    job_plan = [
        dict(
            job_type="create_structured_data",
            spec={},
            job_id=None,
        ),
        # NOTE: train_es, train_knn, build_es_projections are EXPLICITLY NOT included
        # The embedding_space path is already set above, so step_session will skip any ES jobs
        dict(
            job_type="train_single_predictor",
            spec=target_spec,
            job_id=None,
            predictor_index=0,
        ),
    ]
    
    logger.info(f"📋 Job plan created with {len(job_plan)} jobs (NO train_es job - reusing foundation ES)")
    
    # Build metadata for the predictor session
    metadata = {
        "foundation_model_id": foundation_model_id,  # Track which foundation model was used
    }
    
    # Add embedding space metadata if available
    if embedding_space_path:
        metadata["foundation_embedding_space_path"] = embedding_space_path
        if embedding_space_md5:
            metadata["foundation_embedding_space_md5"] = embedding_space_md5
        if embedding_space_ctime:
            metadata["foundation_embedding_space_ctime"] = embedding_space_ctime
    
    session_doc = dict(
        created_at=created_at,
        session_type="predictor",
        session_id=session_id,
        name=name,
        status=SessionStatus.READY,
        job_plan=job_plan,
        # CRITICAL: Use foundation model's embedding space and strings cache (already resolved to absolute paths)
        embedding_space=embedding_space_path,  # Absolute path to existing ES - NO new ES will be created
        strings_cache=strings_cache_path,  # Absolute path to existing cache
        # Paths to artifacts that are shared across jobs
        input_data=input_filename,
        sqlite_db=None,  # Will be created by create_structured_data
        vector_db=None,  # Not needed for predictor training
        projections=None,  # Not needed for predictor training
        preview_png=None,  # Not needed for predictor training
        single_predictors=[None],  # Will be populated when training completes
        training_metrics=[None],  # Will be populated when training completes
        foundation_model_id=foundation_model_id,  # Track which foundation model was used
        metadata=metadata,  # Embedding space path, MD5, and ctime metadata
    )
    
    logger.info(f"✅ Foundation model session created: {session_id}")
    logger.info(f"   Reusing ES from: {foundation_model_id}")
    logger.info(f"   ES path: {embedding_space_path}")
    logger.info(f"   Target column: {target_spec.get('target_column')}")
    
    return session_doc


class NodeUpgradingException(Exception):
    """Exception raised when a node is currently upgrading and cannot accept training requests."""
    pass


def create_embedding_space_session(name: str, s3_training_path: str, s3_validation_path: str, session_id: str = None, user_metadata: dict = None):
    """Create a session specifically for training an embedding space from S3 datasets."""
    # Check if node is upgrading
    upgrade_lock_file = Path("/tmp/auto-upgrade.lock")
    if upgrade_lock_file.exists():
        hostname = socket.gethostname()
        raise NodeUpgradingException(f"Node {hostname} is currently upgrading and cannot accept training requests. Please try again in a few minutes or use a different node.")

    # Ensure all necessary directories exist before creating sessions
    ensure_directories_exist()

    created_at = datetime.now(tz=ZoneInfo("America/New_York"))

    if session_id is None:
        unique_string = str(uuid4())[:6]
        session_timestamp = created_at.strftime('%Y%m%d-%H%M%S')
        # New format: {uuid}-{yyyy}{mm}{dd}-{hh}{mm}{ss}
        session_id = f"{unique_string}-{session_timestamp}"

    # Create job plan for embedding space training
    job_plan = [
        {
            "job_type": "create_structured_data",
            "spec": {
                "s3_training_dataset": s3_training_path,
                "s3_validation_dataset": s3_validation_path,
                "name": name
            },
            "job_id": None,
        },
        {
            "job_type": "train_es",
            "spec": {
                "user_metadata": user_metadata,
            },
            "job_id": None,
        },
        {
            "job_type": "train_knn",
            "spec": {},
            "job_id": None,
        },
        {
            "job_type": "build_es_projections",
            "spec": {},
            "job_id": None,
        }
    ]

    # Create session directory and data path
    # New structure: es_train/{session_id}/
    es_train_dir = config.data_dir / "es_train"
    es_train_dir.mkdir(parents=True, exist_ok=True)
    session_dir = es_train_dir / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    
    # For embedding space sessions, the data will be downloaded by the first job
    input_data_path = session_dir / f"{name}_training_data.csv"

    session_doc = {
        "session_id": session_id,
        "session_type": "embedding_space",
        "name": name,
        "status": SessionStatus.READY,
        "created_at": created_at,
        "input_data": str(input_data_path),
        "s3_training_dataset": s3_training_path,
        "s3_validation_dataset": s3_validation_path,
        "job_plan": job_plan,
        "projected_points": str(session_dir / "projection.json"),
        "preview_png": str(session_dir / "preview.png"),
        "embedding_space": str(session_dir / "embedding_space.json"),
        "sqlite_db": str(session_dir / "embedding_space.db"),
        "vector_db": str(session_dir / "vector_db.lance"),
        "projections": str(session_dir / "projections.json"),
        "strings_cache": str(session_dir / "strings_cache.pkl"),
        "user_metadata": user_metadata,  # User metadata for identification
    }

    return session_doc


def create_fine_tune_embedding_space_session(
    name: str,
    parent_session_id: str | None = None,
    parent_embedding_space_path: str | None = None,
    s3_training_path: str = None,
    s3_validation_path: str = None,
    session_id: str = None,
    user_metadata: dict = None
):
    """Create a session for fine-tuning an existing embedding space on new data.
    
    Args:
        name: Name for the fine-tuned embedding space
        parent_session_id: Session ID of the parent embedding space (optional)
        parent_embedding_space_path: Direct path to parent embedding space pickle file (optional)
        s3_training_path: S3 URL for new training dataset
        s3_validation_path: S3 URL for new validation dataset
        session_id: Optional session ID (will be generated if not provided)
    
    Returns:
        Session document with fine-tuning job plan
        
    The number of epochs is automatically calculated as: original_epochs / F
    where F = len(new_dataset) / len(old_dataset)
    """    
    logger.info(f"🔧 Creating fine-tuning session for embedding space: {name}")
    
    # Check if node is upgrading
    upgrade_lock_file = Path("/tmp/auto-upgrade.lock")
    if upgrade_lock_file.exists():
        hostname = socket.gethostname()
        raise NodeUpgradingException(f"Node {hostname} is currently upgrading and cannot accept training requests. Please try again in a few minutes or use a different node.")
    
    # Ensure all necessary directories exist
    ensure_directories_exist()
    
    created_at = datetime.now(tz=ZoneInfo("America/New_York"))
    
    if session_id is None:
        unique_string = str(uuid4())[:6]
        session_timestamp = created_at.strftime('%Y%m%d-%H%M%S')
        session_id = f"{unique_string}-{session_timestamp}"
    
    # Load parent embedding space to get original training info
    parent_es_path = None
    if parent_session_id:
        logger.info(f"📂 Loading parent embedding space from session: {parent_session_id}")
        parent_session = load_session(parent_session_id)
        if not parent_session:
            raise ValueError(f"Parent session {parent_session_id} not found")
        parent_es_path = parent_session.get("embedding_space")
        if not parent_es_path or not Path(parent_es_path).exists():
            raise FileNotFoundError(f"Parent embedding space not found at: {parent_es_path}")
    elif parent_embedding_space_path:
        parent_es_path = parent_embedding_space_path
        if not Path(parent_es_path).exists():
            raise FileNotFoundError(f"Parent embedding space not found at: {parent_embedding_space_path}")
    
    logger.info(f"📂 Loading parent embedding space from: {parent_es_path}")
    
    # Load the parent embedding space
    from lib.utils import load_embedded_space
    parent_es = load_embedded_space(parent_es_path, force_cpu=True)
    
    if not parent_es:
        raise ValueError(f"Failed to load parent embedding space from: {parent_es_path}")
    
    # Get original dataset size and epoch count
    original_train_size = len(parent_es.train_input_data) if hasattr(parent_es, 'train_input_data') else 0
    original_epochs = parent_es.n_epochs if hasattr(parent_es, 'n_epochs') and parent_es.n_epochs else 0
    
    # If n_epochs is not set, try to get from training_info
    if not original_epochs:
        training_info = getattr(parent_es, 'training_info', {})
        progress_info = training_info.get('progress_info', {})
        original_epochs = progress_info.get('epoch_total', 0)
        if not original_epochs:
            # Try to get from loss_history length
            loss_history = progress_info.get('loss_history', [])
            if loss_history:
                original_epochs = len(loss_history)
    
    logger.info(f"📊 Parent embedding space info:")
    logger.info(f"   Original training dataset size: {original_train_size} rows")
    logger.info(f"   Original epochs: {original_epochs}")
    
    if original_epochs == 0:
        logger.warning(f"⚠️  Could not determine original epoch count, defaulting to 100")
        original_epochs = 100
    
    # Download new datasets to get their sizes
    logger.info(f"📥 Downloading new datasets to calculate size...")
    
    s3_client = boto3.client('s3')
    
    def download_and_count_rows(s3_url: str) -> int:
        """Download S3 file temporarily and count rows."""
        import tempfile
        import os
        
        # Parse S3 URL
        if not s3_url.startswith('s3://'):
            raise ValueError(f"Invalid S3 URL: {s3_url}")
        
        s3_path = s3_url[5:]  # Remove 's3://'
        bucket, key = s3_path.split('/', 1)
        
        # Download to temp file
        with tempfile.NamedTemporaryFile(mode='w+b', suffix='.csv', delete=False) as tmp_file:
            tmp_path = tmp_file.name
            try:
                s3_client.download_file(bucket, key, tmp_path)
                
                # Count rows (read CSV)
                df = pd.read_csv(tmp_path, nrows=0)  # Just get header
                # Count non-header rows by reading file
                with open(tmp_path, 'r') as f:
                    row_count = sum(1 for line in f) - 1  # Subtract header
                return row_count
            finally:
                if os.path.exists(tmp_path):
                    os.unlink(tmp_path)
    
    try:
        new_train_size = download_and_count_rows(s3_training_path)
        new_val_size = download_and_count_rows(s3_validation_path)
        new_total_size = new_train_size + new_val_size
    except Exception as e:
        logger.warning(f"⚠️  Could not download datasets to count rows: {e}")
        logger.warning(f"   Will use default F=1.0 (same dataset size)")
        new_total_size = original_train_size
        new_train_size = int(original_train_size * 0.8)
        new_val_size = int(original_train_size * 0.2)
    
    logger.info(f"📊 New dataset info:")
    logger.info(f"   New training dataset size: {new_train_size} rows")
    logger.info(f"   New validation dataset size: {new_val_size} rows")
    logger.info(f"   New total size: {new_total_size} rows")
    
    # Calculate F = len(new_dataset) / len(old_dataset)
    F = new_total_size / original_train_size if original_train_size > 0 else 1.0
    logger.info(f"📐 Dataset size ratio F = {new_total_size} / {original_train_size} = {F:.4f}")
    
    # Calculate new epochs = original_epochs / F
    new_epochs = int(original_epochs / F) if F > 0 else original_epochs
    if new_epochs < 1:
        new_epochs = 1
    elif new_epochs > 1000:
        logger.warning(f"⚠️  Calculated epochs ({new_epochs}) is very high, capping at 1000")
        new_epochs = 1000
    
    logger.info(f"🎯 Fine-tuning will use {new_epochs} epochs (original: {original_epochs}, F: {F:.4f})")
    
    # Create job plan for fine-tuning
    job_plan = [
        {
            "job_type": "create_structured_data",
            "spec": {
                "s3_training_dataset": s3_training_path,
                "s3_validation_dataset": s3_validation_path,
                "name": name
            },
            "job_id": None,
        },
        {
            "job_type": "train_es",
            "spec": {
                "epochs": new_epochs,
                "parent_embedding_space_path": str(Path(parent_es_path).resolve()),
                "fine_tune": True,
                "user_metadata": user_metadata,
            },
            "job_id": None,
        },
        {
            "job_type": "train_knn",
            "spec": {},
            "job_id": None,
        },
        {
            "job_type": "build_es_projections",
            "spec": {},
            "job_id": None,
        }
    ]
    
    # Create session directory
    es_train_dir = config.data_dir / "es_train"
    es_train_dir.mkdir(parents=True, exist_ok=True)
    session_dir = es_train_dir / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    
    input_data_path = session_dir / f"{name}_training_data.csv"
    
    session_doc = {
        "session_id": session_id,
        "session_type": "embedding_space_finetune",
        "name": name,
        "status": SessionStatus.READY,
        "created_at": created_at,
        "input_data": str(input_data_path),
        "s3_training_dataset": s3_training_path,
        "s3_validation_dataset": s3_validation_path,
        "parent_session_id": parent_session_id,
        "parent_embedding_space_path": str(Path(parent_es_path).resolve()),
        "fine_tune_info": {
            "original_train_size": original_train_size,
            "original_epochs": original_epochs,
            "new_train_size": new_train_size,
            "new_val_size": new_val_size,
            "new_total_size": new_total_size,
            "F": F,
            "calculated_epochs": new_epochs,
        },
        "job_plan": job_plan,
        "projected_points": str(session_dir / "projection.json"),
        "preview_png": str(session_dir / "preview.png"),
        "embedding_space": str(session_dir / "embedding_space.json"),
        "sqlite_db": str(session_dir / "embedding_space.db"),
        "vector_db": str(session_dir / "vector_db.lance"),
        "projections": str(session_dir / "projections.json"),
        "strings_cache": str(session_dir / "strings_cache.pkl"),
        "user_metadata": user_metadata,  # User metadata for identification
    }
    
    logger.info(f"✅ Created fine-tuning session: {session_id}")
    
    return session_doc


def queue_es_training_for_json_column(col_name: str, json_values: List[Any], current_host: str = None, json_column_sessions: dict = None, parent_es_name: str = None, parent_session_id: str = None):
    """
    Queue an embedding space training job for a JSON column if one doesn't exist.
    Checks local ES first, then API, then queues training locally if needed.
    Detects dependencies on other JSON column ES's.
    
    Args:
        col_name: Name of the JSON column
        json_values: List of all JSON values from the column
        current_host: Current compute host name (for targeting)
        json_column_sessions: Map of other JSON column names to their session IDs (for dependency detection)
        parent_es_name: Name of the parent ES session (for naming: <parent_es_name>-<col_name>)
        parent_session_id: Parent session ID (deprecated - child ES sessions are now top-level)
        
    Returns:
        Session ID if job was queued, None otherwise
    """
    logger.info(f"🔍 Checking for existing ES for JSON column '{col_name}'...")
    
    # Extract schema fields from JSON values
    schema_fields = set()
    parsed_values = []
    
    for value in json_values[:1000]:  # Sample first 1000 for schema extraction
        try:
            # Parse JSON value
            if isinstance(value, str):
                value = value.strip()
                if value.startswith('{'):
                    try:
                        parsed = json.loads(value)
                    except:
                        try:
                            parsed = ast.literal_eval(value)
                        except:
                            continue
                elif value.startswith('['):
                    try:
                        parsed = json.loads(value)
                        if isinstance(parsed, list) and len(parsed) > 0:
                            parsed = parsed[0]
                    except:
                        try:
                            parsed = ast.literal_eval(value)
                            if isinstance(parsed, list) and len(parsed) > 0:
                                parsed = parsed[0]
                        except:
                            continue
                else:
                    continue
            elif isinstance(value, dict):
                parsed = value
            elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict):
                parsed = value[0]
            else:
                continue
            
            if isinstance(parsed, dict):
                schema_fields.update(parsed.keys())
                parsed_values.append(parsed)
        except Exception:
            continue
    
    if not schema_fields:
        logger.warning(f"⚠️ Could not extract schema from JSON column '{col_name}'")
        return None
    
    schema_fields = sorted(list(schema_fields))
    schema_fields_set = set(schema_fields)
    logger.info(f"📋 Extracted schema fields for '{col_name}': {schema_fields}")
    
    # CRITICAL: Do NOT check for existing sessions before acquiring lock
    # This prevents race conditions where multiple calls all see "no session" and create duplicates
    # Lock will be acquired first, THEN we check for existing sessions
    sessions_dir = config.session_dir
    
    # First, check local ES sessions
    try:
        sessions_dir = config.session_dir
        if sessions_dir.exists():
            logger.info(f"🔍 Checking local ES sessions...")
            for session_file in sessions_dir.glob("*.session"):
                try:
                    session = load_session(session_file.stem)
                    if session.get("session_type") != "embedding_space":
                        continue
                    
                    embedding_space_path = session.get("embedding_space")
                    if not embedding_space_path or not Path(embedding_space_path).exists():
                        continue
                    
                    # Load ES to check schema
                    with open(embedding_space_path, 'rb') as f:
                        es = pickle.load(f)
                    
                    es_columns = set(es.get_column_names())
                    
                    # Check if schemas match (80% overlap)
                    matched_fields = schema_fields_set.intersection(es_columns)
                    match_score = len(matched_fields) / len(schema_fields_set) if schema_fields_set else 0
                    
                    if match_score >= 0.8:
                        logger.info(f"✅ Found local ES matching '{col_name}': {embedding_space_path} (match: {match_score:.2f})")
                        return None  # ES already exists, no need to train
                except (ValueError, FileNotFoundError, json.JSONDecodeError) as e:
                    # Corrupted or missing session file - skip silently (don't let one bad file break the system)
                    logger.debug(f"Skipping corrupted/missing session {session_file.stem} (checking ES for column '{col_name}'): {e}")
                    continue
                except Exception as e:
                    # Other errors - log at debug level and continue
                    logger.debug(f"Failed to check session {session_file.stem} (checking ES for column '{col_name}'): {e}")
                    continue
    except Exception as e:
        logger.warning(f"⚠️ Failed to check local ES: {e}")
    
    # Check via API for remote ES (optional - nice to know if one exists elsewhere)
    try:
        from config import config as app_config
        api_base = getattr(app_config, 'api_base_url', 'http://localhost:8000')
        if not api_base.startswith('http'):
            api_base = f"http://{api_base}"
        
        schema_fields_str = ','.join(schema_fields)
        url = f"{api_base}/api-sphere/json-encoders"
        params = {"schema_fields": schema_fields_str}
        
        logger.info(f"🔍 Querying API for existing ES: {url}")
        response = requests.get(url, params=params, timeout=10)
        
        if response.status_code == 200:
            result = response.json()
            matched_es = result.get("matched_es")
            
            if matched_es:
                es_path = matched_es.get("embedding_space_path")
                logger.info(f"ℹ️ Found remote ES for '{col_name}': {es_path} (but training locally)")
                # Note: We could return this, but user wants to train locally
    except Exception as e:
        logger.debug(f"API check failed (non-critical): {e}")
    
    # Detect dependencies: check if any schema fields reference other JSON columns
    # Store as {col_name: session_id} mapping
    required_child_es_mapping = {}
    if json_column_sessions:
        for field_name in schema_fields:
            # Check if this field name matches another JSON column name
            if field_name in json_column_sessions:
                dep_session_id = json_column_sessions[field_name]
                required_child_es_mapping[field_name] = dep_session_id
                logger.info(f"🔗 Detected dependency: '{col_name}' requires ES from column '{field_name}' (session: {dep_session_id})")
    
    # No existing ES found - train locally
    logger.info(f"🚀 No local ES found for '{col_name}', training new one locally...")
    if required_child_es_mapping:
        logger.info(f"   Dependencies: {len(required_child_es_mapping)} required columns: {list(required_child_es_mapping.keys())}")
    
    # Prepare training data: convert JSON values to records
    all_parsed = []
    for value in json_values:
        try:
            if isinstance(value, str):
                value = value.strip()
                if value.startswith('{'):
                    try:
                        parsed = json.loads(value)
                    except:
                        try:
                            parsed = ast.literal_eval(value)
                        except:
                            continue
                elif value.startswith('['):
                    try:
                        parsed = json.loads(value)
                        if isinstance(parsed, list):
                            all_parsed.extend(parsed if all(isinstance(x, dict) for x in parsed) else [])
                            continue
                    except:
                        try:
                            parsed = ast.literal_eval(value)
                            if isinstance(parsed, list):
                                all_parsed.extend(parsed if all(isinstance(x, dict) for x in parsed) else [])
                                continue
                        except:
                            continue
                else:
                    continue
            elif isinstance(value, dict):
                parsed = value
            elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict):
                all_parsed.extend(value)
                continue
            else:
                continue
            
            if isinstance(parsed, dict):
                all_parsed.append(parsed)
        except Exception:
            continue
    
    if len(all_parsed) < 10:
        logger.warning(f"⚠️ Not enough JSON records for training ES (got {len(all_parsed)}, need at least 10)")
        return None
    
    # For small datasets (<200 rows), put ALL rows in BOTH training and validation sets
    # For larger datasets, split 80/20
    if len(all_parsed) < 200:
        logger.info(f"📊 Small dataset ({len(all_parsed)} rows): Using all rows for both training and validation")
        train_records = all_parsed
        val_records = all_parsed.copy()  # Copy so both sets have all rows
    else:
        # Split into train/validation (80/20)
        train_size = int(len(all_parsed) * 0.8)
        train_records = all_parsed[:train_size]
        val_records = all_parsed[train_size:]
        logger.info(f"📊 Dataset split: {len(train_records)} train, {len(val_records)} val (80/20)")
    
    # Save to JSONL files (JSON Lines - one JSON object per line)
    # Child ES sessions are top-level - use their own session_id for directory
    # Add UUID to make session ID unique
    unique_id = str(uuid4())
    es_training_session_id = f"json-es-{col_name}-{unique_id}"
    es_name = f"json-es-{col_name}"
    
    # Create training files in the child ES session's own directory
    child_session_output_dir = config.output_dir / es_training_session_id
    temp_dir = child_session_output_dir / "json_es_training"
    temp_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"📁 Creating JSON ES training files in child session directory: {temp_dir}")
    
    train_jsonl = temp_dir / f"{es_name}_training.jsonl"
    val_jsonl = temp_dir / f"{es_name}_validation.jsonl"
    
    # Write JSONL files (one JSON object per line)
    with open(train_jsonl, 'w') as f:
        for record in train_records:
            f.write(json.dumps(record) + '\n')
    
    with open(val_jsonl, 'w') as f:
        for record in val_records:
            f.write(json.dumps(record) + '\n')
    
    logger.info(f"📊 Created training data: {len(train_records)} train, {len(val_records)} val JSON records")
    
    # Train ES locally using the training functions directly
    try:
        # CRITICAL: Use a lock file to prevent race conditions when creating sessions
        # Check for existing session again WITH lock to prevent duplicates
        lock_file_path = config.session_dir / f"json-es-{col_name}.create_lock"
        
        # Try to acquire lock (non-blocking)
        try:
            lock_file = open(lock_file_path, 'x')  # 'x' mode creates file exclusively, fails if exists
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            logger.debug(f"🔒 Acquired lock for creating session for column '{col_name}'")
        except FileExistsError:
            # Another process is creating a session for this column - wait and check
            logger.info(f"⏳ Another process is creating session for '{col_name}', waiting...")
            lock_file = open(lock_file_path, 'r')
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH)  # Shared lock to wait
            # Lock released, check if session was created
            # Reuse any existing session - train_es is part of the same session
            for session_file in sessions_dir.glob(f"json-es-{col_name}-*.session"):
                try:
                    existing_session = load_session(session_file.stem)
                    session_name = existing_session.get("name", "")
                    if session_name == f"json-es-{col_name}":
                        logger.info(f"✅ Session for column '{col_name}' was created by another process: {session_file.stem}")
                        logger.info(f"   Reusing existing session - train_es is part of the same session")
                        lock_file.close()
                        # CRITICAL: Step the session to ensure jobs are queued
                        try:
                            step_session(session_id=session_file.stem)
                            logger.info(f"✅ Stepped existing session to ensure jobs are queued")
                        except Exception as step_e:
                            logger.warning(f"⚠️ Failed to step existing session {session_file.stem}: {step_e}")
                            # Continue anyway - session exists, jobs may already be queued
                        return session_file.stem
                except Exception:
                    continue
            lock_file.close()
            # Still no session, try again with exclusive lock
            lock_file = open(lock_file_path, 'w')
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        
        try:
            # Double-check with lock held (prevents race condition)
            # Reuse any existing session - train_es is part of the same session as create_structured_data
            for session_file in sessions_dir.glob(f"json-es-{col_name}-*.session"):
                try:
                    existing_session = load_session(session_file.stem)
                    session_name = existing_session.get("name", "")
                    if session_name == f"json-es-{col_name}":
                        logger.info(f"✅ Session for column '{col_name}' already exists (checked with lock): {session_file.stem}")
                        logger.info(f"   Reusing existing session - train_es is part of the same session")
                        # CRITICAL: Step the session to ensure jobs are queued
                        try:
                            step_session(session_id=session_file.stem)
                            logger.info(f"✅ Stepped existing session to ensure jobs are queued")
                        except Exception as step_e:
                            logger.warning(f"⚠️ Failed to step existing session {session_file.stem}: {step_e}")
                            # Continue anyway - session exists, jobs may already be queued
                        return session_file.stem
                except Exception:
                    continue
            
            # Create session for ES training
            # Use format: json-es-<column_name>-<uuid> to make it unique
            unique_id = str(uuid4())
            es_training_session_id = f"json-es-{col_name}-{unique_id}"
            logger.info(f"🔍 queue_es_training_for_json_column: Created ES training session_id={es_training_session_id} for column '{col_name}'")
            
            # Create session directory and copy JSONL files there
            # New structure: es_train/{es_training_session_id}/
            es_train_dir = config.data_dir / "es_train"
            es_train_dir.mkdir(parents=True, exist_ok=True)
            session_dir = es_train_dir / es_training_session_id
            session_dir.mkdir(parents=True, exist_ok=True)
            
            # Copy JSONL files to session directory
            session_train_jsonl = session_dir / f"{es_name}_training.jsonl"
            session_val_jsonl = session_dir / f"{es_name}_validation.jsonl"
            shutil.copy(train_jsonl, session_train_jsonl)
            shutil.copy(val_jsonl, session_val_jsonl)
            
            logger.info(f"📁 Copied training JSONL files to session directory: {session_dir}")
            
            # Create session - use local paths (they'll be handled as local files, not S3)
            # The job handler checks for s3:// prefix, so local paths work fine
            logger.debug(f"🔍 queue_es_training_for_json_column: Creating embedding space session with session_id={es_training_session_id}")
            session = create_embedding_space_session(
                name=es_name,
                s3_training_path=str(session_train_jsonl.resolve()),  # Local JSONL path - job handler will use it directly
                s3_validation_path=str(session_val_jsonl.resolve()),  # Local JSONL path
                session_id=es_training_session_id
            )
            logger.debug(f"✅ queue_es_training_for_json_column: Created session document with session_id={session.get('session_id')}")
            
            # Update session to use local files (not S3)
            # The job handler will copy input_data to the job working directory
            session['input_data'] = str(session_train_jsonl.resolve())
            
            # Update job spec to handle local files (not S3)
            # Modify the first job's spec to use local paths
            if session.get('job_plan') and len(session['job_plan']) > 0:
                first_job = session['job_plan'][0]
                if first_job.get('job_type') == 'create_structured_data':
                    # Update spec to use local paths - the handler checks for s3:// prefix
                    first_job['spec']['s3_training_dataset'] = str(session_train_jsonl.resolve())
                    first_job['spec']['s3_validation_dataset'] = str(session_val_jsonl.resolve())
            
            # Store required column dependencies in session as {col_name: session_id} mapping
            if required_child_es_mapping:
                session['required_child_es_mapping'] = required_child_es_mapping
                logger.info(f"📋 Stored {len(required_child_es_mapping)} column dependencies in session: {required_child_es_mapping}")
            
            # Child ES sessions are top-level independent sessions - no parent_session_id needed
            # They use their own session_id for output directories
            
            # Save session
            logger.debug(f"🔍 queue_es_training_for_json_column: Saving session with session_id={es_training_session_id}")
            save_session(session_id=es_training_session_id, session_doc=session, exist_ok=False)
        finally:
            # Always release lock
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
            lock_file.close()
            # Remove lock file
            try:
                lock_file_path.unlink()
            except Exception:
                pass
        
        # Verify session was saved correctly
        verify_session = load_session(es_training_session_id)
        logger.debug(f"✅ Verified session {es_training_session_id} was saved correctly")
        
        logger.info(f"✅ Created ES session: {es_training_session_id}")
        
        # Queue the jobs - CRITICAL: Use es_training_session_id explicitly
        logger.info(f"🚀 Queueing ES training jobs for session {es_training_session_id}...")
        logger.debug(f"🔍 queue_es_training_for_json_column: About to call step_session with es_training_session_id={es_training_session_id}")
        try:
            step_session(session_id=es_training_session_id)
            logger.debug(f"✅ queue_es_training_for_json_column: step_session completed for es_training_session_id={es_training_session_id}")
        except ValueError as step_e:
            # Check if this is a queue-full error
            if "queue is full" in str(step_e) or "max" in str(step_e).lower():
                logger.warning(f"⚠️ Queue is full, cannot queue jobs for session {es_training_session_id} right now")
                logger.warning(f"   Session created successfully - jobs will be queued when queue has space")
                logger.warning(f"   Error: {step_e}")
                # Don't raise - session is created, jobs can be queued later
            else:
                # Other ValueError - re-raise
                logger.error(f"❌ step_session failed for {es_training_session_id}: {step_e}")
                logger.error(traceback.format_exc())
                raise
        except Exception as step_e:
            logger.error(f"❌ step_session failed for {es_training_session_id}: {step_e}")
            logger.error(traceback.format_exc())
            raise
        
        # Get job IDs from the session's job plan for logging
        session = load_session(es_training_session_id)
        job_plan = session.get("job_plan", [])
        job_ids = []
        for job_desc in job_plan:
            job_id = job_desc.get("job_id")
            job_type = job_desc.get("job_type")
            if job_id:
                job_ids.append(f"{job_type}:{job_id}")
                # Verify job file exists
                queue_name = "cpu_data_tasks" if job_type == "create_structured_data" else job_type
                job_file = config.queue_dir / queue_name / f"{job_id}.job"
                if job_file.exists():
                    logger.info(f"   ✅ Job file exists: {job_file}")
                else:
                    logger.warning(f"   ⚠️ Job file missing: {job_file}")
        
        if not job_ids:
            logger.error(f"❌ CRITICAL: No jobs were queued for session {es_training_session_id}!")
            logger.error(f"   This means step_session did not create any jobs")
            logger.error(f"   Session job_plan has {len(job_plan)} jobs")
            for i, job_desc in enumerate(job_plan):
                logger.error(f"     Job {i}: type={job_desc.get('job_type')}, job_id={job_desc.get('job_id')}")
            logger.error(f"   Session status: {session.get('status')}")
            logger.error(f"   This session will NOT start automatically - jobs must be manually queued")
            return None
        
        # Jobs are queued asynchronously - they will run when workers pick them up
        # We don't wait here because:
        # 1. This function is called from within a job execution, so waiting would block the job
        # 2. Waiting would cause logs to go to the wrong stdout.log file
        # 3. Jobs are designed to run asynchronously via the queue system
        logger.info(f"✅ ES training jobs queued for '{col_name}' (session: {es_training_session_id})")
        logger.info(f"   Job IDs: {', '.join(job_ids)}")
        logger.info(f"   ℹ️  Jobs will run asynchronously - ES will be available when training completes")
        
        logger.info(f"✅ queue_es_training_for_json_column: Returning es_training_session_id={es_training_session_id}")
        return es_training_session_id
        
    except Exception as e:
        logger.error(f"❌ Error training ES for '{col_name}': {e}")
        logger.error(traceback.format_exc())
        return None


def _wait_and_report_on_es_training(session_id: str, col_name: str, max_wait_time: int = 3600 * 4, check_interval: int = 30):
    """
    Wait for ES training job to complete and report on progress.
    
    Args:
        session_id: Session ID to wait for
        col_name: Column name being trained (for logging context)
        max_wait_time: Maximum time to wait in seconds (default: 4 hours)
        check_interval: How often to check status in seconds (default: 30s)
    """
    start_time = time.time()
    last_log_time = 0
    log_interval = 60  # Log every 60 seconds
    check_count = 0  # Track number of checks for adaptive interval
    
    # Initial delay to allow job to start and update status
    time.sleep(2)
    
    while time.time() - start_time < max_wait_time:
        try:
            # Reload session to get current status
            session = load_session(session_id)
            session_status = session.get("status")
            job_plan = session.get("job_plan", [])
            
            # Check each job in the plan
            jobs_status = {}
            for idx, job_desc in enumerate(job_plan):
                job_type = job_desc.get("job_type")
                job_id = job_desc.get("job_id")
                
                if job_id:
                    queue_name = "cpu_data_tasks" if job_type == "create_structured_data" else job_type
                    # Retry loading job file (handles race condition where file is being written)
                    job = None
                    for retry in range(3):
                        try:
                            job = load_job(queue_name, job_id)
                            break
                        except FileNotFoundError:
                            if retry < 2:  # Retry up to 2 times (3 total attempts)
                                time.sleep(0.1)  # Brief delay for filesystem sync
                                continue
                            # After retries, mark as not_found
                            jobs_status[job_type] = {"status": "NOT_FOUND", "job_id": job_id}
                            logger.debug(f"Job file {job_id} not found in {queue_name} after {retry+1} attempts")
                            break
                        except Exception as e:
                            logger.debug(f"Error loading job {job_id} from {queue_name}: {e}")
                            jobs_status[job_type] = {"status": "ERROR", "job_id": job_id}
                            break
                    
                    if job is not None:
                        job_status = job.get("status")
                        # Convert JobStatus enum to string
                        if hasattr(job_status, 'value'):
                            status_str = job_status.value.upper()
                        else:
                            status_str = str(job_status).upper()
                        
                        # Check if job is actually running by looking at working directory
                        # This helps detect when status hasn't been updated yet but job is running
                        is_actually_running = False
                        if status_str == "READY":
                            # Check if working directory exists (indicates job started)
                            try:
                                session_for_job = load_session(session_id)
                                if session_for_job:
                                    work_dir = config.output_dir / session_id / job_type / job_id
                                    if work_dir.exists() and (work_dir / "logs").exists():
                                        # Job has started - check if there's a log file or output
                                        log_file = work_dir / "logs" / "stdout.log"
                                        if log_file.exists() and log_file.stat().st_size > 0:
                                            is_actually_running = True
                                            status_str = "RUNNING"  # Override status
                            except Exception:
                                pass  # Ignore errors when checking working directory
                        
                        jobs_status[job_type] = {
                            "status": status_str,
                            "job_id": job_id,
                            "actually_running": is_actually_running
                        }
                else:
                    jobs_status[job_type] = {"status": "not_queued", "job_id": None}
            
            # Check if ES training is complete (train_es job is DONE or embedding_space exists)
            train_es_status = jobs_status.get("train_es", {}).get("status")
            embedding_space_path = session.get("embedding_space")
            es_exists = embedding_space_path and Path(embedding_space_path).exists()
            
            # Log progress periodically
            current_time = time.time()
            if current_time - last_log_time >= log_interval:
                elapsed_minutes = (current_time - start_time) / 60
                
                # Build status summary
                status_parts = []
                ready_jobs = []
                for job_type, info in jobs_status.items():
                    status_str = info["status"].upper()  # Normalize to uppercase for comparison
                    status_emoji = {
                        "DONE": "✅",
                        "RUNNING": "🔄",
                        "READY": "⏳",
                        "FAILED": "❌",
                        "NOT_QUEUED": "⏸️",
                        "NOT_FOUND": "❓",
                        "ERROR": "⚠️"
                    }.get(status_str, "❓")
                    status_parts.append(f"{status_emoji} {job_type}: {info['status']}")
                    if status_str == "READY":
                        ready_jobs.append((job_type, info.get("job_id")))
                
                logger.info(f"📊 ES training progress for '{col_name}' (session: {session_id}) - {elapsed_minutes:.1f} min elapsed")
                logger.info(f"   Session status: {session_status.value if hasattr(session_status, 'value') else session_status}")
                logger.info(f"   Jobs: {' | '.join(status_parts)}")
                
                # Check if jobs are stuck in READY state (no workers running)
                if ready_jobs and elapsed_minutes > 2:  # After 2 minutes, warn if still READY
                    logger.warning(f"   ⚠️  Jobs still READY after {elapsed_minutes:.1f} min - workers may not be running")
                    # Check if job files exist
                    for job_type, job_id in ready_jobs:
                        if job_id:
                            queue_name = "cpu_data_tasks" if job_type == "create_structured_data" else job_type
                            job_file = config.queue_dir / queue_name / f"{job_id}.job"
                            if job_file.exists():
                                logger.info(f"      ✅ Job file exists: {queue_name}/{job_id}.job")
                            else:
                                logger.error(f"      ❌ Job file missing: {queue_name}/{job_id}.job")
                
                if es_exists:
                    logger.info(f"   ✅ Embedding space exists: {embedding_space_path}")
                
                last_log_time = current_time
            
            # Adaptive check interval: check more frequently initially, then slow down
            check_count += 1
            if check_count <= 5:
                # First 5 checks: every 2 seconds
                time.sleep(2)
            elif check_count <= 10:
                # Next 5 checks: every 5 seconds
                time.sleep(5)
            else:
                # After that: use normal interval
                time.sleep(check_interval)
            
            # Check if training is complete
            if session_status == SessionStatus.DONE:
                logger.info(f"✅ ES training completed for '{col_name}' (session: {session_id})")
                if es_exists:
                    logger.info(f"   📁 Embedding space: {embedding_space_path}")
                return True
            
            # Check if train_es job is done (even if session status isn't updated yet)
            if train_es_status == "DONE" or (train_es_status == "RUNNING" and es_exists):
                logger.info(f"✅ ES training job completed for '{col_name}' (session: {session_id})")
                if es_exists:
                    logger.info(f"   📁 Embedding space: {embedding_space_path}")
                return True
            
            # Check if any critical job failed
            if train_es_status == "FAILED" and not es_exists:
                logger.error(f"❌ ES training job failed for '{col_name}' (session: {session_id})")
                logger.error(f"   Job ID: {jobs_status.get('train_es', {}).get('job_id', 'unknown')}")
                return False
            
            # Wait before next check
            time.sleep(check_interval)
            
        except Exception as e:
            logger.warning(f"⚠️ Error checking ES training status: {e}")
            time.sleep(check_interval)
    
    # Timeout
    elapsed_hours = (time.time() - start_time) / 3600
    logger.error(f"⏰ Timeout waiting {elapsed_hours:.1f} hours for ES training to complete for '{col_name}' (session: {session_id})")
    return False


def create_session(session_type: str, session_id: str = None, start: bool = True, input_filename: str | None = None, name: str = None, session_name_prefix: str = None, target_spec: dict = None, single_predictors: list = None, epochs: int = None, column_overrides: dict = None, string_list_delimiter: str = "|", movie_frame_interval: int = 3, weightwatcher_save_every: int = 5, important_columns_for_visualization: list = None, user_metadata: dict = None):

    # TODO: add different session types

    # Check if node is upgrading
    upgrade_lock_file = Path("/tmp/auto-upgrade.lock")
    if upgrade_lock_file.exists():
        hostname = socket.gethostname()
        raise NodeUpgradingException(f"Node {hostname} is currently upgrading and cannot accept training requests. Please try again in a few minutes or use a different node.")

    # Ensure all necessary directories exist before creating sessions
    ensure_directories_exist()

    created_at = datetime.now(tz=ZoneInfo("America/New_York"))

    # Sanitize session_name_prefix: replace slashes and dots with underscores
    if session_name_prefix:
        session_name_prefix = session_name_prefix.replace('/', '_').replace('.', '_')
        logger.debug(f"📋 Sanitized session_name_prefix: {session_name_prefix}")

    if session_id is None:
        # Generate full UUID
        full_uuid = str(uuid4())
        
        # If prefix is provided, format as <prefix>-<full-uuid>
        # For embedding_space sessions, always use new format: <uuid6>-<timestamp>
        # Otherwise, use the original format: <timestamp>_<uuid6>
        if session_name_prefix:
            session_id = f"{session_name_prefix}-{full_uuid}"
        elif session_type == "embedding_space":
            # Embedding space sessions use new format: {uuid}-{timestamp}
            unique_string = full_uuid[:6]
            session_timestamp = created_at.strftime('%Y%m%d-%H%M%S')
            session_id = f"{unique_string}-{session_timestamp}"
        else:
            unique_string = full_uuid[:6]
            session_timestamp = created_at.strftime('%Y%m%d-%H%M%S')
            session_id = "_".join([session_timestamp, unique_string])
    else:
        # session_id was provided - if prefix is also provided, ensure session_id uses it with UUID
        if session_name_prefix:
            # Prefix is provided - session_id MUST be in format {prefix}-{uuid}
            # If it's not, create a new one with the prefix
            if not session_id.startswith(f"{session_name_prefix}-"):
                full_uuid = str(uuid4())
                session_id = f"{session_name_prefix}-{full_uuid}"
                logger.info(f"📋 Created session_id with prefix and UUID: {session_id}")
            # If it does start with prefix, verify it has UUID-like suffix (contains dashes after prefix)
            elif len(session_id) <= len(session_name_prefix) + 1:
                # Too short - probably just the prefix, add UUID
                full_uuid = str(uuid4())
                session_id = f"{session_name_prefix}-{full_uuid}"
                logger.info(f"📋 Added UUID suffix to session_id: {session_id}")

    if session_type == "echo":
        session_doc = create_echo_session(
            session_id=session_id,
            created_at=created_at,
            input_filename=input_filename,
        )
    elif session_type == "sphere":
        session_doc = create_sphere_session(
            session_id=session_id,
            created_at=created_at,
            input_filename=input_filename,
            name=name,
            single_predictors=single_predictors,
            epochs=epochs,
            column_overrides=column_overrides,
            string_list_delimiter=string_list_delimiter,
            movie_frame_interval=movie_frame_interval,
            weightwatcher_save_every=weightwatcher_save_every,
            important_columns_for_visualization=important_columns_for_visualization,
            user_metadata=user_metadata,
        )
    elif session_type == "predictor":
        session_doc = create_predictor_session(
            session_id=session_id,
            created_at=created_at,
            input_filename=input_filename,
            target_spec=target_spec,
        )
    elif session_type == "embedding_space":
        session_doc = create_embedding_space_session(
            name=input_filename,
            s3_training_path=target_spec.get("training_dataset"),
            s3_validation_path=target_spec.get("validation_dataset"),
            session_id=session_id,
        )
    else:
        raise ValueError(f"Unsupported session type {session_type}")


    save_session(session_id=session_id, session_doc=session_doc, exist_ok=False)

    # Always step sessions to start processing
    step_session(session_id=session_id)

    return session_doc


def add_job_id_to_session(session_doc, job_type: str, job_id: str):
    job_plan = session_doc.get("job_plan")

    # For jobs that can have multiple instances (like train_single_predictor),
    # we need to find the first job of that type without a job_id
    for job_description in job_plan:
        if job_description.get("job_type") == job_type and job_description.get("job_id") is None:
            job_description["job_id"] = job_id
            break

    return session_doc

def clean_numpy_values(data):
    """
    Recursively clean NaN, Inf, and other non-JSON-serializable values from data.
    Converts them to None which is JSON serializable. Based on existing _clean_numpy_values 
    from the client code.
    
    Args:
        data: Data structure to clean (dict, list, or primitive)
        
    Returns:
        Cleaned data structure
    """
    import numpy as np
    
    if isinstance(data, dict):
        return {k: clean_numpy_values(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [clean_numpy_values(v) for v in data]
    elif isinstance(data, (float, np.floating)):
        if math.isnan(data) or math.isinf(data):
            return None
        return float(data)  # Convert numpy floats to Python floats
    elif isinstance(data, (int, np.integer)):
        return int(data)  # Convert numpy ints to Python ints
    elif isinstance(data, (bool, np.bool_)):
        return bool(data)  # Convert numpy bools to Python bools
    elif isinstance(data, np.ndarray):
        return clean_numpy_values(data.tolist())  # Convert arrays to lists
    elif data is None or isinstance(data, (str, bool)):
        return data
    else:
        # Handle other numpy types or unknown types
        try:
            # Try to convert to a basic Python type
            if hasattr(data, 'item'):  # numpy scalar
                value = data.item()
                if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
                    return None
                return value
            else:
                return data
        except:
            # If all else fails, convert to string
            return str(data)

def serialize_session(session_doc: dict):

    session_doc = deepcopy(session_doc)
        
    # Validation
    if session_doc.get("created_at") is None:
        raise ValueError("Session document must have a created_at timestamp")
    if session_doc.get("status") is None:
        raise ValueError("Session document must have a status")
    if session_doc.get("job_plan") is None or len(session_doc.get("job_plan")) == 0:
        raise ValueError("Session document must have a job plan")
    if session_doc.get("input_data") is None:
        raise ValueError("Session document must have a data field")
    if session_doc.get("session_id") is None:
        raise ValueError("Session document must have a session_id")
    if session_doc.get("session_type") is None:
        raise ValueError("Session document must have a session_type")


    session_doc['created_at'] = convert_to_iso(session_doc['created_at'])
    session_doc['status'] = session_doc['status'].value

    # Convert any numpy types to native Python types for JSON serialization
    serialized_session = clean_numpy_values(session_doc)

    return serialized_session

def save_session(session_id: str, session_doc: dict, exist_ok: bool = False):
    session_path = config.session_dir / f"{session_id}.session"
    lock_path = config.session_dir / f"{session_id}.lock"

    # If the session already exists and we're not overwriting, raise an error.
    if session_path.exists() and not exist_ok:
        raise FileExistsError(f"Session {session_id} already exists")
    
    serialized_session = serialize_session(session_doc)

    # Use file locking to prevent concurrent writes
    # But handle permission errors gracefully - if we can't create lock file, save without lock
    lock_file = None
    try:
        logger.debug(f"🔒 Acquiring EXCLUSIVE lock for session {session_id}")
        lock_file = open(lock_path, 'w')
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)  # Exclusive lock
        logger.debug(f"✅ Acquired EXCLUSIVE lock for session {session_id}")
    except PermissionError as perm_err:
        # Lock file permission issue - log warning but continue without lock
        logger.warning(f"⚠️  Cannot create lock file for session {session_id}: {perm_err}")
        logger.warning(f"   Saving session file without lock (may cause issues if another process is writing)")
        lock_file = None
    
    try:
        # Similar approach to jobs: use a temp file and atomic rename.
        temp_session_filename = session_id + ".tmp"
        temp_session_path = config.session_dir / temp_session_filename

        # Write JSON to temp file
        json_str = json.dumps(serialized_session, indent=4)
        temp_session_path.write_text(json_str)
        
        # Verify the temp file is valid JSON before atomically renaming
        try:
            json.loads(temp_session_path.read_text())
        except json.JSONDecodeError as e:
            logger.error(f"💥 Serialization created invalid JSON for session {session_id}")
            logger.error(f"   Error: {e}")
            # Clean up bad temp file
            if temp_session_path.exists():
                temp_session_path.unlink()
            raise ValueError(f"Failed to serialize session {session_id} to valid JSON") from e
        
        # Atomic rename - this ensures we never have a partially written session file
        temp_session_path.rename(session_path)
        logger.info(f"Session {session_id} saved to {session_path}")
        
    except Exception as e:
        # Clean up temp file if anything goes wrong
        if temp_session_path.exists():
            try:
                temp_session_path.unlink()
            except:
                pass
        raise
    finally:
        if lock_file:
            try:
                logger.debug(f"🔓 Releasing EXCLUSIVE lock for session {session_id}")
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)  # Release lock
                lock_file.close()
                logger.debug(f"✅ Released EXCLUSIVE lock for session {session_id}")
            except Exception as unlock_err:
                logger.warning(f"⚠️  Error releasing lock for session {session_id}: {unlock_err}")


def resolve_session_path(session_id: str, path_from_session: str) -> Path:
    """
    Resolve a path from session metadata, checking published location first.
    
    When a session is published, files are moved to /sphere/published/<sessionId>/output/.
    This function checks the published location first, then falls back to the path
    stored in session metadata (which may be in featrix_output).
    
    Args:
        session_id: Session ID
        path_from_session: Path as stored in session metadata (may be absolute or relative)
        
    Returns:
        Resolved Path object, checking published location first, then original path
    """
    if not path_from_session:
        raise ValueError(f"Empty path provided for session {session_id}")
    
    original_path = Path(path_from_session)
    
    # If path is absolute and exists, use it (handles both published and regular locations)
    if original_path.is_absolute() and original_path.exists():
        logger.debug(f"📦 Using absolute path from session: {original_path}")
        return original_path
    
    # Extract filename/relative path from the original path
    # If it's an absolute path, get just the filename
    # If it's relative, use it as-is
    if original_path.is_absolute():
        filename = original_path.name
        # Try to extract relative path from common output directories
        for base_dir in [config.output_dir, Path("/sphere/published")]:
            try:
                relative = original_path.relative_to(base_dir)
                if str(relative).startswith(session_id):
                    # Path is under session directory, extract the part after session_id
                    parts = relative.parts
                    if len(parts) > 1 and parts[0] == session_id:
                        relative_path = Path(*parts[1:])
                    else:
                        relative_path = Path(filename)
                    break
            except ValueError:
                continue
        else:
            relative_path = Path(filename)
    else:
        relative_path = original_path
    
    # Check published location first
    published_dir = Path("/sphere/published") / session_id
    published_output_dir = published_dir / "output"
    
    # Try multiple possible locations in published directory
    published_candidates = [
        published_output_dir / relative_path,
        published_output_dir / relative_path.name,  # Just filename
        published_dir / relative_path,
        published_dir / relative_path.name,  # Just filename
    ]
    
    for candidate in published_candidates:
        if candidate.exists():
            logger.debug(f"📦 Resolved {path_from_session} from published location: {candidate}")
            return candidate
    
    # Fall back to original path (may be in featrix_output)
    if original_path.exists():
        logger.debug(f"📦 Using original path from session: {original_path}")
        return original_path
    
    # Try to construct path in regular output directory
    regular_path = config.output_dir / session_id / relative_path
    if regular_path.exists():
        logger.debug(f"📦 Resolved {path_from_session} from regular location: {regular_path}")
        return regular_path
    
    # If neither exists, return the original path (will fail when accessed, but preserves expected behavior)
    logger.warning(f"⚠️  Path {path_from_session} not found in published or regular location, returning original path")
    return original_path


def load_session(session_id: str):
    """Load a session by session_id. Checks both regular and published locations."""
    session_path = config.session_dir / f"{session_id}.session"
    lock_path = config.session_dir / f"{session_id}.lock"
    
    # Check if it's a symlink to published location
    if session_path.is_symlink():
        session_path = session_path.resolve()
        logger.debug(f"📦 Session {session_id} is symlinked to published location: {session_path}")

    # If not found in regular location, check published location
    if not session_path.is_file():
        published_session_file = Path("/sphere/published") / session_id / f"{session_id}.session"
        if published_session_file.exists():
            session_path = published_session_file
            logger.debug(f"📦 Loading session {session_id} from published location: {session_path}")
        else:
            raise FileNotFoundError(f"Session {session_id} does not exist")

    # Use file locking to prevent reading while another process is writing
    # But handle permission errors gracefully - if we can't create lock file, read without lock
    lock_file = None
    try:
        logger.debug(f"🔒 Acquiring SHARED lock for session {session_id}")
        lock_file = open(lock_path, 'w')
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_SH)  # Shared lock for reading
        logger.debug(f"✅ Acquired SHARED lock for session {session_id}")
    except PermissionError as perm_err:
        # Lock file permission issue - log warning but continue without lock
        logger.warning(f"⚠️  Cannot create lock file for session {session_id}: {perm_err}")
        logger.warning(f"   Reading session file without lock (may cause issues if another process is writing)")
        lock_file = None
    
    try:
        try:
            session_doc = json.loads(session_path.read_text())
        except json.JSONDecodeError as e:
            # Log at WARNING level - corrupted files are often just old data that can be skipped
            # ERROR level is too noisy when iterating through many sessions
            # Include session_id in the message so we know which session is corrupted
            logger.warning(f"⚠️ Corrupted session file: {session_path} (session_id: {session_id})")
            logger.warning(f"   JSON error: {e}")
            logger.warning(f"   File size: {session_path.stat().st_size} bytes")
            
            # Try to read the file and show context around the error (only at DEBUG level to reduce noise)
            if logger.isEnabledFor(logging.DEBUG):
                try:
                    content = session_path.read_text()
                    lines = content.split('\n')
                    error_line = e.lineno - 1  # 0-indexed
                    start = max(0, error_line - 3)
                    end = min(len(lines), error_line + 4)
                    
                    logger.debug(f"   Context around line {e.lineno}:")
                    for i in range(start, end):
                        marker = " --> " if i == error_line else "     "
                        logger.debug(f"   {marker}{i+1}: {lines[i]}")
                except Exception as debug_err:
                    logger.debug(f"   Could not read file for debugging: {debug_err}")
            
            raise ValueError(f"Session {session_id} file is corrupted. JSON parse error at line {e.lineno}, column {e.colno}") from e

        session_doc['created_at'] = convert_from_iso(session_doc['created_at'])
        session_doc['status'] = SessionStatus(session_doc['status'])

        # TODO: validate the session document

        return session_doc
    finally:
        if lock_file:
            try:
                logger.debug(f"🔓 Releasing SHARED lock for session {session_id}")
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
                lock_file.close()
                logger.debug(f"✅ Released SHARED lock for session {session_id}")
            except Exception as unlock_err:
                logger.warning(f"⚠️  Error releasing lock for session {session_id}: {unlock_err}")


def serialize_session_private(session_private_doc: dict):
    session_private_doc = deepcopy(session_private_doc)

    # TODO: Check if valid email

    return session_private_doc


def save_session_private(session_id: str, session_private_doc: dict, exist_ok: bool = False):
    session_private_path = config.session_private_dir / f"{session_id}.session_private"

    # If the session already exists and we're not overwriting, raise an error.
    if session_private_path.exists() and not exist_ok:
        raise FileExistsError(f"Private session {session_id} already exists")
    
    serialized_session_private = serialize_session_private(session_private_doc)

    # Similar approach to jobs: use a temp file and atomic rename.
    temp_session_filename = session_id + ".tmp"
    temp_session_path = config.session_private_dir / temp_session_filename

    temp_session_path.write_text(json.dumps(serialized_session_private, indent=4))
    temp_session_path.rename(session_private_path)

    logger.info(f"Private session data for {session_id} saved to {session_private_path}")


def publish_session(session_id: str) -> dict:
    """
    Publish a session by moving it to /sphere/published/<sessionId>.
    Moves both the session file and output directory.
    Handles cross-filesystem moves by copying then deleting.
    
    Returns:
        dict with published_path, output_path, and status
    """    
    logger.info(f"📦 Publishing session {session_id}")
    
    # Load session first to verify it exists
    session = load_session(session_id)
    
    # Check if already published
    if session.get("published"):
        published_path = session.get("published_path")
        logger.warning(f"⚠️  Session {session_id} is already published at {published_path}")
        return {
            "status": "already_published",
            "published_path": published_path,
            "session_id": session_id
        }
    
    # Define paths
    published_base = Path("/sphere/published")
    published_dir = published_base / session_id
    published_session_file = published_dir / f"{session_id}.session"
    published_output_dir = published_dir / "output"
    
    # Source paths
    source_session_file = config.session_dir / f"{session_id}.session"
    source_output_dir = config.output_dir / session_id
    
    # Create published directory
    published_dir.mkdir(parents=True, exist_ok=True)
    
    # Move session file (handle cross-filesystem)
    try:
        if source_session_file.exists():
            if source_session_file.stat().st_dev == published_dir.stat().st_dev:
                # Same filesystem - use move
                shutil.move(str(source_session_file), str(published_session_file))
                logger.info(f"✅ Moved session file to {published_session_file}")
            else:
                # Different filesystem - copy then delete
                shutil.copy2(str(source_session_file), str(published_session_file))
                source_session_file.unlink()
                logger.info(f"✅ Copied session file to {published_session_file} (cross-filesystem)")
        else:
            logger.warning(f"⚠️  Session file not found at {source_session_file}")
    except Exception as e:
        logger.error(f"❌ Failed to move session file: {e}")
        raise
    
    # Move output directory (handle cross-filesystem)
    try:
        if source_output_dir.exists():
            if source_output_dir.stat().st_dev == published_dir.stat().st_dev:
                # Same filesystem - use move
                shutil.move(str(source_output_dir), str(published_output_dir))
                logger.info(f"✅ Moved output directory to {published_output_dir}")
            else:
                # Different filesystem - copy then delete
                shutil.copytree(str(source_output_dir), str(published_output_dir), dirs_exist_ok=True)
                shutil.rmtree(str(source_output_dir))
                logger.info(f"✅ Copied output directory to {published_output_dir} (cross-filesystem)")
        else:
            logger.warning(f"⚠️  Output directory not found at {source_output_dir}")
            published_output_dir.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        logger.error(f"❌ Failed to move output directory: {e}")
        raise
    
    # Update session metadata
    session["published"] = True
    session["published_path"] = str(published_dir)
    session["published_at"] = convert_to_iso(datetime.now(tz=ZoneInfo("America/New_York")))
    
    # Update paths in session document to point to published location
    # Update embedding_space path if it exists
    if "embedding_space" in session:
        old_path = Path(session["embedding_space"])
        if old_path.exists():
            # Try to update relative to published directory
            relative_path = old_path.relative_to(config.output_dir) if str(old_path).startswith(str(config.output_dir)) else old_path.name
            session["embedding_space"] = str(published_output_dir / relative_path)
    
    # Save session to published location
    save_session(session_id, session, exist_ok=True)
    # Also save to original location as a symlink or reference
    # For now, we'll just update the original location to point to published
    try:
        source_session_file.parent.mkdir(parents=True, exist_ok=True)
        source_session_file.symlink_to(published_session_file)
        logger.info(f"✅ Created symlink from {source_session_file} to {published_session_file}")
    except Exception as e:
        logger.warning(f"⚠️  Could not create symlink: {e}")
    
    logger.info(f"✅ Session {session_id} published to {published_dir}")
    
    return {
        "status": "published",
        "published_path": str(published_dir),
        "output_path": str(published_output_dir),
        "session_id": session_id
    }


def deprecate_session(session_id: str, warning_message: str, expiration_date: str) -> dict:
    """
    Deprecate a published session with a warning message and expiration date.
    The session remains available until the expiration date.
    
    Args:
        session_id: Session ID to deprecate
        warning_message: Warning message to display about deprecation
        expiration_date: ISO format date string when session will be removed
        
    Returns:
        dict with deprecation status
    """
    logger.info(f"⚠️  Deprecating session {session_id}")
    
    # Load session (may be in published location)
    try:
        session = load_session(session_id)
    except FileNotFoundError:
        # Try published location
        published_session_file = Path("/sphere/published") / session_id / f"{session_id}.session"
        if published_session_file.exists():
            session = json.loads(published_session_file.read_text())
        else:
            raise FileNotFoundError(f"Session {session_id} not found")
    
    # Parse expiration date
    try:
        expiration_dt = convert_from_iso(expiration_date)
    except Exception as e:
        raise ValueError(f"Invalid expiration_date format: {expiration_date}. Use ISO format.") from e
    
    # Update session metadata
    session["deprecated"] = True
    session["deprecation_warning"] = warning_message
    session["deprecation_expiration"] = expiration_date
    session["deprecated_at"] = convert_to_iso(datetime.now(tz=ZoneInfo("America/New_York")))
    
    # Save session
    save_session(session_id, session, exist_ok=True)
    
    logger.info(f"✅ Session {session_id} deprecated. Expires: {expiration_date}")
    
    return {
        "status": "deprecated",
        "session_id": session_id,
        "warning_message": warning_message,
        "expiration_date": expiration_date
    }


def unpublish_session(session_id: str) -> dict:
    """
    Unpublish a session by moving it back from /sphere/published/<sessionId>.
    
    Returns:
        dict with unpublish status
    """
    logger.info(f"📤 Unpublishing session {session_id}")
    
    # Try to load from published location first
    published_session_file = Path("/sphere/published") / session_id / f"{session_id}.session"
    
    if not published_session_file.exists():
        # Try regular location
        regular_session_file = config.session_dir / f"{session_id}.session"
        if regular_session_file.exists():
            logger.info(f"✅ Session {session_id} is not published, already in regular location")
            return {
                "status": "not_published",
                "session_id": session_id
            }
        else:
            raise FileNotFoundError(f"Session {session_id} not found in published or regular location")
    
    # Load session
    session = json.loads(published_session_file.read_text())
    
    # Define paths
    published_dir = Path("/sphere/published") / session_id
    published_output_dir = published_dir / "output"
    
    # Destination paths
    dest_session_file = config.session_dir / f"{session_id}.session"
    dest_output_dir = config.output_dir / session_id
    
    # Remove symlink if it exists
    if dest_session_file.exists() and dest_session_file.is_symlink():
        dest_session_file.unlink()
        logger.info(f"✅ Removed symlink at {dest_session_file}")
    
    # Move session file back
    try:
        if published_session_file.stat().st_dev == dest_session_file.parent.stat().st_dev:
            # Same filesystem - use move
            shutil.move(str(published_session_file), str(dest_session_file))
            logger.info(f"✅ Moved session file back to {dest_session_file}")
        else:
            # Different filesystem - copy then delete
            shutil.copy2(str(published_session_file), str(dest_session_file))
            published_session_file.unlink()
            logger.info(f"✅ Copied session file back to {dest_session_file} (cross-filesystem)")
    except Exception as e:
        logger.error(f"❌ Failed to move session file back: {e}")
        raise
    
    # Move output directory back
    try:
        if published_output_dir.exists():
            if published_output_dir.stat().st_dev == dest_output_dir.parent.stat().st_dev:
                # Same filesystem - use move
                shutil.move(str(published_output_dir), str(dest_output_dir))
                logger.info(f"✅ Moved output directory back to {dest_output_dir}")
            else:
                # Different filesystem - copy then delete
                shutil.copytree(str(published_output_dir), str(dest_output_dir), dirs_exist_ok=True)
                shutil.rmtree(str(published_output_dir))
                logger.info(f"✅ Copied output directory back to {dest_output_dir} (cross-filesystem)")
    except Exception as e:
        logger.error(f"❌ Failed to move output directory back: {e}")
        raise
    
    # Update session metadata
    session["published"] = False
    if "published_path" in session:
        del session["published_path"]
    if "published_at" in session:
        del session["published_at"]
    
    # Save session
    save_session(session_id, session, exist_ok=True)
    
    # Clean up published directory if empty
    try:
        if published_dir.exists():
            # Check if directory is empty (except maybe .session file which we already moved)
            remaining = list(published_dir.iterdir())
            if not remaining or (len(remaining) == 1 and remaining[0].name.endswith('.session') and not remaining[0].exists()):
                published_dir.rmdir()
                logger.info(f"✅ Removed empty published directory {published_dir}")
    except Exception as e:
        logger.warning(f"⚠️  Could not remove published directory: {e}")
    
    logger.info(f"✅ Session {session_id} unpublished")
    
    return {
        "status": "unpublished",
        "session_id": session_id
    }


def load_session_private(session_id: str, empty_ok=False):
    session_private_path = config.session_private_dir / f"{session_id}.session_private"

    if not session_private_path.is_file():
        if empty_ok:
            # If the file doesn't yet exist, we return an empty dict.
            return dict()
        else:
            raise FileNotFoundError(f"Private session data for {session_id} does not exist")

    session_private_doc = json.loads(session_private_path.read_text())

    return session_private_doc


def add_notification_email_to_session(session_id: str, email: str):
    session_doc = load_session(session_id)

    if "notification_email" not in session_doc:
        session_doc["notification_email"] = []

    session_doc["notification_email"].append(email)

    save_session(session_id, session_doc, exist_ok=True)

    logger.info(f"Notification email {email} added to session {session_id}")


def send_session_done_notification(session_id: str):
    session_doc = load_session(session_id)

    notification_emails = session_doc.get("notification_email", [])

    if len(notification_emails) == 0:
        logger.info(f"No email found for session {session_id}. Notification not sent.")
        return

    # TODO: send notification email
    # send_email(notification_emails[0], f"Session {session_id} is done")


def load_session_jobs(session_doc: dict):
    jobs = dict()

    for job_info in session_doc.get("job_plan", []):
        job_type = job_info.get("job_type")
        job_id = job_info.get("job_id")
        if job_id is not None:
            try:
                job = load_job(queue_name=job_type, job_id=job_id)
                jobs[job_id] = job
            except FileNotFoundError:
                # Job file no longer exists (deleted, cleaned up, etc.)
                # This is expected behavior - completed jobs are often cleaned up
                # Log at debug level to avoid noise, but don't fail the entire session load
                session_id = session_doc.get("session_id", "unknown")
                logger.debug(f"Job {job_id} referenced in session {session_id} but file not found in queue {job_type} - skipping (expected after cleanup)")
                # Create a placeholder job entry so the session shows the job existed
                jobs[job_id] = {
                    "job_id": job_id,
                    "type": job_type,
                    "status": "deleted",
                    "error": "Job file no longer exists"
                }

    return jobs


def iterate_over_sessions():
    session_dir = config.session_dir

    if not session_dir.exists():
        logger.warning(f"Session directory {config.session_dir.resolve()} does not exist")
        return

    for session_path in session_dir.iterdir():
        # Skip non-session files.
        if session_path.suffix != ".session":
            continue

        # Just the name, without the extension
        session_id = session_path.stem

        yield load_session(session_id)


def print_session(session_doc: dict):
    serialized_session = serialize_session(session_doc)

    logger.info(f"Name: {session_doc.get('session_id')}")
    logger.info(json.dumps(serialized_session, indent=4))


def get_session_info(session_id: str):
    """
    Get session information with optimized queue scanning.
    
    OPTIMIZATION: Batch queue scans to avoid redundant file I/O.
    - Groups jobs by queue name
    - Scans each queue only once
    - Reuses scan results for position calculation and detailed info
    """
    session = load_session(session_id)
    session_jobs = load_session_jobs(session)
    session_jobs_queue_positions = dict()
    
    # Enhanced queue information for better user experience
    detailed_queue_info = {}
    
    # OPTIMIZATION: Group jobs by queue to batch queue scans
    jobs_by_queue = {}
    for job_id, job in session_jobs.items():
        job_queue = job.get("type")
        if job_queue not in jobs_by_queue:
            jobs_by_queue[job_queue] = []
        jobs_by_queue[job_queue].append((job_id, job))
    
    # OPTIMIZATION: Scan each queue only once and cache results
    queue_cache = {}  # queue_name -> {'ready_jobs': [...], 'all_jobs': [...], 'ready_job_ids': [...]}
    
    for queue_name in jobs_by_queue.keys():
        # Scan queue once to get all ready jobs and all jobs
        ready_jobs = get_ready_jobs(queue_name, order="oldest_first")
        all_jobs = list(iterate_over_jobs_in_queue(queue_name))
        
        queue_cache[queue_name] = {
            'ready_jobs': ready_jobs,
            'all_jobs': all_jobs,
            'ready_job_ids': [j.get("job_id") for j in ready_jobs]
        }
    
    # Now process each job using cached queue data
    for job_id, job in session_jobs.items():
        job_queue = job.get("type")
        job_status = job.get("status")
        
        # Skip deleted/placeholder jobs when getting queue position
        if job_status == "deleted":
            session_jobs_queue_positions[job_id] = None
            detailed_queue_info[job_id] = {
                "position": None,
                "status": "deleted",
                "queue_name": job_queue,
                "queue_status": "deleted",
                "estimated_wait_message": "Job file no longer exists"
            }
            continue
        
        # OPTIMIZATION: Calculate position using cached queue data
        job_position = None
        if job_status == JobStatus.READY and job_queue in queue_cache:
            cache = queue_cache[job_queue]
            try:
                job_position = cache['ready_job_ids'].index(job_id)
            except ValueError:
                # Job not in ready list (shouldn't happen, but handle gracefully)
                job_position = None
        elif job_status != JobStatus.READY:
            # Not a READY job, so not in queue
            job_position = None
        
        session_jobs_queue_positions[job_id] = job_position
        
        # Collect detailed queue information
        detailed_info = {
            "position": job_position,
            "status": job_status.value if hasattr(job_status, 'value') else str(job_status),
            "queue_name": job_queue,
        }
        
        # Add detailed status information based on job state
        if job_status == JobStatus.READY and job_position is not None:
            # OPTIMIZATION: Use cached queue data instead of re-scanning
            cache = queue_cache.get(job_queue, {})
            ready_jobs = cache.get('ready_jobs', [])
            all_jobs = cache.get('all_jobs', [])
            total_ready = len(ready_jobs)
            
            # Get currently running jobs from cached all_jobs
            running_jobs = []
            for queue_job in all_jobs:
                if queue_job.get("status") == JobStatus.RUNNING:
                    running_jobs.append({
                        "job_id": queue_job.get("job_id"),
                        "started_at": convert_to_iso(queue_job.get("started_at")),  # Serialize datetime
                        "session_id": queue_job.get("session_id")
                    })
            
            detailed_info.update({
                "queue_status": "waiting",
                "position_in_queue": job_position,
                "total_jobs_ahead": job_position,
                "total_ready_jobs": total_ready,
                "currently_running_jobs": running_jobs,
                "worker_available": len(running_jobs) == 0,  # Simplified - assumes 1 worker per queue
                "estimated_wait_message": _generate_wait_message(job_position, running_jobs, job_queue),
                "next_to_run": job_position == 0,
            })
            
        elif job_status == JobStatus.RUNNING:
            # Job is currently running
            started_at = job.get("started_at")
            progress = job.get("progress", 0)
            
            detailed_info.update({
                "queue_status": "running",
                "started_at": convert_to_iso(started_at),  # Serialize datetime
                "progress": progress,
                "progress_percent": round(progress * 100, 1) if progress else 0,
                "estimated_wait_message": "Currently running",
                "current_epoch": job.get("current_epoch"),
                "current_loss": job.get("current_loss"),
                "validation_loss": job.get("validation_loss"),
            })
            
        elif job_status == JobStatus.DONE:
            finished_at = job.get("finished_at")
            detailed_info.update({
                "queue_status": "completed",
                "finished_at": convert_to_iso(finished_at),  # Serialize datetime
                "estimated_wait_message": "Completed successfully",
                "progress_percent": 100,
            })
            
        elif job_status == JobStatus.FAILED:
            finished_at = job.get("finished_at")
            detailed_info.update({
                "queue_status": "failed",
                "finished_at": convert_to_iso(finished_at),  # Serialize datetime
                "estimated_wait_message": "Job failed - check logs",
                "progress_percent": 0,
            })
        
        elif job_status == JobStatus.PAUSED:
            paused_at = job.get("paused_at")
            pause_reason = job.get("pause_reason", "Paused by user")
            detailed_info.update({
                "queue_status": "paused",
                "paused_at": paused_at,  # Already ISO format string
                "pause_reason": pause_reason,
                "estimated_wait_message": f"Job paused: {pause_reason}",
                "progress_percent": job.get("progress", 0) * 100 if job.get("progress") else 0,
            })
        
        detailed_queue_info[job_id] = detailed_info

    return {
        "session": session,
        "jobs": session_jobs,
        "job_queue_positions": session_jobs_queue_positions,
        "detailed_queue_info": detailed_queue_info,  # New enhanced information
    }


def _generate_wait_message(position: int, running_jobs: list, queue_name: str) -> str:
    """Generate a human-readable message about queue wait status."""
    
    if position == 0:
        if len(running_jobs) == 0:
            return "🚀 Next in line - starting shortly"
        else:
            running_job = running_jobs[0]
            running_session = running_job.get("session_id", "unknown")
            return f"🚀 Next in line - waiting for current job to complete (session: {running_session})"
    
    elif position == 1:
        return f"🥈 Second in line - 1 job ahead of you"
    
    elif position <= 5:
        return f"⏳ Position {position + 1} in queue - {position} jobs ahead of you"
    
    else:
        return f"⏳ Position {position + 1} in queue - {position} jobs ahead of you (busy queue)"


def get_queue_summary() -> dict:
    """Get a summary of all queue states for admin/debugging purposes."""
    
    queue_summary = {}
    
    # Check all known queue types
    known_queues = [
        "create_structured_data",
        "train_es", 
        "train_knn",
        "build_es_projections",
        "train_single_predictor"
    ]
    
    for queue_name in known_queues:
        queue_path = config.queue_dir / queue_name
        
        if not queue_path.exists():
            queue_summary[queue_name] = {
                "exists": False,
                "ready_jobs": 0,
                "running_jobs": 0,
                "total_jobs": 0,
                "status": "queue_not_found"
            }
            continue
        
        ready_jobs = []
        running_jobs = []
        total_jobs = 0
        
        try:
            for job in iterate_over_jobs_in_queue(queue_name):
                total_jobs += 1
                status = job.get("status")
                
                if status == JobStatus.READY:
                    ready_jobs.append(job)
                elif status == JobStatus.RUNNING:
                    running_jobs.append(job)
        except Exception as e:
            queue_summary[queue_name] = {
                "exists": True,
                "error": str(e),
                "status": "error_reading_queue"
            }
            continue
        
        queue_summary[queue_name] = {
            "exists": True,
            "ready_jobs": len(ready_jobs),
            "running_jobs": len(running_jobs),
            "total_jobs": total_jobs,
            "status": "active" if len(running_jobs) > 0 else "idle",
            "next_job_id": ready_jobs[0].get("job_id") if ready_jobs else None,
            "currently_running": [job.get("job_id") for job in running_jobs],
        }
    
    return queue_summary


def list_sessions(show_jobs: bool = False):
    """List all sessions with optional job details."""

    sessions = list(iterate_over_sessions())

    # Sort sessions from newest to oldest
    sessions = sorted(sessions, key=lambda x: x.get("created_at"), reverse=True)

    if len(sessions) == 0:
        logger.info(f"No sessions found in {config.session_dir.resolve()}")
        return

    for session in sessions:
        print_session(session)

        if show_jobs:
            session_jobs = load_session_jobs(session)

            if len(session_jobs) == 0:
                logger.info("No jobs in this session")
            else:
                logger.info("Session jobs:")
                for job in session_jobs:
                    print_job(job)

            logger.info("")  # Empty line for readability


# def clear_sessions():
#     """Delete all sessions."""
#     clear_directory(config.session_dir)


def step_session(session_id: str):
    # Load a session, check what other work needs to be done, and 
    # schedule the next job.
    
    logger.info(f"🔍 step_session() called for session_id={session_id}")
    logger.info(f"📋 step_session: Processing session {session_id}")

    session = load_session(session_id=session_id)
    loaded_session_id = session.get("session_id")
    if loaded_session_id != session_id:
        logger.error(f"❌ MISMATCH: step_session called with session_id={session_id} but loaded session has session_id={loaded_session_id}")
    logger.debug(f"✅ Session {session_id} loaded successfully (verified: {loaded_session_id})")
    
    session_type = session.get('session_type')

    # Check if session is already complete - if so, don't queue more jobs
    session_status = session.get("status")
    if session_status == SessionStatus.DONE:
        logger.info(f"Session {session_id} is already DONE, skipping step_session")
        return
    
    new_job_id = None
    queue_name = None

    job_plan = session.get("job_plan")
    if job_plan is None or len(job_plan) == 0:
        raise ValueError(f"Session {session_id} does not have a job plan")
    
    # Track which job plan index got a new job_id (for jobs with multiple instances of same type)
    job_plan_index_to_update = None
    job_type_to_update = None
    
    # Track if we broke because of a READY job (session should remain RUNNING)
    has_ready_job = False

    # Job descriptions in the job plan are in the order in which they should
    # be executed, so we just iterate over them.
    for idx, job_description in enumerate(job_plan):
        job_type = job_description.get("job_type")
        
        # Map job types to queue names
        # create_structured_data runs on cpu_data_tasks queue
        queue_name_map = {
            "create_structured_data": "cpu_data_tasks",
            "generate_movie_frame": "cpu_data_tasks",
        }
        queue_name = queue_name_map.get(job_type, job_type)  # Default to job_type if not mapped
        
        job_spec = job_description.get("spec")
        job_id = job_description.get("job_id")

        if job_id is None:
            # CRITICAL: Use the function parameter session_id explicitly to avoid any variable shadowing
            # Store it in a local variable to ensure we're using the correct one
            es_session_id = session_id  # Explicitly capture the parameter
            
            logger.info(f"🎯 step_session: Found job in plan with job_id=None (job_type={job_type}, index={idx})")
            
            # Skip train_es, train_knn, and build_es_projections if embedding space already exists (foundation model)
            # This allows training predictors on sessions with imported embedding spaces
            if job_type == "train_es":
                embedding_space_path = session.get("embedding_space")
                foundation_model_id = session.get("foundation_model_id")
                
                # CRITICAL: Never train a new ES if we have a foundation model or existing ES
                if foundation_model_id:
                    logger.warning(f"🚫 BLOCKED train_es job - session {session_id} is based on foundation model {foundation_model_id}")
                    logger.warning(f"   This session MUST reuse the foundation's embedding space - NO new ES training allowed")
                    job_plan[idx]["job_id"] = "skipped-foundation-model"
                    continue
                elif embedding_space_path and Path(embedding_space_path).exists():
                    logger.info(f"⏭️  Skipping train_es job - embedding space already exists: {embedding_space_path}")
                    # Mark job as skipped by setting a placeholder job_id (won't be queued)
                    job_plan[idx]["job_id"] = "skipped-foundation-model"
                    continue
            elif job_type == "train_knn":
                embedding_space_path = session.get("embedding_space")
                vector_db_path = session.get("vector_db")
                if embedding_space_path and Path(embedding_space_path).exists():
                    if vector_db_path and Path(vector_db_path).exists():
                        logger.info(f"⏭️  Skipping train_knn job - vector_db already exists: {vector_db_path}")
                        job_plan[idx]["job_id"] = "skipped-foundation-model"
                        continue
                    else:
                        # ES exists but no vector_db - KNN training might still be needed
                        # But for foundation models, we often don't need KNN, so skip it
                        logger.info(f"⏭️  Skipping train_knn job - embedding space exists (foundation model, KNN not required)")
                        job_plan[idx]["job_id"] = "skipped-foundation-model"
                        continue
            elif job_type == "build_es_projections":
                embedding_space_path = session.get("embedding_space")
                if embedding_space_path and Path(embedding_space_path).exists():
                    logger.info(f"⏭️  Skipping build_es_projections job - embedding space already exists (foundation model): {embedding_space_path}")
                    job_plan[idx]["job_id"] = "skipped-foundation-model"
                    continue
            
            logger.info(f"   This job needs to be created and queued")
            
            # Check queue depth before creating a new job
            ready_count = 0
            running_count = 0
            for queue_job in iterate_over_jobs_in_queue(queue_name):
                status = queue_job.get("status")
                if status == JobStatus.READY:
                    ready_count += 1
                elif status == JobStatus.RUNNING:
                    running_count += 1
            
            total_queued = ready_count + running_count
            MAX_QUEUE_DEPTH = 64
            
            if total_queued >= MAX_QUEUE_DEPTH:
                raise ValueError(
                    f"Cannot queue new job: queue {queue_name} is full "
                    f"({total_queued} jobs queued/running, max {MAX_QUEUE_DEPTH}). "
                    f"Please wait for some jobs to complete."
                )
            
            logger.info(f"   Queue depth check: {total_queued}/{MAX_QUEUE_DEPTH} - proceeding to create job")
            
            # Child ES sessions are top-level - use their own session_id
            job_session_id = es_session_id
            logger.info(f"📋 Creating job for top-level child ES session: {job_session_id}")
            
            # Use job_session_id directly - it should already be the user-facing session_id (with prefix)
            logger.info(f"🚀 step_session: Creating job for session_id={job_session_id}, job_type={job_type}, queue_name={queue_name}")
            logger.info(f"   Job spec: {json.dumps(job_spec, indent=2) if job_spec else 'None'}")
            
            try:
                new_job_id = add_job_to_queue(queue_name=queue_name, job_type=job_type, session_id=job_session_id, job_spec=job_spec)
                logger.info(f"✅ step_session: Successfully created job {new_job_id} with session_id={job_session_id}")
                logger.info(f"   Job file: {config.queue_dir / queue_name / f'{new_job_id}.job'}")
            except Exception as create_error:
                logger.error(f"❌ CRITICAL: Failed to create job in step_session!")
                logger.error(f"   session_id: {job_session_id}")
                logger.error(f"   job_type: {job_type}")
                logger.error(f"   queue_name: {queue_name}")
                logger.error(f"   Error: {create_error}")
                logger.error(f"   Traceback: {traceback.format_exc()}")
                raise  # Re-raise to surface the error
            # You'd think we can just do `job_description["job_id"] = new_job_id` here,
            # but we're iterating over the job plan, and mutating the structure we're iterating over
            # is a bad idea, and leads to bugs.
            # Instead, we keep track of which job plan entry needs updating, and do it after the loop.
            job_plan_index_to_update = idx
            job_type_to_update = job_type

            # Break out of the loop after scheduling the first job.
            # We only schedule one job per iteration.
            break

        else:
            # Handle skipped jobs (foundation model case)
            if job_id == "skipped-foundation-model":
                logger.debug(f"⏭️  Job {job_type} was skipped (foundation model) - continuing to next job")
                continue
            
            # Try to load the job, but if it doesn't exist (was deleted or cleaned up),
            # treat it as completed and continue to the next job
            try:
                job = load_job(queue_name, job_id)
            except FileNotFoundError:
                logger.warning(f"⚠️  Job {job_id} not found in queue {queue_name} - treating as completed")
                continue
            
            job_status = job.get("status")
            
            if job_status == JobStatus.DONE:
                # Skip done jobs to follow the plan.
                continue

            elif job_status == JobStatus.RUNNING:
                # Special case: If it's an ES training job that created the embedding space,
                # it actually completed even if status wasn't updated to DONE
                if job_type == "train_es":
                    embedding_space_path = session.get("embedding_space")
                    if embedding_space_path and Path(embedding_space_path).exists():
                        logger.warning(f"⚠️  ES job {job_id} status is RUNNING but embedding space exists - treating as DONE")
                        # Update the job status to DONE so this doesn't happen again
                        try:
                            update_job_status(queue_name="train_es", job_id=job_id, status=JobStatus.DONE)
                            logger.info(f"✅ Updated ES job {job_id} status to DONE")
                        except Exception as e:
                            logger.warning(f"Could not update ES job status: {e}")
                        continue
                
                # Allow queueing additional single predictor jobs even if one is already running
                # This enables training multiple predictors in parallel or queued on a session
                if job_type == "train_single_predictor":
                    logger.info(f"Single predictor job {job_id} is running, continuing to check for more jobs to queue")
                    continue
                
                # For other job types (train_es, train_knn, etc.), if a job is already running,
                # just continue to check remaining jobs in the plan
                # This allows scheduling dependent jobs that don't require the running job to complete
                logger.info(f"Job {job_id} in queue {queue_name} is running - continuing to check remaining jobs in plan")
                continue  # Check the next job in the plan

            elif job_status in [JobStatus.FAILED, JobStatus.CANCELLED]:
                # For single predictor jobs, allow continuing to add more predictors
                # even if previous ones failed - this enables multiple single predictors per session
                if job_type == "train_single_predictor":
                    logger.warning(f"Skipping failed single predictor job {job_id}, continuing with other jobs")
                    continue
                # For train_es jobs, if embedding space already exists, skip the failed job
                # This allows training predictors on sessions even if an ES training job failed
                elif job_type == "train_es":
                    embedding_space_path = session.get("embedding_space")
                    if embedding_space_path and Path(embedding_space_path).exists():
                        logger.warning(f"⚠️  ES job {job_id} failed but embedding space exists - skipping to allow predictor training")
                        continue
                    else:
                        # No embedding space exists, so we can't continue
                        raise ValueError(f"Job Plan error. Job {job_id} in queue {queue_name} failed or was cancelled. Skipping.")
                # For train_knn jobs, if vector_db already exists, skip the failed job
                # This allows training predictors on sessions even if a KNN training job failed
                elif job_type == "train_knn":
                    vector_db_path = session.get("vector_db")
                    if vector_db_path and Path(vector_db_path).exists():
                        logger.warning(f"⚠️  KNN job {job_id} failed but vector_db exists - skipping to continue")
                        continue
                    else:
                        # No vector_db exists, so we can't continue
                        raise ValueError(f"Job Plan error. Job {job_id} in queue {queue_name} failed or was cancelled. Skipping.")
                else:
                    # For other job types, still block on failures
                    raise ValueError(f"Job Plan error. Job {job_id} in queue {queue_name} failed or was cancelled. Skipping.")
            
            elif job_status == JobStatus.READY:
                # READY jobs are normal - they're waiting to be picked up by workers
                # For single predictor jobs, continue to check if we can queue more predictors
                if job_type == "train_single_predictor":
                    logger.info(f"Single predictor job {job_id} is ready and waiting, continuing to check for more jobs to queue")
                    continue
                # For other job types, don't schedule a new job, just wait for this one
                logger.info(f"Job {job_id} in queue {queue_name} is ready and waiting for worker")
                has_ready_job = True  # Mark that we have a READY job waiting
                break
            else: 
                raise ValueError(f"Job Plan error. {job_id} in queue {queue_name} has unexpected status: {job_status}")

    if new_job_id is None:
        # Check if we have a READY job waiting - if so, session is still RUNNING
        if has_ready_job:
            logger.info(f"Session {session_id} has a READY job waiting - keeping status as RUNNING")
            session["status"] = SessionStatus.RUNNING
        else:
            # If we get here, all jobs in the plan are done.
            logger.info(f"Session {session_id} is complete")
            session["status"] = SessionStatus.DONE
    else:
        logger.info(f"Session {session_id} is running. New job {new_job_id} on queue {queue_name} added to queue.")
        session["status"] = SessionStatus.RUNNING

    # Update the specific job plan entry with the new job_id
    if job_plan_index_to_update is not None:
        session["job_plan"][job_plan_index_to_update]["job_id"] = new_job_id
        logger.info(f"✅ Assigned job_id {new_job_id} to job plan index {job_plan_index_to_update} (type: {job_type_to_update})")

    logger.debug(f"💾 Saving session {session_id} after stepping...")
    save_session(session_id=session_id, session_doc=session, exist_ok=True)
    logger.debug(f"✅ Session {session_id} saved successfully after stepping")
    
    # Send notification when the session is done.
    if session["status"] == SessionStatus.DONE:
        send_slack_message(f"Session {session_id} is DONE.")
        send_session_done_notification(session_id)


def get_session_progress(session_id):

    # load session

    # load jobs for the session

    # calculate progress

    # return
    pass


##############################################################################
# 
# QUEUE WORKERS
# 
##############################################################################


def watch_queue(queue_name: str, skip_existing: bool = False, num_workers: int = None):
    # 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

    import warnings
    from queue import Queue
    from threading import Thread
    from concurrent.futures import ThreadPoolExecutor

    from watchdog.events import FileSystemEventHandler
    from watchdog.observers import Observer

    # Suppress Pydantic warnings in workers - they're noisy and not actionable
    warnings.filterwarnings("ignore", message=".*Field.*has conflict with protected namespace.*")
    warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")

    # Alias old queue names for backward compatibility
    if queue_name == 'create_structured_data':
        logger.info(f"📝 Queue name 'create_structured_data' is deprecated, using 'cpu_data_tasks'")
        queue_name = 'cpu_data_tasks'

    # Disable debug logging from watchdog to prevent "write to closed file" errors
    # when log rotation happens. Watchdog's internal threads try to write debug
    # messages to file handles that get closed during job log rotation.
    logging.getLogger("watchdog").setLevel(logging.WARNING)

    # For cpu_data_tasks, use thread pool for better I/O-bound task handling
    # For other queues (especially train_es), use single thread due to special exit behavior
    use_thread_pool = (queue_name == 'cpu_data_tasks')
    
    if num_workers is None:
        # Default: 4 workers for cpu_data_tasks (matching previous 4-process setup)
        # Single worker for other queues
        num_workers = 4 if use_thread_pool else 1
    
    if use_thread_pool:
        logger.info(f"🧵 Using thread pool with {num_workers} workers for {queue_name} queue")
        executor = ThreadPoolExecutor(max_workers=num_workers)
    else:
        logger.info(f"🧵 Using single worker thread for {queue_name} queue")
        executor = None

    # Initialize thread variables (will be set conditionally below)
    queue_processor = None
    worker_thread = None

    job_queue = Queue()

    def run_job_wrapper(job_id: str):
        """Wrapper for running jobs in thread pool"""
        try:
            logger.info(f"running job {job_id} in {queue_name}")
            err = run_job(queue_name=queue_name, job_id=job_id)
            
            # Don't restart worker for cpu_data_tasks - just log and continue
            if err == 2:
                logger.warning(f"🚫 Job {job_id} was aborted (exit code 2)")
                # Don't kill worker - just continue processing next job
            
            if err is not None:
                logger.error(f"job {job_id} failed: {err}")
            
            logger.info(f"job {job_id} done")
        except Exception as e:
            logger.error(f"An unrecoverable error occurred while running job {job_id}")
            logger.error(f"Exception: {str(e)}")
            logger.error(f"Traceback:\n{traceback.format_exc()}")

    def worker():
        while True:
            job_id = job_queue.get()
            if job_id is None:
                break

            logger.info(f"running job {job_id} in {queue_name}")
            try:
                # Get session_id before running job (in case job fails)
                try:
                    job = load_job(queue_name, job_id)
                    # Skip PAUSED jobs
                    if job.get("status") == JobStatus.PAUSED:
                        logger.info(f"⏸️  Job {job_id} is PAUSED (reason: {job.get('pause_reason', 'unknown')}) - skipping")
                        job_queue.task_done()
                        continue
                    session_id = job.get("session_id")
                    job_spec = job.get("spec", {})
                    # Check if this is a child ES training job that needs to step a different session
                    child_session_id = job_spec.get("_child_session_id")
                except Exception:
                    session_id = None
                    child_session_id = None
                
                err = run_job(queue_name=queue_name, job_id=job_id, auto_step=False)
                
                # After job completes successfully, step the session to queue next job
                # This is done OUTSIDE of run_job() to avoid deadlocks
                logger.info(f"🔍 Worker: Job {job_id} finished with err={err}, session_id={session_id}")
                if err is None:
                    if session_id:
                        logger.info(f"🚀 Job {job_id} completed successfully, calling step_session for session {session_id}")
                        try:
                            time.sleep(0.1)  # Small delay to ensure locks are released
                            step_session(session_id=session_id)
                            logger.info(f"✅ step_session completed successfully for session {session_id}")
                        except Exception as step_error:
                            logger.error(f"❌ CRITICAL: Failed to step session {session_id} after job {job_id}: {step_error}")
                            logger.error(f"   This means the next job in the plan was NOT queued!")
                            logger.error(f"   Session will be stuck until step_session is called manually!")
                            logger.error(f"   Traceback: {traceback.format_exc()}")
                            # Don't fail the job - stepping is a separate concern
                            # But log it prominently so we can see it
                            logger.error(f"🔴🔴🔴 SESSION STEP FAILED - MANUAL INTERVENTION NEEDED 🔴🔴🔴")
                    else:
                        logger.error(f"❌ CRITICAL: Job {job_id} has no session_id, cannot step session!")
                        logger.error(f"   This means the next job in the plan was NOT queued!")
                        try:
                            job_data = load_job(queue_name, job_id)
                            logger.error(f"   Job file contents: {json.dumps(job_data, indent=2, default=str)}")
                        except Exception as e:
                            logger.error(f"   Could not load job file: {e}")
                elif err is not None:
                    logger.warning(f"⚠️  Job {job_id} completed with error {err}, skipping step_session")
                
                # Check if job was aborted (exit code 2)
                # Only restart worker for train_es, not for cpu_data_tasks
                if err == 2 and queue_name != "cpu_data_tasks":
                    logger.warning(f"🚫 Job {job_id} was aborted - triggering worker restart")
                    job_queue.task_done()
                    # Signal supervisor to restart by exiting
                    os.kill(os.getpid(), signal.SIGTERM)
                    break
                elif err == 2:
                    logger.warning(f"🚫 Job {job_id} was aborted (exit code 2) - continuing worker")
                    # Don't restart cpu_data_tasks worker, just continue
                
                if err is not None:
                    logger.error(f"job {job_id} failed: {err}")
                    
                    # Track job restarts and pause if too many restarts in last hour
                    try:
                        job = load_job(queue_name, job_id)
                        restart_count = _count_recent_restarts(job, hours=1)
                        if restart_count >= 20:
                            logger.warning(f"⏸️  Job {job_id} has restarted {restart_count} times in the last hour - pausing job")
                            job['status'] = JobStatus.PAUSED.value
                            job['pause_reason'] = f'Too many restarts: {restart_count} in last hour (limit: 20)'
                            job['paused_at'] = datetime.now(tz=ZoneInfo("America/New_York")).isoformat()
                            save_job(queue_name, job_id, job, exist_ok=True)
                            logger.warning(f"⏸️  Job {job_id} paused - will skip until manually resumed")
                            continue  # Skip to next job
                    except Exception as e:
                        logger.debug(f"Could not check restart count for job {job_id}: {e}")

            except SystemExit as e:
                # Catch sys.exit() calls from training code
                # For train_es, always exit with code 0 so worker restarts cleanly
                if queue_name == "train_es":
                    exit_code = getattr(e, 'code', 0)
                    if exit_code != 0:
                        logger.info(f"🔄 train_es job {job_id} exited with code {exit_code}, converting to 0 for clean restart")
                    
                    # Track restart before exiting
                    try:
                        job = load_job(queue_name, job_id)
                        restart_count = _count_recent_restarts(job, hours=1)
                        if restart_count >= 20:
                            logger.warning(f"⏸️  Job {job_id} has restarted {restart_count} times in the last hour - pausing job")
                            job['status'] = JobStatus.PAUSED.value  # Convert enum to string for JSON serialization
                            job['pause_reason'] = f'Too many restarts: {restart_count} in last hour (limit: 20)'
                            job['paused_at'] = datetime.now(tz=ZoneInfo("America/New_York")).isoformat()
                            save_job(queue_name, job_id, job, exist_ok=True)
                            logger.warning(f"⏸️  Job {job_id} paused - will skip until manually resumed")
                    except Exception as e:
                        logger.debug(f"Could not check restart count for job {job_id}: {e}")
                    
                    # Exit with 0 so supervisor restarts the worker
                    sys.exit(0)
                else:
                    # For other queues, preserve the exit code
                    raise
            except Exception as e:
                logger.error(f"An unrecoverable error occurred while running job {job_id}")
                logger.error(f"Exception: {str(e)}")
                logger.error(f"Traceback:\n{traceback.format_exc()}")
                # For train_es, continue to exit with 0 even after exceptions
                # so the worker restarts cleanly
                if queue_name == "train_es":
                    logger.info(f"🔄 train_es job {job_id} had exception but will exit with code 0 for clean restart")
            
            logger.info(f"job {job_id} done")

            job_queue.task_done()
            
            # For train_es queue: Exit after each job so supervisor restarts the worker
            # This ensures clean state between embedding space training jobs
            # Always exit with code 0 (success) so supervisor restarts cleanly
            if queue_name == "train_es":
                logger.info(f"🔄 train_es job {job_id} complete - exiting worker for clean restart (exit code 0)")
                # Exit with code 0 so supervisor restarts the worker
                sys.exit(0)

    if use_thread_pool:
        # Thread pool mode: submit jobs to executor
        def process_job_queue():
            while True:
                job_id = job_queue.get()
                if job_id is None:
                    break

                # Submit job to thread pool
                executor.submit(run_job_wrapper, job_id)
                job_queue.task_done()
        
        queue_processor = Thread(target=process_job_queue)
        queue_processor.start()
    else:
        # Single-threaded mode: use original worker
        worker_thread = Thread(target=worker)
        worker_thread.start()

    # For train_es: Clean up dangling multiprocess workers before doing anything
    if queue_name == "train_es":
        logger.info("🧹 Cleaning up dangling pt_data_worker processes...")
        try:
            import psutil
            killed_count = 0
            current_pid = os.getpid()
            
            # Find all pt_data_worker processes
            # These are PyTorch DataLoader worker processes that may be left over from previous training
            for proc in psutil.process_iter(['pid', 'name', 'ppid', 'cmdline']):
                try:
                    proc_info = proc.info
                    cmdline_str = ' '.join(proc_info['cmdline']) if proc_info['cmdline'] else ''
                    
                    # Check if this is a pt_data_worker process
                    # PyTorch DataLoader workers show up as 'pt_data_worker' in process name or cmdline
                    is_pt_worker = (
                        proc_info['name'] == 'pt_data_worker' or
                        'pt_data_worker' in cmdline_str
                    )
                    
                    if is_pt_worker:
                        # Check if it's a child of a Python process (likely from previous training)
                        try:
                            parent = psutil.Process(proc_info['ppid'])
                            parent_name = parent.name()
                            parent_cmdline = ' '.join(parent.cmdline()) if parent.cmdline() else ''
                            
                            # Kill if parent is Python (likely from previous training run)
                            if 'python' in parent_name.lower() or 'python' in parent_cmdline.lower():
                                logger.info(f"   Killing dangling pt_data_worker PID {proc_info['pid']} (parent: {parent.pid} - {parent_name})")
                                proc.kill()
                                killed_count += 1
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            # Parent already gone or can't access - kill the orphaned worker anyway
                            logger.info(f"   Killing orphaned pt_data_worker PID {proc_info['pid']} (parent gone)")
                            try:
                                proc.kill()
                                killed_count += 1
                            except (psutil.NoSuchProcess, psutil.AccessDenied):
                                # Process already gone - skip
                                pass
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    # Process already gone or can't access - skip
                    continue
            
            if killed_count > 0:
                logger.info(f"✅ Cleaned up {killed_count} dangling pt_data_worker process(es)")
            else:
                logger.info("✅ No dangling pt_data_worker processes found")
        except ImportError:
            logger.warning("⚠️  psutil not available - cannot clean up pt_data_worker processes")
        except Exception as e:
            logger.warning(f"⚠️  Error cleaning up pt_data_worker processes: {e}")
    
    # Check for interrupted jobs when starting up
    logger.info(f"🔍 {queue_name} worker starting - checking for interrupted jobs...")
    interrupted_jobs = _recover_queue_interrupted_jobs(queue_name)
    
    if interrupted_jobs:
        logger.info(f"✅ {queue_name} recovered {len(interrupted_jobs)} interrupted jobs")
    else:
        logger.info(f"✅ {queue_name} startup clean - no interrupted jobs found")

    
    class MyHandler(FileSystemEventHandler):

        def _process_job_file(self, job_file_path: Path):
            """Helper to process a job file - checks if it's READY and adds to queue."""
            try:
                if not job_file_path.exists():
                    return
                
                if job_file_path.suffix != ".job":
                    return
                
                job_id = job_file_path.stem
                job = load_job(queue_name, job_id)

                is_ready = job.get("status") == JobStatus.READY

                if not is_ready:
                    return
                
                logger.info("""
***********************************************
New job detected
***********************************************
""")
                logger.info(f"   Job ID: {job_id}")
                logger.info(f"   Queue: {queue_name}")
                logger.info(f"   Status: {job.get('status')}")

                job_queue.put(job_id)
            except Exception as e:
                logger.error(f"Error processing job file {job_file_path}: {e}")
                logger.error(f"Traceback: {traceback.format_exc()}")

        def on_moved(self, event):
            # Don't log debug messages here - can cause "write to closed file" errors
            # when log rotation happens during job execution
            # logger.debug(f"moved {event.src_path} to {event.dest_path}")

            src_file = Path(event.src_path)
            dest_file = Path(event.dest_path)

            is_job_file = (src_file.suffix == ".tmp") and (dest_file.suffix == ".job")

            if not is_job_file:
                return
            
            self._process_job_file(dest_file)

        def on_created(self, event):
            """Handle job files created directly (shouldn't happen but catch it anyway)."""
            if event.is_directory:
                return
            
            job_file = Path(event.src_path)
            if job_file.suffix == ".job":
                logger.info(f"🔍 on_created: Detected new .job file: {job_file.name}")
                self._process_job_file(job_file)


    # get all pending jobs in the queue
    # Bail is too many jobs to catch up?
    ready_jobs = get_ready_jobs(queue_name)
    
    # execute the  pending jobs
    if skip_existing:
        logger.info("Skipping existing jobs")
    else:
        for job in ready_jobs:
            job_id = job["job_id"]

            job_queue.put(job_id)

    queue_path = config.queue_dir / queue_name

    # watch the queue for new jobs
    event_handler = MyHandler()
    observer = Observer()
    observer.schedule(event_handler, queue_path, recursive=False)
    observer.start()

    # HACK: For train_es, monitor for stuck jobs and restart if needed
    if queue_name == "train_es":
        restart_flag_file = Path("/tmp/train_es_restart_flag")
        first_restart_time_file = Path("/tmp/train_es_first_restart_time")
        
        def monitor_stuck_jobs():
            """Monitor for new jobs that don't start within 60 seconds and trigger restart."""
            last_job_start_time = {}  # Track when each job actually started running
            new_job_detected_time = {}  # Track when new jobs were detected
            restart_time_written = False
            
            while True:
                try:
                    time.sleep(5)  # Check every 5 seconds
                    
                    # Check for new job files in the queue directory
                    today = datetime.now().date()
                    for job_file in queue_path.glob("*.job"):
                        job_id = job_file.stem
                        job_mtime = datetime.fromtimestamp(job_file.stat().st_mtime)
                        job_date = job_mtime.date()
                        
                        # Only consider jobs created today
                        if job_date != today:
                            continue
                        
                        # Check if job is READY (new job waiting)
                        try:
                            job = load_job(queue_name, job_id)
                            if job.get("status") == JobStatus.READY:
                                # New ready job detected
                                if job_id not in new_job_detected_time:
                                    new_job_detected_time[job_id] = time.time()
                                    logger.info(f"🔍 HACK: New train_es job {job_id} detected (created today)")
                            elif job.get("status") == JobStatus.RUNNING:
                                # Job started running
                                if job_id not in last_job_start_time:
                                    last_job_start_time[job_id] = time.time()
                                    logger.info(f"✅ HACK: train_es job {job_id} started running")
                                    # Clear detection time since it started
                                    new_job_detected_time.pop(job_id, None)
                        except Exception as e:
                            # Skip if we can't load the job
                            continue
                    
                    # Check if any new job has been waiting > 60 seconds without starting
                    current_time = time.time()
                    for job_id, detected_time in list(new_job_detected_time.items()):
                        wait_time = current_time - detected_time
                        if wait_time > 60:
                            # Job stuck - write restart time and trigger restart
                            if not restart_time_written:
                                first_restart_time = datetime.now().isoformat()
                                try:
                                    first_restart_time_file.write_text(first_restart_time)
                                    restart_flag_file.touch()
                                    logger.warning(f"🚨 HACK: train_es job {job_id} stuck for {wait_time:.0f}s - writing restart flag")
                                    restart_time_written = True
                                except Exception as e:
                                    logger.error(f"Failed to write restart flag: {e}")
                            
                            # Restart once per minute
                            if restart_flag_file.exists():
                                flag_mtime = restart_flag_file.stat().st_mtime
                                if current_time - flag_mtime > 60:
                                    logger.warning(f"🔄 HACK: Restarting train_es worker (stuck job {job_id})")
                                    restart_flag_file.touch()  # Update timestamp
                                    # Exit with code 0 to trigger supervisor restart
                                    sys.exit(0)
                    
                except Exception as e:
                    logger.error(f"Error in stuck job monitor: {e}")
                    time.sleep(10)  # Wait longer on error
        
        # Start the monitor thread
        monitor_thread = Thread(target=monitor_stuck_jobs, daemon=True)
        monitor_thread.start()
        logger.info("🔧 HACK: train_es stuck job monitor started")

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()
    observer.join()

    # Gracefully stop the worker thread(s)
    job_queue.put(None)
    if use_thread_pool:
        queue_processor.join()
        executor.shutdown(wait=True)
    else:
        worker_thread.join()


##############################################################################
# 
# SIMILARITY SEARCH
# 
##############################################################################


def find_closest_points(session_id: str, query_record: Dict[str, Any], k: int = 5) -> Dict[str, Any]:
    logger.debug(sys.path)
    pp = str(Path("./lib").resolve())
    if pp not in sys.path:
        sys.path.insert(0, pp)
    from lib.utils import load_embedded_space
    from lib.vector_db import CSVtoLanceDB

    session = load_session(session_id)

    # validate faiss index path
    vector_db_path = session.get("vector_db")
    if vector_db_path is None:
        raise ValueError(f"Session {session_id} does not have a vector_db")
    
    vector_db_path = Path(vector_db_path)
    if not vector_db_path.is_dir():
        raise FileNotFoundError(f"LanceDB data at {vector_db_path} does not exist")


    # validate embedding space path
    embedding_space_path = session.get("embedding_space")
    if embedding_space_path is None:
        raise ValueError(f"Session {session_id} does not have an embedding space")

    embedding_space_path = Path(embedding_space_path)
    if not embedding_space_path.is_file():
        raise FileNotFoundError(f"Embedding space {embedding_space_path} does not exist")

    # load the ebmedding space
    try: 
        embedded_space = load_embedded_space(embedding_space_path)
    except Exception as e:
        raise ValueError(f"Error loading embedding space: {e}")


    # validate sqlite_db_path
    sqlite_db_path = session.get("sqlite_db")
    if sqlite_db_path is None:
        raise ValueError(f"Session {session_id} does not have an sqlite db")
    
    sqlite_db_path = Path(sqlite_db_path)
    if not sqlite_db_path.is_file():
        raise FileNotFoundError(f"Sqlite db at {sqlite_db_path} does not exist")


    # load the faiss index
    try:
        vec_db = CSVtoLanceDB(
            featrix_es=embedded_space, 
            sqlite_db_path=sqlite_db_path, 
            lancedb_path=vector_db_path)
        vec_db.load_existing()
    except Exception as e:
        raise ValueError(f"Error loading faiss index: {e}")

    # get the nearest neighbors
    try:
        result = vec_db.search(query_record, k=k)
    except Exception as e:
        raise ValueError(f"Error querying nearest neighbors: {e}")

    # result shape:
    # {
    #     "results": [
    #         {
    #             coord: [0., 0.2, 0.4],
    #             __featrix_row_offset: 1234,
    #             distance: 0.2
    #             original: {
    #                 ...  # original key/value pairs
    #             }
    #         }
    #     ],
    #     "stats": {
    #         "query_time_seconds": 0.123456,
    #         "records_searched": 10000,
    #         "brute_force_comparisons": 1280000,
    #         "vector_dimensions": 128,
    #         "results_returned": 5,
    #         "avg_distance": 0.456789,
    #         "std_distance": 0.123456,
    #         "min_distance": 0.123456,
    #         "max_distance": 0.789012
    #     }
    # }

    # return the closest points, their embeddings, and query statistics
    # Add server metadata for debugging    
    result["_meta"] = {
        "compute_cluster": socket.gethostname(),
        "compute_cluster_time": datetime.utcnow().isoformat() + "Z",
    }
    
    return result


def recover_interrupted_jobs():
    """
    Smart job recovery that distinguishes between planned upgrades and crashes.
    
    Behavior:
    - Planned upgrades (UPGRADE_SPHERE flag exists): Always recover jobs
    - Crashes: Limit retries to prevent infinite loops
    - Development: Always recover (based on hostname/environment detection)
    
    Returns:
        dict: Summary of recovery actions taken
    """
    recovery_summary = {
        "restart_reason": None,
        "queues_checked": [],
        "jobs_recovered": [],
        "train_es_resumed_from_checkpoint": [],
        "jobs_restarted": [],
        "jobs_blocked_by_retry_limit": [],
        "total_recovered": 0
    }
    
    # Determine restart reason
    upgrade_flag_path = Path("/tmp/UPGRADE_SPHERE")
    restart_reason = _determine_restart_reason(upgrade_flag_path)
    recovery_summary["restart_reason"] = restart_reason
    
    logger.info(f"🔍 Server restart detected - Reason: {restart_reason}")
    
    # Clean up upgrade flag if it exists
    if upgrade_flag_path.exists():
        try:
            with open(upgrade_flag_path, 'r') as f:
                flag_content = f.read().strip()
            logger.info(f"📋 Upgrade flag contents:\n{flag_content}")
            upgrade_flag_path.unlink()  # Remove the flag
            logger.info("🗑️  Upgrade flag cleaned up")
        except Exception as e:
            logger.warning(f"Failed to read/remove upgrade flag: {e}")
    
    # Check all known queue types
    known_queues = [
        "create_structured_data",
        "train_es", 
        "train_knn",
        "build_es_projections",
        "train_single_predictor"
    ]
    
    for queue_name in known_queues:
        queue_path = config.queue_dir / queue_name
        recovery_summary["queues_checked"].append(queue_name)
        
        if not queue_path.exists():
            logger.debug(f"Queue {queue_name} does not exist, skipping recovery check")
            continue
            
        interrupted_jobs = []
        
        try:
            # Find all jobs with RUNNING status or recently FAILED status
            for job in iterate_over_jobs_in_queue(queue_name):
                job_id = job.get("job_id")
                status = job.get("status")
                
                # Always include RUNNING jobs
                if status == JobStatus.RUNNING:
                    interrupted_jobs.append(job_id)
                    
                # Also include FAILED jobs that failed within the last hour
                elif status == JobStatus.FAILED:
                    if _should_retry_failed_job(job):
                        interrupted_jobs.append(job_id)
                        logger.info(f"🔄 Including recently failed job {job_id} for retry")
                    else:
                        continue  # Skip old failed jobs
                        
                else:
                    continue  # Skip other statuses
                
                # Check if job should be recovered based on retry limits
                should_recover = _should_recover_job(job, restart_reason)
                
                if not should_recover:
                    recovery_summary["jobs_blocked_by_retry_limit"].append({
                        "queue": queue_name, 
                        "job_id": job_id,
                        "reason": "retry_limit_exceeded"
                    })
                    _mark_job_as_failed_due_to_retry_limit(queue_name, job_id, job)
                    logger.warning(f"🚫 Job {job_id} exceeded retry limit - marked as FAILED")
                    continue
                
                # Special case: train_es job with RUNNING status but embedding space exists
                if queue_name == "train_es" and status == JobStatus.RUNNING:
                    # Check if the job actually completed successfully
                    session_id = job.get("session_id")
                    if session_id:
                        try:
                            session = load_session(session_id)
                            embedding_space_path = session.get("embedding_space")
                            if embedding_space_path and Path(embedding_space_path).exists():
                                logger.info(f"✅ train_es job {job_id} has RUNNING status but embedding space exists - marking as DONE")
                                update_job_status(queue_name="train_es", job_id=job_id, status=JobStatus.DONE)
                                logger.info(f"✅ Updated job {job_id} status to DONE - no recovery needed")
                                continue  # Skip recovery for this job
                        except Exception as e:
                            logger.warning(f"Could not check session for job {job_id}: {e}")
                
                # Special handling for train_es jobs - try to resume from checkpoint
                if queue_name == "train_es":
                    resume_result = _recover_train_es_job(job_id, job, restart_reason)
                    if resume_result["resumed"]:
                        recovery_summary["train_es_resumed_from_checkpoint"].append({
                            "job_id": job_id,
                            "resume_epoch": resume_result["resume_epoch"],
                            "checkpoint_path": resume_result["checkpoint_path"]
                        })
                        logger.info(f"🎯 train_es job {job_id} will resume from epoch {resume_result['resume_epoch']}")
                    else:
                        recovery_summary["jobs_restarted"].append({"queue": queue_name, "job_id": job_id})
                        logger.info(f"🔄 train_es job {job_id} will restart from beginning")
                else:
                    # For other job types, just reset to READY
                    _reset_job_to_ready(queue_name, job_id, job, restart_reason)
                    recovery_summary["jobs_restarted"].append({"queue": queue_name, "job_id": job_id})
                    logger.info(f"🔄 {queue_name} job {job_id} reset to READY for retry")
                    
        except Exception as e:
            logger.error(f"Error recovering jobs from queue {queue_name}: {e}")
            continue
            
        if interrupted_jobs:
            recovery_summary["jobs_recovered"].extend([
                {"queue": queue_name, "job_id": job_id} for job_id in interrupted_jobs
            ])
            logger.info(f"✅ Processed {len(interrupted_jobs)} interrupted jobs from {queue_name}")
        else:
            logger.debug(f"No interrupted jobs found in {queue_name}")
    
    recovery_summary["total_recovered"] = len(recovery_summary["jobs_recovered"])
    blocked_count = len(recovery_summary["jobs_blocked_by_retry_limit"])
    
    if recovery_summary["total_recovered"] > 0:
        resumed_count = len(recovery_summary["train_es_resumed_from_checkpoint"])
        restarted_count = len(recovery_summary["jobs_restarted"])
        
        logger.info(f"🎯 Job recovery complete: {recovery_summary['total_recovered']} jobs processed")
        if resumed_count > 0:
            logger.info(f"   📁 {resumed_count} train_es jobs will resume from checkpoints")
        if restarted_count > 0:
            logger.info(f"   🔄 {restarted_count} jobs will restart from beginning")
        if blocked_count > 0:
            logger.warning(f"   🚫 {blocked_count} jobs blocked by retry limits")
            
        # Send appropriate Slack notification
        if restart_reason == "planned_upgrade":
            send_slack_message(f"🔄 Planned upgrade complete: recovered {recovery_summary['total_recovered']} jobs ({resumed_count} resumed, {restarted_count} restarted)")
        elif restart_reason == "development":
            send_slack_message(f"🔄 Dev restart: recovered {recovery_summary['total_recovered']} jobs ({resumed_count} resumed, {restarted_count} restarted)")
        else:
            send_slack_message(f"⚠️ Server crash recovery: {recovery_summary['total_recovered']} jobs ({resumed_count} resumed, {restarted_count} restarted, {blocked_count} blocked)")
    else:
        logger.info("✅ No interrupted jobs found - clean server restart")
    
    return recovery_summary


def _determine_restart_reason(upgrade_flag_path: Path) -> str:
    """
    Determine why the server restarted.
    
    Returns:
        str: "planned_upgrade", "development", or "production_crash"
    """
    # Check for planned upgrade flag
    if upgrade_flag_path.exists():
        return "planned_upgrade"
    
    # Check if we're in development environment
    hostname = socket.gethostname()
    dev_indicators = [
        hostname.endswith('.local'),
        'dev' in hostname.lower(),
        'test' in hostname.lower(),
        hostname in ['localhost', '127.0.0.1'],
        hostname.startswith('MacBook'),
        hostname.startswith('admin'),  # Common dev machine names
    ]
    
    if any(dev_indicators):
        return "development"
    
    return "production_crash"


def _should_retry_failed_job(job: dict) -> bool:
    """
    Determine if a FAILED job should be retried based on when it failed.
    
    Args:
        job: Job dictionary with FAILED status
        
    Returns:
        bool: True if job failed within the last hour and should be retried
    """
    try:
        finished_at = job.get("finished_at")
        if not finished_at:
            return False
            
        # Convert to datetime if it's a string
        if isinstance(finished_at, str):
            finished_at = convert_from_iso(finished_at)
        elif not isinstance(finished_at, datetime):
            return False
            
        # Check if it failed within the last hour
        now = datetime.now(tz=ZoneInfo("America/New_York"))
        time_since_failure = now - finished_at
        
        # Retry if failed within last hour (3600 seconds)
        if time_since_failure.total_seconds() < 3600:
            logger.info(f"Job failed {time_since_failure.total_seconds():.0f} seconds ago - eligible for retry")
            return True
        else:
            logger.debug(f"Job failed {time_since_failure.total_seconds():.0f} seconds ago - too old for retry")
            return False
            
    except Exception as e:
        logger.error(f"Error checking failed job retry eligibility: {e}")
        return False


def _should_recover_job(job: dict, restart_reason: str) -> bool:
    """
    Determine if a job should be recovered based on retry limits and restart reason.
    
    Args:
        job: Job dictionary
        restart_reason: "planned_upgrade", "development", or "production_crash"
        
    Returns:
        bool: True if job should be recovered
    """
    # CRITICAL: Never recover manually killed jobs
    job_spec = job.get("spec", {})
    if job_spec.get("manual_kill", False):
        logger.info(f"Job {job.get('job_id')} was manually killed - will not recover")
        return False
    
    # Check for manual kill in recovery_info as well
    recovery_info = job.get("recovery_info", [])
    for recovery in recovery_info:
        if recovery.get("manual_kill", False) or recovery.get("reason") == "manual_kill":
            logger.info(f"Job {job.get('job_id')} has manual_kill in recovery history - will not recover")
            return False
    
    # Always recover for planned upgrades and development
    if restart_reason in ["planned_upgrade", "development"]:
        return True
    
    # For production crashes, check retry limits
    now = datetime.now(tz=ZoneInfo("America/New_York"))
    
    # Count recent recovery attempts (last 24 hours)
    recent_recoveries = []
    for recovery in recovery_info:
        try:
            recovery_time = recovery.get("recovered_at")
            if isinstance(recovery_time, str):
                recovery_time = convert_from_iso(recovery_time)
            elif isinstance(recovery_time, datetime):
                pass  # Already datetime
            else:
                continue  # Skip invalid entries
                
            time_diff = now - recovery_time
            if time_diff.total_seconds() < 86400:  # 24 hours
                recent_recoveries.append(recovery)
        except Exception:
            continue  # Skip invalid recovery entries
    
    # Limits for production crashes
    max_retries_24h = 100  # Max 100 retries in 24 hours
    max_retries_1h = 100   # Max 100 retries in 1 hour
    
    # Count retries in last hour
    recent_1h_recoveries = [
        r for r in recent_recoveries 
        if (now - convert_from_iso(r.get("recovered_at", "1970-01-01"))).total_seconds() < 3600
    ]
    
    if len(recent_1h_recoveries) >= max_retries_1h:
        logger.warning(f"Job {job.get('job_id')} hit 1-hour retry limit: {len(recent_1h_recoveries)}/{max_retries_1h}")
        return False
        
    if len(recent_recoveries) >= max_retries_24h:
        logger.warning(f"Job {job.get('job_id')} hit 24-hour retry limit: {len(recent_recoveries)}/{max_retries_24h}")
        return False
    
    return True


def _count_recent_restarts(job: dict, hours: int = 1) -> int:
    """
    Count how many times a job has restarted in the last N hours.
    
    Args:
        job: Job dictionary
        hours: Number of hours to look back (default: 1)
        
    Returns:
        Number of restarts in the last N hours
    """
    recovery_info = job.get("recovery_info", [])
    if not recovery_info:
        return 0
    
    now = datetime.now(tz=ZoneInfo("America/New_York"))
    cutoff_seconds = hours * 3600
    
    count = 0
    for recovery in recovery_info:
        try:
            recovery_time = recovery.get("recovered_at")
            if isinstance(recovery_time, str):
                recovery_time = convert_from_iso(recovery_time)
            elif not isinstance(recovery_time, datetime):
                continue
            
            time_diff = now - recovery_time
            if time_diff.total_seconds() < cutoff_seconds:
                count += 1
        except Exception:
            continue
    
    return count


def _mark_job_as_failed_due_to_retry_limit(queue_name: str, job_id: str, job: dict):
    """Mark a job as failed due to retry limit exceeded."""
    try:
        job_content = job.copy()
        job_content["status"] = JobStatus.FAILED
        job_content["finished_at"] = datetime.now(tz=ZoneInfo("America/New_York"))
        
        # Add failure reason to spec
        job_spec = job_content.get("spec", {}).copy()
        job_spec["failure_reason"] = "retry_limit_exceeded"
        job_spec["failure_details"] = "Job failed too many times after server restarts - possible recurring issue"
        job_content["spec"] = job_spec
        
        save_job(queue_name, job_id, job_content, exist_ok=True)
        
    except Exception as e:
        logger.error(f"Failed to mark job {job_id} as failed: {e}")


def _mark_job_as_stale(queue_name: str, job_id: str, job: dict):
    """Mark an old job as failed because it's stale (newer job was running)."""
    try:
        job_content = job.copy()
        job_content["status"] = JobStatus.FAILED
        job_content["finished_at"] = datetime.now(tz=ZoneInfo("America/New_York"))
        
        # Add failure reason to spec
        job_spec = job_content.get("spec", {}).copy()
        job_spec["failure_reason"] = "stale_job"
        job_spec["failure_details"] = "Job was superseded by a newer job in the same queue"
        job_content["spec"] = job_spec
        
        save_job(queue_name, job_id, job_content, exist_ok=True)
        
    except Exception as e:
        logger.error(f"Failed to mark job {job_id} as stale: {e}")


def _recover_train_es_job(job_id: str, job: dict, restart_reason: str) -> dict:
    """
    Try to recover a train_es job by resuming from the latest checkpoint.
    
    Returns:
        dict: {"resumed": bool, "resume_epoch": int, "checkpoint_path": str}
    """
    result = {"resumed": False, "resume_epoch": None, "checkpoint_path": None}
    
    try:
        # Check for work directory using helper function (handles new structure)
        work_dir = get_job_output_path(job_id)
        if not work_dir.exists():
            logger.info(f"No work directory found for {job_id} - will restart from beginning")
            _reset_job_to_ready("train_es", job_id, job, restart_reason)
            return result
        
        # CRITICAL FIX: Check for ABORT file before resuming
        abort_file = work_dir / "ABORT"
        if abort_file.exists():
            logger.warning(f"🚫 ABORT file detected for job {job_id} - marking as FAILED instead of resuming")
            job_content = job.copy()
            job_content["status"] = JobStatus.FAILED
            job_content["finished_at"] = datetime.now(tz=ZoneInfo("America/New_York"))
            job_content["error_message"] = "Job aborted - ABORT file detected during recovery"
            save_job("train_es", job_id, job_content, exist_ok=True)
            
            # Delete the job file so a new job can be created
            try:
                job_file_path = config.queue_dir / "train_es" / f"{job_id}.json"
                if job_file_path.exists():
                    job_file_path.unlink()
                    logger.info(f"🗑️ Deleted job file {job_file_path} to allow new job creation")
            except Exception as delete_err:
                logger.warning(f"Failed to delete job file: {delete_err}")
            
            return result
        
        # Check for FINISH file before resuming
        finish_file = work_dir / "FINISH"
        if finish_file.exists():
            # FINISH flag exists - but check if job actually finished successfully
            finished_file = work_dir / "FINISHED"
            
            if finished_file.exists():
                logger.warning(f"🏁 FINISH flag detected for job {job_id} - FINISHED flag exists, job completed successfully")
                logger.info(f"🏁 Job {job_id} completed gracefully - not resuming")
                # Job finished successfully after FINISH flag - don't resume
                return result
            else:
                logger.warning(f"🏁 FINISH flag detected for job {job_id} BUT FINISHED flag does NOT exist")
                logger.warning(f"🏁 Job may have crashed before completing - will resume to complete and save")
                logger.info(f"🏁 Once resumed, job will detect FINISH flag and complete gracefully")
                # Let the resume continue - job will detect FINISH flag, complete, and save model
        
        # Look for training state checkpoints: training_state_e-{epoch}.pth
        checkpoint_files = list(work_dir.glob("training_state_e-*.pth"))
        
        if not checkpoint_files:
            logger.info(f"No training checkpoints found for {job_id} - will restart from beginning")
            _reset_job_to_ready("train_es", job_id, job, restart_reason)
            return result
        
        # Find the latest checkpoint by epoch number
        latest_epoch = -1
        latest_checkpoint = None
        
        for checkpoint_file in checkpoint_files:
            # Extract epoch from filename like "training_state_e-42.pth"
            match = re.search(r'training_state_e-(\d+)\.pth', checkpoint_file.name)
            if match:
                epoch = int(match.group(1))
                if epoch > latest_epoch:
                    latest_epoch = epoch
                    latest_checkpoint = checkpoint_file
        
        if latest_checkpoint and latest_epoch >= 0:
            logger.info(f"📁 Found latest checkpoint: {latest_checkpoint.name} (epoch {latest_epoch})")
            
            # Update job spec to resume from this checkpoint
            job_spec = job.get("spec", {}).copy()
            job_spec["resume_from_epoch"] = latest_epoch
            job_spec["resume_from_checkpoint"] = str(latest_checkpoint)
            
            # Reset job to READY with resume information
            job_content = job.copy()
            job_content["status"] = JobStatus.READY.value  # Convert enum to string for JSON serialization
            job_content["started_at"] = None
            job_content["spec"] = job_spec
            
            # Add recovery metadata
            now = datetime.now(tz=ZoneInfo("America/New_York"))
            if "recovery_info" not in job_content:
                job_content["recovery_info"] = []
            
            version_info = get_version_info()
            job_content["recovery_info"].append({
                "recovered_at": convert_to_iso(now),
                "reason": f"server_restart_{restart_reason}_resume_checkpoint",
                "previous_status": "running",
                "restart_reason": restart_reason,
                "resume_epoch": latest_epoch,
                "checkpoint_file": latest_checkpoint.name,
                "version": version_info,
            })
            
            save_job("train_es", job_id, job_content, exist_ok=True)
            
            result["resumed"] = True
            result["resume_epoch"] = latest_epoch
            result["checkpoint_path"] = str(latest_checkpoint)
            
            logger.info(f"✅ train_es job {job_id} prepared to resume from epoch {latest_epoch}")
            
        else:
            logger.warning(f"Checkpoint files found but couldn't parse epoch - restarting from beginning")
            _reset_job_to_ready("train_es", job_id, job, restart_reason)
            
    except Exception as e:
        logger.error(f"Error recovering train_es job {job_id}: {e}")
        _reset_job_to_ready("train_es", job_id, job, restart_reason)
    
    return result


def _scan_for_restart_flags(queue_name: str) -> list:
    """
    Scan job directories for RESTART flags, sorted by newest first.
    
    When a RESTART flag is found:
    - Reset the job to READY status (or create new job if it doesn't exist)
    - Clear all recovery info and progress
    - Rename RESTART to started.RESTART.<date> to prevent re-processing
    - Queue the job for restart
    
    Args:
        queue_name: Name of the queue to scan
        
    Returns:
        list: Job IDs that were queued for restart (sorted by newest first)
    """
    restarted_jobs = []
    
    try:
        from featrix.neural.embedded_space import check_restart_files
        
        # Scan job directories for RESTART flags
        job_dirs = []
        
        # Check session-based structure: featrix_output/{session_id}/{job_type}/{job_id}
        session_dirs = config.output_dir.glob("*")
        for session_dir in session_dirs:
            if session_dir.is_dir():
                job_type_dir = session_dir / queue_name
                if job_type_dir.exists():
                    job_dirs.extend(job_type_dir.iterdir())
        
        # Also check direct structure: featrix_output/{queue_name}/{job_id}
        queue_base = config.output_dir / queue_name
        if queue_base.exists():
            job_dirs.extend(queue_base.iterdir())
        
        # Also check old structure: featrix_output/{job_id}
        for job_dir in config.output_dir.iterdir():
            if job_dir.is_dir() and job_dir.name not in [d.name for d in session_dirs]:
                # Check if this directory has a RESTART file
                if (job_dir / "RESTART").exists():
                    job_dirs.append(job_dir)
        
        # Find all RESTART files with their modification times
        restart_candidates = []
        for job_dir in job_dirs:
            if not job_dir.is_dir():
                continue
            
            restart_file = job_dir / "RESTART"
            if restart_file.exists():
                # Check if there's a started.RESTART.* file from a previous restart
                started_restart_files = list(job_dir.glob("started.RESTART.*"))
                if started_restart_files:
                    # If RESTART file is newer than the most recent started.RESTART file, process it
                    # This allows users to create a new RESTART file after a previous restart was processed
                    restart_mtime = restart_file.stat().st_mtime
                    latest_started_mtime = max(f.stat().st_mtime for f in started_restart_files)
                    if restart_mtime <= latest_started_mtime:
                        # RESTART file is older than or same age as started.RESTART - skip it
                        logger.debug(f"Skipping {job_dir.name} - RESTART file is older than started.RESTART")
                        continue
                    # RESTART file is newer - remove old started.RESTART files and process the new one
                    logger.info(f"🔄 Found newer RESTART file for {job_dir.name} (removing old started.RESTART files)")
                    for old_file in started_restart_files:
                        try:
                            old_file.unlink()
                            logger.debug(f"   Removed old {old_file.name}")
                        except Exception as e:
                            logger.warning(f"   Failed to remove {old_file.name}: {e}")
                
                # Get modification time for sorting
                mtime = restart_file.stat().st_mtime
                job_id = job_dir.name
                restart_candidates.append((mtime, job_id, restart_file, job_dir))
        
        # Sort by newest first (highest mtime first)
        restart_candidates.sort(key=lambda x: x[0], reverse=True)
        
        if not restart_candidates:
            logger.debug(f"No RESTART flags found in {queue_name} directories")
            return restarted_jobs
        
        logger.info(f"🔄 Found {len(restart_candidates)} RESTART flags in {queue_name}, processing newest first...")
        
        now = datetime.now(tz=ZoneInfo("America/New_York"))
        date_str = now.strftime("%Y%m%d-%H%M%S")
        
        for mtime, job_id, restart_file, job_dir in restart_candidates:
            try:
                logger.info(f"🔄 Processing RESTART flag for job {job_id}")
                
                # Try to load existing job
                job = None
                try:
                    job = load_job(queue_name, job_id)
                except FileNotFoundError:
                    # Job doesn't exist in queue - skip it (user needs to recreate manually)
                    logger.warning(f"   ⚠️  Skipping {job_id} - job not found in queue (may need manual recreation)")
                    # Still rename the RESTART file so it doesn't get processed again
                    started_restart_file = job_dir / f"started.RESTART.{date_str}"
                    restart_file.rename(started_restart_file)
                    logger.info(f"   ✅ Renamed RESTART to {started_restart_file.name}")
                    continue
                
                # Reset job to READY status
                job_content = job.copy()
                job_content["status"] = JobStatus.READY.value  # Convert enum to string for JSON serialization
                job_content["started_at"] = None
                job_content["finished_at"] = None
                job_content["progress"] = None
                job_content["error_message"] = None
                
                # Clear any existing resume information
                job_spec = job_content.get("spec", {}).copy()
                job_spec.pop("resume_from_epoch", None)
                job_spec.pop("resume_checkpoint_path", None)
                job_spec.pop("resume_from_checkpoint", None)
                job_spec.pop("manual_kill", None)
                job_content["spec"] = job_spec
                
                # Clear recovery info
                job_content["recovery_info"] = []
                
                # Add restart metadata
                previous_status = job.get("status")
                # Convert JobStatus enum to string if needed
                if hasattr(previous_status, "value"):
                    previous_status = previous_status.value
                version_info = get_version_info()
                job_content["recovery_info"].append({
                    "recovered_at": convert_to_iso(now),
                    "reason": "manual_restart_via_restart_flag",
                    "previous_status": previous_status,
                    "restart_reason": "RESTART flag detected",
                    "version": version_info,
                })
                
                # Save the updated job
                save_job(queue_name, job_id, job_content, exist_ok=True)
                
                # Rename RESTART to started.RESTART.<date> to prevent re-processing
                started_restart_file = job_dir / f"started.RESTART.{date_str}"
                restart_file.rename(started_restart_file)
                
                restarted_jobs.append(job_id)
                logger.info(f"✅ Job {job_id} reset to READY and queued for restart")
                logger.info(f"   RESTART flag renamed to {started_restart_file.name}")
                
            except Exception as e:
                logger.error(f"Error processing RESTART flag for job {job_id}: {e}")
                traceback.print_exc()
                # Continue with next job
                continue
                
    except Exception as e:
        logger.error(f"Error scanning for RESTART flags in {queue_name}: {e}")
        traceback.print_exc()
    
    return restarted_jobs


def _reset_job_to_ready(queue_name: str, job_id: str, job: dict, restart_reason: str):
    """Reset a job to READY status for retry from beginning."""
    try:
        # Check for ABORT file - if it exists, mark job as FAILED instead of READY
        if queue_name == "train_es":
            try:
                from featrix.neural.embedded_space import check_abort_files
                if check_abort_files(job_id):
                    logger.warning(f"🚫 Job {job_id} has ABORT file - marking as FAILED instead of READY")
                    job_content = job.copy()
                    job_content["status"] = JobStatus.FAILED.value  # Convert enum to string for JSON serialization
                    job_content["finished_at"] = datetime.now(tz=ZoneInfo("America/New_York"))
                    job_content["error_message"] = "Job aborted due to ABORT file"
                    
                    # Add recovery metadata
                    now = datetime.now(tz=ZoneInfo("America/New_York"))
                    if "recovery_info" not in job_content:
                        job_content["recovery_info"] = []
                    
                    version_info = get_version_info()
                    job_content["recovery_info"].append({
                        "recovered_at": convert_to_iso(now),
                        "reason": "aborted_due_to_abort_file",
                        "previous_status": "running",
                        "restart_reason": restart_reason,
                        "version": version_info,
                    })
                    
                    save_job(queue_name, job_id, job_content, exist_ok=True)
                    return
            except Exception as e:
                logger.warning(f"Could not check ABORT file for job {job_id}: {e}")
                # Continue with normal recovery if check fails
        
        job_content = job.copy()
        job_content["status"] = JobStatus.READY.value  # Convert enum to string for JSON serialization
        job_content["started_at"] = None
        job_content["finished_at"] = None
        job_content["progress"] = None
        
        # Clear any existing resume information
        job_spec = job_content.get("spec", {}).copy()
        job_spec.pop("resume_from_epoch", None)
        job_spec.pop("resume_from_checkpoint", None)
        job_content["spec"] = job_spec
        
        # Add recovery metadata with restart reason
        now = datetime.now(tz=ZoneInfo("America/New_York"))
        if "recovery_info" not in job_content:
            job_content["recovery_info"] = []
        
        version_info = get_version_info()
        job_content["recovery_info"].append({
            "recovered_at": convert_to_iso(now),
            "reason": f"server_restart_{restart_reason}",
            "previous_status": "running",
            "restart_reason": restart_reason,
            "version": version_info,
        })
        
        # Check if job should be paused due to too many restarts
        restart_count = _count_recent_restarts(job_content, hours=1)
        if restart_count >= 20:
            logger.warning(f"⏸️  Job {job_id} has restarted {restart_count} times in the last hour - pausing job")
            job_content['status'] = JobStatus.PAUSED.value  # Convert enum to string for JSON serialization
            job_content['pause_reason'] = f'Too many restarts: {restart_count} in last hour (limit: 20)'
            job_content['paused_at'] = convert_to_iso(now)
        # else: status is already set to READY above
        
        save_job(queue_name, job_id, job_content, exist_ok=True)
        
    except Exception as e:
        logger.error(f"Failed to reset job {job_id} to READY: {e}")


def _recover_queue_interrupted_jobs(queue_name: str) -> list:
    """
    Recover interrupted jobs for a specific queue.
    
    This is called by each queue worker on startup to handle any jobs
    that were running when the server was previously shut down.
    
    Also scans for RESTART flags on jobs from the last 2 days.
    
    Returns:
        list: Job IDs that were recovered
    """
    recovered_jobs = []
    
    try:
        # Determine restart reason
        upgrade_flag_path = Path("/tmp/UPGRADE_SPHERE")
        restart_reason = _determine_restart_reason(upgrade_flag_path)
        
        # Clean up upgrade flag on first queue to process it
        if upgrade_flag_path.exists():
            try:
                with open(upgrade_flag_path, 'r') as f:
                    flag_content = f.read().strip()
                logger.info(f"📋 Upgrade flag detected: {restart_reason}")
                logger.debug(f"Flag contents:\n{flag_content}")
                upgrade_flag_path.unlink()  # Remove the flag
            except Exception as e:
                logger.warning(f"Failed to read/remove upgrade flag: {e}")
        
        queue_path = config.queue_dir / queue_name
        if not queue_path.exists():
            logger.debug(f"Queue {queue_name} does not exist, no recovery needed")
            return recovered_jobs
        
        # FIRST: Scan for RESTART flags in job directories (sorted by newest first)
        restarted_jobs = _scan_for_restart_flags(queue_name)
        if restarted_jobs:
            logger.info(f"🔄 Found {len(restarted_jobs)} jobs with RESTART flag")
            for job_id in restarted_jobs:
                logger.info(f"   ↻ {job_id} will be restarted")
            # Add to recovered_jobs list
            recovered_jobs.extend(restarted_jobs)
        
        # THEN: Find jobs that need recovery: RUNNING jobs + recently FAILED jobs
        jobs_to_recover = []
        for job in iterate_over_jobs_in_queue(queue_name):
            status = job.get("status")
            
            # Always include RUNNING jobs (unless they were already restarted via RESTART flag)
            if status == JobStatus.RUNNING:
                jobs_to_recover.append(job)
                
            # Also include FAILED jobs that failed within the last hour
            elif status == JobStatus.FAILED:
                if _should_retry_failed_job(job):
                    jobs_to_recover.append(job)
                    logger.info(f"🔄 Including recently failed job {job.get('job_id')} for retry")
        
        if not jobs_to_recover:
            # Return early, but we still might have restarted some jobs
            return recovered_jobs
            
        # Sort by started_at to get the most recent
        # CRITICAL FIX: Normalize started_at to ISO strings to avoid datetime/string comparison errors
        def safe_started_at(job):
            started_at = job.get("started_at")
            if started_at is None:
                return "1970-01-01T00:00:00Z"  # Very old timestamp
            elif isinstance(started_at, str):
                return started_at  # Already an ISO string
            else:
                # Convert datetime to ISO string
                return convert_to_iso(started_at)
        
        jobs_to_recover.sort(key=safe_started_at, reverse=True)
        most_recent_job = jobs_to_recover[0]
        
        # Mark all other old jobs as failed (they're stale)
        for old_job in jobs_to_recover[1:]:
            old_job_id = old_job.get("job_id")
            _mark_job_as_stale(queue_name, old_job_id, old_job)
            logger.info(f"🗑️  Marked stale job {old_job_id} as FAILED")
        
        # Recover only the most recent job
        job = most_recent_job
        job_id = job.get("job_id")
        
        # Check retry limits
        should_recover = _should_recover_job(job, restart_reason)
        
        if not should_recover:
            _mark_job_as_failed_due_to_retry_limit(queue_name, job_id, job)
            logger.warning(f"🚫 {queue_name} job {job_id} exceeded retry limit - marked as FAILED")
            return recovered_jobs
        
        # Recover based on queue type
        if queue_name == "train_es":
            resume_result = _recover_train_es_job(job_id, job, restart_reason)
            if resume_result["resumed"]:
                logger.info(f"🎯 train_es job {job_id} will resume from epoch {resume_result['resume_epoch']}")
            else:
                logger.info(f"🔄 train_es job {job_id} will restart from beginning")
        else:
            _reset_job_to_ready(queue_name, job_id, job, restart_reason)
            logger.info(f"🔄 {queue_name} job {job_id} reset to READY for retry")
        
        recovered_jobs.append(job_id)
        
        # Send notification if jobs were recovered or restarted
        if recovered_jobs or restarted_jobs:
            total_actions = len(recovered_jobs)
            restart_count = len(restarted_jobs) if restarted_jobs else 0
            recovered_count = total_actions - restart_count
            
            if restart_reason == "planned_upgrade":
                msg_parts = []
                if recovered_count > 0:
                    msg_parts.append(f"recovered {recovered_count} interrupted jobs")
                if restart_count > 0:
                    msg_parts.append(f"restarted {restart_count} flagged jobs")
                if msg_parts:
                    send_slack_message(f"🔄 {queue_name} worker: {', '.join(msg_parts)} after planned upgrade")
            elif restart_reason == "development":
                # Skip Slack for dev restarts
                pass
            else:
                msg_parts = []
                if recovered_count > 0:
                    msg_parts.append(f"recovered {recovered_count} interrupted jobs")
                if restart_count > 0:
                    msg_parts.append(f"restarted {restart_count} flagged jobs")
                if msg_parts:                    
                    hostname = get_hostname()
                    version_info = get_version()
                    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    
                    slack_msg = f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n"
                    slack_msg += f"⚠️ *{queue_name} worker recovery on `{hostname}`*\n"
                    slack_msg += f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n\n"
                    slack_msg += f"```\n"
                    slack_msg += f"Version:   {version_info.semantic_version}\n"
                    slack_msg += f"Date:      {current_time}\n"
                    slack_msg += f"```\n\n"
                    slack_msg += f"```\n"
                    slack_msg += f"Field              Value\n"
                    slack_msg += f"─────────────────────────────────────\n"
                    slack_msg += f"Queue              {queue_name}\n"
                    slack_msg += f"Status             {', '.join(msg_parts)} after server crash\n"
                    slack_msg += f"```"
                    send_slack_message(slack_msg, skip_hostname_prefix=True)  # Hostname already in message
                
    except Exception as e:
        logger.error(f"Error recovering interrupted jobs for {queue_name}: {e}")
        traceback.print_exc()
    
    return recovered_jobs

