api/worker.py

"""
Background worker for processing document generation jobs using batched Claude API.
Runs as RQ worker process.
"""

import asyncio
import io
import json
import os
import pathlib
import tempfile
import time
import traceback
import zipfile
import shutil
from typing import Dict, Any, List, Callable
from datetime import datetime

# Add worker startup logging
from .config import settings

from .supabase_client import supabase_client
from .google_drive import GoogleDriveClient
from .utils import (
    download_seed_images,
    build_prompt,
    extract_html_documents_from_response,
    extract_ground_truth,
    extract_css_from_html,
    render_html_to_pdf,
    extract_bboxes_from_rendered_pdf,
    pdf_to_base64,
    process_stage3_complete,
    process_stage4_ocr,
    process_stage5_complete,
    validate_html_structure,
    validate_pdf,
    validate_bboxes
)
from docgenie.generation.pipeline_01.claude_batching import ClaudeBatchedClient
from docgenie import ENV


# ==================== Worker Logging Configuration ====================
# Read from environment variable, default to False for cleaner logs
VERBOSE_LOGGING = os.getenv('WORKER_VERBOSE_LOGGING', 'false').lower() in ('true', '1', 'yes')

def log_verbose(message: str):
    """Log message only if verbose logging is enabled"""
    if VERBOSE_LOGGING:
        print(message)


# ==================== Startup Validation ====================
def validate_worker_config():
    """Validate worker configuration at startup"""
    print("=" * 60)
    print("🔧 Worker Configuration Check")
    print("=" * 60)
    
    # Check Anthropic API
    if settings.ANTHROPIC_API_KEY:
        print("✓ ANTHROPIC_API_KEY: Set")
    else:
        print("✗ ANTHROPIC_API_KEY: NOT SET (REQUIRED)")
    
    # Check Supabase
    if settings.SUPABASE_URL and settings.SUPABASE_KEY:
        print(f"✓ SUPABASE: {settings.SUPABASE_URL[:30]}...")
    else:
        print("✗ SUPABASE: NOT SET (REQUIRED)")
    
    # Check Google OAuth (optional, for token refresh)
    if settings.GOOGLE_CLIENT_ID and settings.GOOGLE_CLIENT_SECRET:
        print(f"✓ GOOGLE_CLIENT_ID: {settings.GOOGLE_CLIENT_ID[:20]}...")
        print("✓ GOOGLE_CLIENT_SECRET: Set")
        print("  → Token auto-refresh: ENABLED")
    else:
        print("⚠ GOOGLE_CLIENT_ID/SECRET: Not set")
        print("  → Token auto-refresh: DISABLED")
        print("  → Users must provide fresh access tokens that don't expire during processing")
    
    print("=" * 60)

# Run validation on module import
validate_worker_config()


def retry_on_network_error(func: Callable, max_retries: int = 3, delay: float = 2.0) -> Any:
    """
    Retry a function on network errors with exponential backoff.
    
    Args:
        func: Function to execute (must be callable with no args)
        max_retries: Maximum number of retry attempts
        delay: Initial delay in seconds (doubles each retry)
    
    Returns:
        Result of the function call
    
    Raises:
        Last exception if all retries fail
    """
    last_exception = None
    for attempt in range(max_retries):
        try:
            return func()
        except Exception as e:
            last_exception = e
            error_str = str(e).lower()
            # Retry on network/DNS errors
            if any(err in error_str for err in ['name resolution', 'connection', 'timeout', 'network']):
                if attempt < max_retries - 1:
                    wait_time = delay * (2 ** attempt)
                    print(f"[Retry {attempt + 1}/{max_retries}] Network error, retrying in {wait_time}s: {e}")
                    time.sleep(wait_time)
                    continue
            # Non-network error or last attempt
            raise
    # All retries exhausted
    raise last_exception


async def process_document_generation_job_async(request_id: str, request_data: Dict[str, Any]):
    """
    Async background job function - processes document generation using batched Claude API.
    
    This function:
    1. Creates Claude batch with single message (generates N documents)
    2. Polls batch until completion
    3. Processes all documents (PDFs, handwriting, etc.)
    4. Uploads ZIP to user's Google Drive
    5. Updates Supabase with results
    
    Args:
        request_id: Document request UUID from Supabase
        request_data: Request parameters dict containing:
            - user_id: int
            - seed_images: List[str] (URLs)
            - prompt_params: Dict (language, doc_type, num_solutions, etc.)
    
    Raises:
        Exception: Any error during processing (logged to Supabase)
    """
    user_id = request_data['user_id']
    google_drive_token = request_data.get('google_drive_token')
    google_drive_refresh_token = request_data.get('google_drive_refresh_token')
    seed_image_urls = request_data['seed_images']
    prompt_params = request_data['prompt_params']
    
    # Validate Google Drive credentials configuration
    if google_drive_refresh_token:
        if not settings.GOOGLE_CLIENT_ID or not settings.GOOGLE_CLIENT_SECRET:
            print(f"[Job {request_id}] ⚠️ WARNING: refresh_token provided but GOOGLE_CLIENT_ID/SECRET not configured")
            print(f"[Job {request_id}] Token auto-refresh will fail. Ensure access token remains valid.")
    
    # Create temporary directories for this job
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = pathlib.Path(tmp_dir)
        batch_dir = tmp_path / "batches"
        message_dir = tmp_path / "messages"
        batch_dir.mkdir(exist_ok=True)
        message_dir.mkdir(exist_ok=True)
        
        # Initialize DatasetExporter for organized structure
        from .dataset_exporter import DatasetExporter
        exporter = DatasetExporter(tmp_path, dataset_name="docgenie_documents")
        
        try:
            # ==================== Update Status: Downloading ====================
            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "downloading"))
            print(f"[Job {request_id}] Status: downloading (fetching seed images)")
            
            # ==================== Step 1: Download Seed Images ====================
            log_verbose(f"[Job {request_id}] Downloading {len(seed_image_urls)} seed images...")
            seed_images_base64 = download_seed_images(seed_image_urls)
            log_verbose(f"[Job {request_id}] Downloaded {len(seed_images_base64)} images")
            
            # ==================== Step 2: Build Prompt ====================
            prompt_template_path = ENV.PROMPT_TEMPLATES_DIR / "ClaudeRefined12" / "seed-based-json.txt"
            
            if not prompt_template_path.exists():
                raise FileNotFoundError(f"Prompt template not found: {prompt_template_path}")
            
            prompt = build_prompt(
                language=prompt_params.get('language', 'English'),
                doc_type=prompt_params.get('doc_type', 'business and administrative'),
                gt_type=prompt_params.get('gt_type', 'Questions and answers'),
                gt_format=prompt_params.get('gt_format', '{"question": "answer"}'),
                num_solutions=prompt_params.get('num_solutions', 1),
                num_seed_images=len(seed_images_base64),
                prompt_template_path=prompt_template_path,
                enable_visual_elements=prompt_params.get('enable_visual_elements', False),
                visual_element_types=prompt_params.get('visual_element_types', [])
            )
            log_verbose(f"[Job {request_id}] Prompt built")
            
            # ==================== Step 3: Create Claude Batch ====================
            log_verbose(f"[Job {request_id}] Creating Claude batch (batched API)...")
            
            client = ClaudeBatchedClient(api_key=settings.ANTHROPIC_API_KEY)
            
            # Send batch with 1 message that generates num_solutions documents
            client.send_batch(
                model=settings.CLAUDE_MODEL,
                prompts=[prompt],  # Single prompt (list of 1)
                images_base64=[seed_images_base64],  # Single image set (list of 1)
                image_docids=[["seed"] * len(seed_images_base64)],  # Dummy doc IDs
                batch_data_directory=batch_dir,
                max_tokens=16384
            )
            
            print(f"[Job {request_id}] ⏳ Batch created, processing for Claude to process...")
            
            # ==================== Step 4: Poll Batch Until Complete ====================
            client.await_batches(
                batch_data_directory=batch_dir,
                message_data_directory=message_dir,
                sleep_seconds_between_batch=2,
                sleep_seconds_iteration=settings.BATCH_POLL_INTERVAL
            )
            
            print(f"[Job {request_id}] ✓ Batch complete")
            
            # ==================== Step 5: Read Batch Results ====================
            message_files = list(message_dir.glob("*.json"))
            
            if not message_files:
                raise RuntimeError("No message results found after batch completion")
            
            message_data = json.loads(message_files[0].read_text())
            
            if message_data.get('result_type') != 'succeeded':
                error_msg = message_data.get('error', 'Unknown error from Claude API')
                raise RuntimeError(f"Claude API error: {error_msg}")
            
            llm_response = message_data['response']
            log_verbose(f"[Job {request_id}] Received LLM response ({len(llm_response)} chars)")
            
            # ==================== Step 6: Extract HTML Documents ====================
            html_documents = extract_html_documents_from_response(llm_response)
            
            if not html_documents:
                raise RuntimeError("No valid HTML documents found in LLM response")
            
            print(f"[Job {request_id}] ✓ Extracted {len(html_documents)} documents")
            
            # ==================== Update Status: Generating ====================
            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "generating"))
            print(f"[Job {request_id}] Status: generating (processing documents)")
            
            # ==================== Step 7: Download Assets from Supabase ====================
            assets_temp_dir = None
            try:
                assets_path = f"{user_id}/{request_id}/assets"
                files = supabase_client.list_files("doc_storage", assets_path)
                
                # Filter out directories
                asset_files = [f for f in files if f.get('id') is not None]
                
                if asset_files:
                    assets_temp_dir = pathlib.Path(tempfile.mkdtemp())
                    print(f"[Job {request_id}] Found {len(asset_files)} assets in storage, downloading...")
                    
                    for file_info in asset_files:
                        file_name = file_info['name']
                        try:
                            file_content = supabase_client.download_file("doc_storage", f"{assets_path}/{file_name}")
                            with open(assets_temp_dir / file_name, 'wb') as f:
                                f.write(file_content)
                            log_verbose(f"  ✓ Downloaded {file_name}")
                        except Exception as download_err:
                            print(f"  ⚠ Failed to download {file_name}: {download_err}")
                else:
                    log_verbose(f"[Job {request_id}] No assets found in {assets_path}")
            except Exception as e:
                print(f"[Job {request_id}] ⚠ Asset check/download failed: {e}")
            
            # ==================== Step 8: Process Each Document ====================
            pdf_files = []
            metadata = []
            
            for idx, html in enumerate(html_documents):
                try:
                    doc_id = f"document_{idx + 1}"
                    log_verbose(f"[Job {request_id}] Processing document {idx + 1}/{len(html_documents)}")
                    
                    # Initialize original_pdf_path
                    original_pdf_path = None
                    
                    # Validate HTML
                    is_valid, error_msg = validate_html_structure(html)
                    if not is_valid:
                        print(f"[Job {request_id}] Document {idx + 1} HTML validation failed: {error_msg}")
                        continue
                    
                    # Extract ground truth and CSS
                    gt, html_clean = extract_ground_truth(html)
                    css, _ = extract_css_from_html(html_clean)
                    
                    # Render to PDF
                    pdf_path = tmp_path / f"{doc_id}.pdf"
                    pdf_path, width_mm, height_mm, geometries = await render_html_to_pdf(
                        html=html_clean,
                        output_pdf_path=pdf_path
                    )
                    
                    # Track original PDF
                    original_pdf_path = pdf_path
                    
                    # Validate PDF
                    is_valid, error_msg = validate_pdf(pdf_path)
                    if not is_valid:
                        print(f"[Job {request_id}] Document {idx + 1} PDF validation failed: {error_msg}")
                        continue
                    
                    # Extract bounding boxes
                    bboxes_raw = extract_bboxes_from_rendered_pdf(pdf_path)
                    
                    # Validate bboxes
                    is_valid, error_msg = validate_bboxes(bboxes_raw, min_bbox_count=1)
                    if not is_valid:
                        print(f"[Job {request_id}] Document {idx + 1} BBox validation warning: {error_msg}")
                    
                    log_verbose(f"[Job {request_id}] Document {idx + 1}: Extracted {len(bboxes_raw)} bboxes")
                    
                    # Process Stage 3 (Handwriting & Visual Elements) if enabled
                    final_image_b64 = None
                    handwriting_regions = []
                    visual_elements = []
                    handwriting_images = {}
                    visual_element_images = {}
                    ocr_results = None
                    pdf_with_handwriting_path = None
                    pdf_final_path = None
                    
                    if prompt_params.get('enable_handwriting') or prompt_params.get('enable_visual_elements'):
                        # Update status: Handwriting
                        if prompt_params.get('enable_handwriting'):
                            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "handwriting"))
                            log_verbose(f"[Job {request_id}] Status: handwriting (generating handwritten text)")
                        
                        log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing handwriting/visual elements...")
                        
                        try:
                            final_image_b64, handwriting_regions, visual_elements, handwriting_images, visual_element_images, pdf_with_handwriting_path, pdf_final_path = await process_stage3_complete(
                                pdf_path=pdf_path,
                                geometries=geometries,
                                ground_truth=gt,
                                bboxes_raw=bboxes_raw,
                                page_width_mm=width_mm,
                                page_height_mm=height_mm,
                                enable_handwriting=prompt_params.get('enable_handwriting', False),
                                handwriting_ratio=prompt_params.get('handwriting_ratio', 0.3),
                                enable_visual_elements=prompt_params.get('enable_visual_elements', False),
                                visual_element_types=prompt_params.get('visual_element_types', []),
                                seed=prompt_params.get('seed'),
                                assets_dir=assets_temp_dir,
                                barcode_number=prompt_params.get('barcode_number')
                            )
                            
                            # Use final PDF if both modifications applied, otherwise use handwriting PDF
                            if pdf_final_path and pdf_final_path.exists():
                                pdf_path = pdf_final_path
                            elif pdf_with_handwriting_path and pdf_with_handwriting_path.exists():
                                pdf_path = pdf_with_handwriting_path
                            
                            log_verbose(f"[Job {request_id}] Document {idx + 1}: {len(handwriting_regions)} handwriting, {len(visual_elements)} visual elements")
                        
                        except Exception as e:
                            print(f"[Job {request_id}] Document {idx + 1}: Stage 3 failed: {str(e)}")
                    
                    # Process Stage 4/5 (OCR) if needed
                    if prompt_params.get('enable_ocr'):
                        # Update status: OCR
                        retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "ocr"))
                        log_verbose(f"[Job {request_id}] Status: ocr (running OCR on documents)")
                        
                        log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing OCR...")
                        
                        try:
                            stage4_image, ocr_results = await process_stage4_ocr(
                                pdf_path=pdf_path,
                                enable_ocr=True,
                                dpi=settings.OCR_DPI
                            )
                            
                            if ocr_results:
                                log_verbose(f"[Job {request_id}] Document {idx + 1}: OCR complete - {len(ocr_results.get('words', []))} words")
                        
                        except Exception as e:
                            print(f"[Job {request_id}] Document {idx + 1}: OCR failed: {str(e)}")
                    
                    # Process Stage 5 (Dataset packaging) if needed
                    stage5_results = {}
                    if any([
                        prompt_params.get('enable_bbox_normalization'),
                        prompt_params.get('enable_gt_verification'),
                        prompt_params.get('enable_analysis'),
                        prompt_params.get('enable_debug_visualization')
                    ]):
                        # Update status: Validation (if GT verification enabled)
                        if prompt_params.get('enable_gt_verification'):
                            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "validation"))
                            log_verbose(f"[Job {request_id}] Status: validation (validating ground truth)")
                        
                        log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing dataset packaging...")
                        
                        try:
                            stage5_results = await process_stage5_complete(
                                document_id=doc_id,
                                pdf_path=pdf_path,
                                image_base64=final_image_b64,
                                ocr_results=ocr_results,
                                ground_truth=gt,
                                has_handwriting=prompt_params.get('enable_handwriting', False),
                                has_visual_elements=prompt_params.get('enable_visual_elements', False),
                                layout_elements=visual_elements,
                                enable_bbox_normalization=prompt_params.get('enable_bbox_normalization', False),
                                enable_gt_verification=prompt_params.get('enable_gt_verification', False),
                                enable_analysis=prompt_params.get('enable_analysis', False),
                                enable_debug_visualization=prompt_params.get('enable_debug_visualization', False)
                            )
                        
                        except Exception as e:
                            print(f"[Job {request_id}] Document {idx + 1}: Stage 5 failed: {str(e)}")
                    
                    # Track PDFs for metadata
                    if original_pdf_path and pdf_path != original_pdf_path:
                        pdf_files.append(original_pdf_path)
                        pdf_files.append(pdf_path)
                    else:
                        pdf_files.append(pdf_path)
                    
                    # Extract bbox_pdf (word + char) from original PDF (ground truth positions)
                    from .utils import extract_all_bboxes_from_pdf, extract_raw_annotations_from_geometries
                    log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Extracting bbox_pdf (word + char level) from original PDF...")
                    
                    try:
                        bboxes_pdf = extract_all_bboxes_from_pdf(original_pdf_path if original_pdf_path else pdf_path)
                        bbox_pdf_word = bboxes_pdf.get('word', [])
                        bbox_pdf_char = bboxes_pdf.get('char', [])
                        log_verbose(f"[Job {request_id}] Document {idx + 1}:   ✓ Extracted {len(bbox_pdf_word)} word bboxes, {len(bbox_pdf_char)} char bboxes from PDF")
                    except Exception as e:
                        print(f"[Job {request_id}] Document {idx + 1}:   ⚠ bbox_pdf extraction failed: {e}")
                        bbox_pdf_word = bboxes_raw  # Fallback to raw bboxes
                        bbox_pdf_char = []
                    
                    # Extract raw_annotations (layout boxes before normalization)
                    raw_annotations = None
                    if geometries:
                        log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Extracting raw_annotations from geometries...")
                        try:
                            raw_annotations = extract_raw_annotations_from_geometries(geometries)
                            log_verbose(f"[Job {request_id}] Document {idx + 1}:   ✓ Extracted {len(raw_annotations)} layout annotations")
                        except Exception as e:
                            print(f"[Job {request_id}] Document {idx + 1}:   ⚠ raw_annotations extraction failed: {e}")
                    
                    # Decode final image to bytes
                    final_image_bytes = None
                    if final_image_b64:
                        import base64
                        final_image_bytes = base64.b64decode(final_image_b64)
                    
                    # Decode debug visualization
                    debug_viz_bytes = None
                    if stage5_results.get('debug_visualization'):
                        import base64
                        debug_viz_dict = stage5_results['debug_visualization']
                        if debug_viz_dict and 'bbox_overlay_base64' in debug_viz_dict:
                            debug_viz_b64 = debug_viz_dict['bbox_overlay_base64']
                            debug_viz_bytes = base64.b64decode(debug_viz_b64)
                    
                    # Prepare token mapping if tokens exist
                    output_detail = prompt_params.get('output_detail', 'minimal')
                    token_mapping_data = None
                    if output_detail in ["dataset", "complete"]:
                        if handwriting_images or visual_element_images:
                            from .utils import create_token_mapping_json
                        token_mapping_data = create_token_mapping_json(
                            handwriting_regions,
                            handwriting_images,
                            visual_elements,
                            visual_element_images
                        )
                        log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Output detail '{output_detail}': Prepared {len(handwriting_images)} handwriting tokens, {len(visual_element_images)} visual elements")
                    
                    # Extract bbox_final_word and bbox_final_segment (from OCR or PDF)
                    bbox_final_word = None
                    bbox_final_segment = None
                    if ocr_results and ocr_results.get('words'):
                        # Use OCR results as final bboxes
                        bbox_final_word = ocr_results.get('words', [])
                        bbox_final_segment = ocr_results.get('lines', [])
                    else:
                        # Fallback to PDF bboxes if no OCR
                        bbox_final_word = bbox_pdf_word
                        bbox_final_segment = []  # No line-level data without OCR
                    
                    # Read PDF bytes for exporter
                    pdf_initial_bytes = original_pdf_path.read_bytes()
                    
                    # Read modified PDFs if they exist
                    pdf_with_handwriting_bytes = None
                    pdf_final_bytes = None
                    pdf_with_visual_elements_bytes = None
                    
                    if pdf_with_handwriting_path and pdf_with_handwriting_path.exists():
                        pdf_with_handwriting_bytes = pdf_with_handwriting_path.read_bytes()
                    
                    if pdf_final_path and pdf_final_path.exists():
                        pdf_final_bytes = pdf_final_path.read_bytes()
                    
                    # Special case: if only visual elements (no handwriting), pdf_final is actually pdf_with_visual_elements
                    if pdf_final_bytes and not pdf_with_handwriting_bytes:
                        pdf_with_visual_elements_bytes = pdf_final_bytes
                        pdf_final_bytes = None
                    
                    # Add document to exporter
                    log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Adding document to dataset exporter...")
                    exporter.add_document(
                        document_id=doc_id,
                        html=html_clean,
                        css=css,
                        pdf_initial=pdf_initial_bytes,
                        pdf_with_handwriting=pdf_with_handwriting_bytes,
                        pdf_with_visual_elements=pdf_with_visual_elements_bytes,
                        pdf_final=pdf_final_bytes,
                        final_image=final_image_bytes,
                        ground_truth=gt,
                        raw_annotations=raw_annotations,
                        bboxes_pdf_word=bbox_pdf_word,
                        bboxes_pdf_char=bbox_pdf_char,
                        bboxes_final_word=bbox_final_word,
                        bboxes_final_segment=bbox_final_segment,
                        bboxes_normalized_word=stage5_results.get('normalized_bboxes_word'),
                        bboxes_normalized_segment=stage5_results.get('normalized_bboxes_segment'),
                        gt_verification=stage5_results.get('gt_verification'),
                        token_mapping=token_mapping_data,
                        handwriting_regions=handwriting_regions,
                        handwriting_images=handwriting_images,
                        visual_elements=visual_elements,
                        visual_element_images=visual_element_images,
                        layout_elements=visual_elements,
                        geometries=geometries,
                        ocr_results=ocr_results,
                        analysis_stats=stage5_results.get('analysis_stats'),
                        debug_visualization=debug_viz_bytes
                    )
                    log_verbose(f"[Job {request_id}] Document {idx + 1}:   ✓ Document {doc_id} added to dataset")
                    
                    # Store comprehensive metadata (matching /generate/pdf format)
                    metadata.append({
                        "document_id": doc_id,
                        "filename": f"{doc_id}.pdf",
                        "bboxes": bboxes_raw,
                        "ground_truth": gt,
                        "geometries": geometries,
                        "page_width_mm": width_mm,
                        "page_height_mm": height_mm,
                        "handwriting_regions": handwriting_regions,
                        "visual_elements": visual_elements,
                        "has_stage3_image": final_image_b64 is not None,
                        "ocr_results": ocr_results,
                        # Stage 5 results
                        "normalized_bboxes_word": stage5_results.get('normalized_bboxes_word'),
                        "normalized_bboxes_segment": stage5_results.get('normalized_bboxes_segment'),
                        "gt_verification": stage5_results.get('gt_verification'),
                        "analysis_stats": stage5_results.get('analysis_stats'),
                        "debug_visualization_available": stage5_results.get('debug_visualization') is not None
                    })
                
                except Exception as e:
                    print(f"[Job {request_id}] Error processing document {idx + 1}: {str(e)}")
                    traceback.print_exc()
                    continue
            
            if not pdf_files:
                raise RuntimeError("Failed to process any documents")
            
            log_verbose(f"[Job {request_id}] Processed {len(pdf_files)} PDF files")
            
            # ==================== Step 8: Finalize Dataset & Create ZIP ====================
            log_verbose(f"[Job {request_id}] 📦 Finalizing dataset export...")
            exporter.finalize(
                request_id=request_id,
                user_id=user_id,
                prompt_params=prompt_params,
                api_mode="async"
            )
            log_verbose(f"[Job {request_id}]   ✓ Dataset structure finalized at {exporter.base_path}")
            
            # ==================== Update Status: Zipping ====================
            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "zipping"))
            print(f"[Job {request_id}] Status: zipping (creating ZIP archive)")
            
            # Create ZIP from organized dataset
            log_verbose(f"[Job {request_id}] 📦 Creating ZIP archive from dataset...")
            zip_path = tmp_path / f"docgenie_{request_id}.zip"
            
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_file:
                # Add all files from exporter.base_path
                for file_path in exporter.base_path.rglob('*'):
                    if file_path.is_file():
                        arcname = file_path.relative_to(exporter.base_path.parent)
                        zip_file.write(file_path, arcname)
            
            zip_size_mb = zip_path.stat().st_size / (1024 * 1024)
            log_verbose(f"[Job {request_id}]   ✓ ZIP created: {zip_size_mb:.2f} MB")
            
            # ==================== Update Status: Uploading ====================
            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "uploading"))
            print(f"[Job {request_id}] Status: uploading (uploading to Google Drive)")
            
            # ==================== Step 9: Upload to Google Drive ====================
            print(f"[Job {request_id}] ⬆️  Uploading to Google Drive...")
            
            google_drive_url = None
            gdrive_failed = False
            # Check if Google Drive token provided
            if not google_drive_token:
                print(f"[Job {request_id}] No Google Drive token provided. Skipping Google Drive upload.")
            else:
                try:
                    drive_client = GoogleDriveClient(
                        access_token=google_drive_token,
                        refresh_token=google_drive_refresh_token
                    )
                    google_drive_url = drive_client.upload_file(
                        file_path=zip_path,
                        filename=f"docgenie_{request_id}.zip",
                        folder_name=settings.GOOGLE_DRIVE_FOLDER_NAME
                    )
                    
                    print(f"[Job {request_id}] ✓ Uploaded to Google Drive: {google_drive_url}")
                
                except Exception as e:
                    print(f"[Job {request_id}] Google Drive upload failed: {str(e)}")
                    gdrive_failed = True
                    # Do not raise an error, just continue so we can still save to Supabase
            
            # ==================== Step 10: Store Results in Supabase ====================
            log_verbose(f"[Job {request_id}] Saving results to Supabase...")
            log_verbose(f"[Job {request_id}]   URL: {google_drive_url}")
            
            # Upload ZIP to Supabase
            zip_url = None
            try:
                zip_storage_path = f"{user_id}/{request_id}/generated/docgenie_{request_id}.zip"
                supabase_client.upload_to_storage("doc_storage", zip_storage_path, zip_path.read_bytes(), "application/zip")
                zip_url = supabase_client.get_public_url("doc_storage", zip_storage_path)
                print(f"[Job {request_id}] ✓ Uploaded ZIP to Supabase: {zip_url}")
            except Exception as e:
                print(f"[Job {request_id}] ⚠ Supabase ZIP upload failed: {e}")

            # Create generated document record
            retry_on_network_error(lambda: supabase_client.create_generated_document(
                request_id=request_id,
                file_url=google_drive_url,
                file_type="application/zip",
                page_count=len(metadata),  # Using document count as page_count
                model_version=settings.LLM_MODEL,
                zip_url=zip_url
            ))
            
            # Update request status
            status = "completed_gdrive_failed" if gdrive_failed else "completed"
            retry_on_network_error(lambda: supabase_client.update_request_status(request_id, status))
            
            # Log analytics
            retry_on_network_error(lambda: supabase_client.log_analytics_event(
                user_id=user_id,
                event_type="document_generation_completed",
                entity_id=request_id
            ))
            
            print(f"[Job {request_id}] ✅ Job completed successfully!")
        
        except Exception as e:
            # Update status to failed with error message
            error_message = f"{type(e).__name__}: {str(e)}"
            print(f"[Job {request_id}] ❌ Job failed: {error_message}")
            traceback.print_exc()
            
            retry_on_network_error(lambda: supabase_client.update_request_status(
                request_id=request_id,
                status="failed",
                error_message=error_message
            ))
            
            # Log analytics
            retry_on_network_error(lambda: supabase_client.log_analytics_event(
                user_id=user_id,
                event_type="document_generation_failed",
                entity_id=request_id
            ))
            
            raise  # Re-raise so RQ marks job as failed
        finally:
            # Clean up assets directory if it exists
            if 'assets_temp_dir' in locals() and assets_temp_dir and assets_temp_dir.exists():
                try:
                    shutil.rmtree(assets_temp_dir, ignore_errors=True)
                    print(f"[Job {request_id}] ✓ Cleaned up assets directory {assets_temp_dir}")
                except:
                    pass


def process_document_generation_job(request_id: str, request_data: Dict[str, Any]):
    """
    Synchronous wrapper for RQ - calls the async function with asyncio.run().
    
    This is the function that RQ worker calls. It runs the async version using asyncio.
    """
    print(f"{'='*60}")
    print(f"🎯 Worker picked up job: {request_id}")
    print(f"   User ID: {request_data.get('user_id', 'N/A')}")
    print(f"   Num documents: {request_data.get('prompt_params', {}).get('num_solutions', 'N/A')}")
    print(f"{'='*60}")
    
    return asyncio.run(process_document_generation_job_async(request_id, request_data))