from celery import Celery
import os
import gc  # Garbage collector for explicit memory cleanup
from src.core.config import settings
from src.core.logging import logger
from src.database.session import SessionLocal
from src.modules.models import BnkAIReqRes
import src.modules.prompt.controller as prompt_controller
from src.modules.prompt.schemas.requests import MultiDocumentVerificationRequest
from datetime import datetime
from src.utils.pusher_config import pusher_client
from src.utils.file_chunk import (
    task_extract_and_chunk,
    task_process_chunks_llm,
    task_consolidate_outputs
)

# Initialize Celery with settings from config
celery_app = Celery(
    'bankruptcy_tasks',
    broker=settings.CELERY_BROKER_URL,
    backend=settings.CELERY_RESULT_BACKEND
)

# Complete Celery configuration
celery_app.conf.update(
    task_serializer='json',
    accept_content=['json'],
    result_serializer='json',
    timezone='UTC',
    enable_utc=True,
    task_track_started=True,
    task_time_limit=30 * 60,  # 30 minutes hard limit
    task_soft_time_limit=25 * 60,  # 25 minutes soft limit
    worker_prefetch_multiplier=1,  # Prefetch only 1 task (fair distribution)
    worker_max_tasks_per_child=40,  # Restart worker after 40 tasks (balance between stability and efficiency)
    worker_max_memory_per_child=400000,  # Restart worker at 400MB (400000 KB) - safe threshold
    task_acks_late=True,  # Acknowledge task only after completion
    task_reject_on_worker_lost=True,
    result_expires=3600,  # Results expire after 1 hour
    broker_pool_limit=50,  # Increase Redis connection pool
    broker_connection_retry_on_startup=True,
    worker_send_task_events=True,  # Enable task events for monitoring
    task_send_sent_event=True,
    # Rate limiting for external API protection
    task_default_rate_limit='100/m',  # Max 100 tasks per minute globally
    # Worker concurrency recommendations:
    # For memory-constrained environments: --concurrency=4
    # For balanced performance: --concurrency=8
    # For high-performance with 16GB+ RAM: --concurrency=12
)


@celery_app.task(bind=True, name='process_verification_task')
def process_verification_task(self, job_id: int):
    """
    Orchestrator task that dispatches the chunking pipeline:
    - Dispatches Task 1 (extract_and_chunk_task)
    - Task 1 auto-triggers Task 2 (process_chunks_llm_task)
    - Task 2 auto-triggers Task 3 (consolidate_outputs_task)
    
    Args:
        self: Task instance (bound)
        job_id: Primary key ID from bnk_ai_req_res table
    """
    db = SessionLocal()
    job = None
    
    try:
        # Fetch job from database
        job = db.query(BnkAIReqRes).filter(BnkAIReqRes.id == job_id).first()
        
        if not job:
            return {"error": f"Job not found: {job_id}"}
        
        # Parse request from stored JSON
        request_data = job.request
        parsed_request = MultiDocumentVerificationRequest.from_raw(request_data)
        
        # Extract file URLs from request
        file_urls = [str(url) for url in parsed_request.file_urls] if parsed_request.file_urls else []
        
        if not file_urls:
            raise Exception("No file URLs provided in request")
        
        # Dispatch Task 1 - which will auto-trigger Task 2 → Task 3
        logger.info(f"[Orchestrator] Dispatching extract_and_chunk_task for job_id={job_id}")
        extract_and_chunk_task.apply_async(args=[job_id, file_urls])
        
        return {"success": True, "job_id": job_id, "status": "pipeline_dispatched"}
        
    except Exception as e:
        # Store error in response
        if job:
            job.response = {"error": str(e), "success": False}
            job.updated_at = datetime.utcnow()
            db.commit()
        
        return {"error": str(e), "job_id": job_id}
        
    finally:
        db.close()
        # Explicit memory cleanup to prevent leaks
        gc.collect()


@celery_app.task(name='health_check')
def health_check():
    """Health check task for Celery worker"""
    return {"status": "healthy", "timestamp": datetime.utcnow().isoformat()}


# ============================================================================
# TASK A: Extract and Chunk Documents (handles 1 or many files)
# ============================================================================
@celery_app.task(
    name='extract_and_chunk_task',
    bind=True,
    autoretry_for=(Exception,),
    retry_backoff=True,
    retry_backoff_max=600,
    retry_jitter=True,
    max_retries=3
)
def extract_and_chunk_task(self, job_id: int, file_urls: list):
    """
    Task 1: Extract text from documents using LlamaParse and store chunks with overlap
    Auto-triggers Task 2 (process_chunks_llm_task) upon completion
    
    Args:
        job_id: Primary key ID from bnk_ai_req_res table
        file_urls: List of document URLs (can be 1 or many)
    """
    db = SessionLocal()
    job = None
    
    try:
        job = db.query(BnkAIReqRes).filter(BnkAIReqRes.id == job_id).first()
        if not job:
            return {"error": f"Job not found: {job_id}"}
        
        # Send start notification
        if pusher_client:
            pusher_client.trigger('notification-channel', 'notification-event', {
                'status': False,
                'message': 'Extracting text from documents',
                'request_id': job.request_id
            })
        
        import asyncio
        result = asyncio.run(task_extract_and_chunk(
            db=db,
            ai_req_res_id=job_id,
            file_urls=file_urls
        ))
        
        if not result.get("success"):
            raise Exception(result.get("error", "Unknown error"))
        
        # Send completion notification
        if pusher_client:
            pusher_client.trigger('notification-channel', 'notification-event', {
                'status': False,
                'message': f'Extraction complete: {result.get("total_chunks", 0)} chunks created',
                'request_id': job.request_id
            })
        
        # Get file_type and data from request for Task 2
        request_data = job.request
        file_type = request_data.get('file_type')
        data = request_data.get('data', {})
        
        if not file_type:
            raise Exception("file_type is required in request")
        
        # Auto-trigger Task 2 (Process Chunks with LLM)
        logger.info(f"[Task 1] Triggering Task 2 (process_chunks_llm_task) for job_id={job_id}")
        process_chunks_llm_task.apply_async(args=[job_id, file_type, data])
        
        return {
            "success": True,
            "job_id": job_id,
            "documents_processed": result.get("documents_processed"),
            "total_chunks": result.get("total_chunks"),
            "next_task": "process_chunks_llm_task"
        }
        
    except Exception as e:
        db.rollback()
        
        # Store error in job
        if job:
            job.response = {"error": str(e), "success": False, "stage": "extract_and_chunk"}
            job.updated_at = datetime.utcnow()
            db.commit()
        
        return {"error": str(e), "job_id": job_id}
        
    finally:
        db.close()
        # Explicit memory cleanup to prevent leaks
        gc.collect()


# ============================================================================
# TASK B: Process Chunks with LLM
# ============================================================================
@celery_app.task(
    name='process_chunks_llm_task',
    bind=True,
    autoretry_for=(Exception,),
    retry_backoff=True,
    retry_backoff_max=600,
    retry_jitter=True,
    max_retries=3
)
def process_chunks_llm_task(self, job_id: int, file_type: str, data: dict):
    """
    Task 2: Loop over chunks, send each to OpenAI, store responses
    Auto-triggers Task 3 (consolidate_outputs_task) upon completion
    
    Args:
        job_id: Primary key ID from bnk_ai_req_res table
        file_type: Document type for prompt selection
        data: Data payload for verification
    """
    db = SessionLocal()
    job = None
    
    try:
        job = db.query(BnkAIReqRes).filter(BnkAIReqRes.id == job_id).first()
        if not job:
            return {"error": f"Job not found: {job_id}"}
        
        # Send processing notification
        if pusher_client:
            pusher_client.trigger('notification-channel', 'notification-event', {
                'status': False,
                'message': 'Processing documents with AI',
                'request_id': job.request_id
            })
        
        logger.info(f"[Task 2] Using file_type={file_type}")
        
        import asyncio
        result = asyncio.run(task_process_chunks_llm(
            db=db,
            ai_req_res_id=job_id,
            file_type=file_type,
            data=data,
            request_id=job.request_id
        ))
        
        if not result.get("success"):
            raise Exception(result.get("error", "Unknown error"))
        
        # Skip completion notification if jailbreak detected
        if pusher_client and not result.get("jailbreak_detected"):
            pusher_client.trigger('notification-channel', 'notification-event', {
                'status': False,
                'message': f'AI analysis complete: {result.get("chunks_processed", 0)} chunks processed',
                'request_id': job.request_id
            })
        
        # Auto-trigger Task 3 (Consolidate Outputs) - pass jailbreak flag
        logger.info(f"[Task 2] Triggering Task 3 (consolidate_outputs_task) for job_id={job_id}")
        consolidate_outputs_task.apply_async(args=[job_id, file_type, result.get("jailbreak_detected", False)])
        
        return {
            "success": True,
            "job_id": job_id,
            "chunks_processed": result.get("chunks_processed"),
            "next_task": "consolidate_outputs_task"
        }
        
    except Exception as e:
        db.rollback()
        
        # Store error in job
        if job:
            job.response = {"error": str(e), "success": False, "stage": "process_chunks_llm"}
            job.updated_at = datetime.utcnow()
            db.commit()
        
        return {"error": str(e), "job_id": job_id}
        
    finally:
        db.close()
        # Explicit memory cleanup to prevent leaks
        gc.collect()


# ============================================================================
# TASK C: Consolidate LLM Outputs
# ============================================================================
@celery_app.task(
    name='consolidate_outputs_task',
    bind=True,
    autoretry_for=(Exception,),
    retry_backoff=True,
    retry_backoff_max=600,
    retry_jitter=True,
    max_retries=3
)
def consolidate_outputs_task(self, job_id: int, file_type: str, jailbreak_detected: bool = False):
    """
    Task 3: Consolidate all chunk LLM outputs into final response
    Sends callback to frontend when complete
    
    Args:
        job_id: Primary key ID from bnk_ai_req_res table
        file_type: Document type for consolidation logic
        jailbreak_detected: If True, skip notifications (only do consolidation)
    """
    db = SessionLocal()
    job = None
    
    try:
        job = db.query(BnkAIReqRes).filter(BnkAIReqRes.id == job_id).first()
        if not job:
            return {"error": f"Job not found: {job_id}"}
        
        # Send consolidation started notification only if no jailbreak
        if pusher_client and not jailbreak_detected:
            pusher_client.trigger('notification-channel', 'notification-event', {
                'status': False,
                'message': 'Consolidating results',
                'request_id': job.request_id
            })
        
        import asyncio
        result = asyncio.run(task_consolidate_outputs(
            db=db,
            ai_req_res_id=job_id,
            file_type=file_type
        ))
        
        if not result.get("success"):
            raise Exception(result.get("error", "Unknown error"))
        
        # Update job response
        consolidated_response = result.get("consolidated_response")
        job.response = consolidated_response
        job.updated_at = datetime.utcnow()
        db.commit()
        
        # Send callback if configured
        request_data = job.request
        callback_url = request_data.get('callback_url') if request_data else None
        
        logger.info(f"[Task 3] Consolidation complete for job_id={job_id}")
        logger.info(f"[Task 3] callback_url: {callback_url}")
        
        if callback_url:
            job.callback_url = callback_url
            db.commit()
            
            try:
                import requests
                payload = consolidated_response
                
                print(f"\n{'='*60}")
                print(f"[Task 3 Callback] Sending to: {callback_url}")
                print(f"{'='*60}\n")
                
                response = requests.post(callback_url, json=payload, timeout=30)
                
                print(f"[Task 3 Callback] Response status: {response.status_code}")
                print(f"[Task 3 Callback] Response body length: {len(response.text)}")
                
                if response.status_code == 200:
                    try:
                        job.callback_res = response.json()
                    except:
                        job.callback_res = {"raw_text": response.text}
                    
                    job.callback_res_header = {"status_code": response.status_code, "headers": dict(response.headers)}
                    db.commit()
                    
                    # Send success notification via Pusher only if no jailbreak
                    if pusher_client and not jailbreak_detected:
                        pusher_client.trigger('notification-channel', 'notification-event', {
                            'status': True,
                            'message': 'Verification completed successfully',
                            'request_id': job.request_id,
                            'timestamp': datetime.utcnow().isoformat()
                        })
                    
                    logger.info(f"[Task 3] Callback successful for job_id={job_id}")
                else:
                    job.callback_res = {"error": f"Callback returned {response.status_code}", "response": response.text[:500]}
                    db.commit()
                    logger.error(f"[Task 3] Callback failed with status {response.status_code}")
                    
            except Exception as callback_error:
                logger.error(f"[Task 3] Callback exception: {callback_error}")
                job.callback_res = {"error": str(callback_error)}
                db.commit()
        else:
            # No callback, but still send success notification
            if pusher_client:
                pusher_client.trigger('notification-channel', 'notification-event', {
                    'status': True,
                    'message': 'Verification completed successfully',
                    'request_id': job.request_id,
                    'timestamp': datetime.utcnow().isoformat()
                })
        
        return {
            "success": True,
            "job_id": job_id,
            "consolidated": True,
            "callback_sent": callback_url is not None
        }
        
    except Exception as e:
        db.rollback()
        logger.error(f"[Task 3] Error: {e}")
        
        # Store error in job
        if job:
            job.response = {"error": str(e), "success": False, "stage": "consolidate_outputs"}
            job.updated_at = datetime.utcnow()
            db.commit()
        
        return {"error": str(e), "job_id": job_id}
        
    finally:
        db.close()
        # Explicit memory cleanup to prevent leaks
        gc.collect()

#celery without DB


@celery_app.task(name='process_verification_task_withoutdb')
def process_verification_task_withoutdb(request_id: str, request_data: dict):
    """
    Celery task to process document verification WITHOUT database
    
    Args:
        request_id: Unique request ID (correlation ID)
        request_data: Full request payload dict (includes callback_url)
    """
    try:
        # DEBUG: Log what we received
        print(f"\n{'='*60}")
        print(f"[NO-DB TASK] Received request_id: {request_id}")
        print(f"[NO-DB TASK] request_data type: {type(request_data)}")
        print(f"[NO-DB TASK] request_data keys: {request_data.keys() if isinstance(request_data, dict) else 'Not a dict'}")
        print(f"[NO-DB TASK] callback_url in request_data: {'callback_url' in request_data if isinstance(request_data, dict) else 'N/A'}")
        print(f"[NO-DB TASK] callback_url value: {request_data.get('callback_url') if isinstance(request_data, dict) else 'N/A'}")
        print(f"{'='*60}\n")
        
        # Parse request from passed data (NO DB)
        parsed_request = MultiDocumentVerificationRequest.from_raw(request_data)
        
        # Process verification (controller is async, pass db=None)
        import asyncio
        result = asyncio.run(prompt_controller.verify_multi_documents_celery_withoutdb(None, parsed_request, None))
        
        # Send to frontend callback (NO DB SAVE)
        callback_url = request_data.get('callback_url')
        print(f"[NO-DB TASK] After processing - callback_url: {callback_url}")
        if callback_url:
            try:
                import requests
                import json
                
                # Send original response payload
                payload = result.dict()
                
                # Log what we're sending
                print(f"\n{'='*60}")
                print(f"[Callback NO-DB] Sending to: {callback_url}")
                print(f"[Callback NO-DB] Payload:")
                print(json.dumps(payload, indent=2, default=str))
                print(f"{'='*60}\n")
                
                response = requests.post(callback_url, json=payload, timeout=30)
                print(f"[Callback NO-DB] Response status: {response.status_code}")
                print(f"[Callback NO-DB] Response body: {response.text[:500]}")
                
                # Log webhook server's response
                if response.status_code == 200:
                    try:
                        webhook_response = response.json()
                        print(f"[Callback NO-DB] Webhook responded with: {webhook_response}")
                    except:
                        pass
            except Exception as callback_error:
                print(f"[Callback NO-DB] Failed: {str(callback_error)}")
        
        return {"success": True, "request_id": request_id}
        
    except Exception as e:
        # Best-effort error callback
        callback_url = request_data.get('callback_url') if isinstance(request_data, dict) else None
        if callback_url:
            try:
                import requests
                error_payload = {
                    "success": False,
                    "message": str(e),
                    "error": str(e)
                }
                requests.post(callback_url, json=error_payload, timeout=30)
            except Exception:
                pass
        
        return {"error": str(e), "request_id": request_id}