app/scraper/workers/pdf_processor.py

"""
Celery Worker - PDF Processing Pipeline
"""
import os
import sys
import json
import logging
import hashlib
import time
from datetime import datetime
from pathlib import Path
from celery import Celery
import redis
from minio import Minio

# Add project root to path to allow importing app configs
root_path = str(Path(__file__).resolve().parent.parent.parent)
if root_path not in sys.path:
    sys.path.insert(0, root_path)

from app.config import get_settings
settings = get_settings()

celery_redis_url = settings.redis_url
if celery_redis_url.startswith("rediss://"):
    if "ssl_cert_reqs=none" in celery_redis_url:
        celery_redis_url = celery_redis_url.replace("ssl_cert_reqs=none", "ssl_cert_reqs=CERT_NONE")
    elif "ssl_cert_reqs" not in celery_redis_url:
        separator = "&" if "?" in celery_redis_url else "?"
        celery_redis_url = f"{celery_redis_url}{separator}ssl_cert_reqs=CERT_NONE"

app = Celery('oag_pdf_processor')
app.config_from_object({
    'broker_url': celery_redis_url,
    'result_backend': celery_redis_url,
    'task_serializer': 'json',
    'accept_content': ['json'],
    'result_serializer': 'json',
    'timezone': 'Africa/Nairobi',
    'enable_utc': True,
    'task_track_started': True,
    'task_time_limit': 600,
    'worker_prefetch_multiplier': 1,
    'beat_schedule': {
        'run-scraper-daily': {
            'task': 'app.scraper.workers.pdf_processor.run_scraper',
            'schedule': 86400.0,  # Run daily (every 24 hours)
            'args': (None,),
        },
    },
})

logger = logging.getLogger(__name__)
redis_client = redis.Redis.from_url(settings.redis_url, decode_responses=True)

STREAMS = {
    'pdf_parse': 'stream:pdf:parse',
    'embed': 'stream:chunk:embed',
    'graph': 'stream:graph:build',
}

@app.task(bind=True, max_retries=3)
def process_pdf(self, report_id, minio_path, s3_url, metadata):
    try:
        logger.info(f"Processing PDF: {report_id}")
        pdf_data = download_from_minio(minio_path)
        if not pdf_data:
            raise Exception("Failed to download PDF")
        
        text_content = extract_text_from_pdf(pdf_data, report_id)
        chunks = chunk_document(text_content, metadata)
        
        # Enrich chunks with entities
        entities = extract_entities(text_content)
        # Add basic entity mapping to chunks (can associate entities per chunk text)
        for chunk in chunks:
            chunk_entities = []
            for ent in entities:
                if ent['text'] in chunk['text']:
                    chunk_entities.append(ent['text'])
            chunk['entities'] = list(set(chunk_entities))
            
        embeddings = generate_embeddings(chunks)
        
        # Save chunks in Qdrant (vector index)
        store_in_qdrant(report_id, chunks, embeddings, metadata)
        
        # Save chunks in PostgreSQL (FTS index & metadata database)
        store_in_postgresql_chunks(report_id, chunks, metadata)
        
        # Save optional Graph relationships
        store_in_neo4j(report_id, entities, metadata, chunks)
        
        # Update PostgreSQL Document status to ready
        update_postgresql_status(report_id, 'embedded', len(chunks))
        return {'status': 'success', 'chunks': len(chunks)}
    except Exception as exc:
        logger.error(f"PDF processing failed: {exc}")
        update_postgresql_status(report_id, 'failed', error=str(exc))
        raise self.retry(exc=exc, countdown=60)

def download_from_minio(minio_path):
    client = Minio(
        os.getenv('MINIO_ENDPOINT', 'localhost:9000'),
        access_key=os.getenv('MINIO_ACCESS_KEY', 'minioadmin'),
        secret_key=os.getenv('MINIO_SECRET_KEY', 'minioadmin'),
        secure=os.getenv('MINIO_SECURE', 'false').lower() == 'true',
    )
    bucket = os.getenv('MINIO_BUCKET', 'oag-raw-pdfs')
    try:
        response = client.get_object(bucket, minio_path)
        return response.read()
    except Exception as e:
        logger.error(f"MinIO download failed: {e}")
        return None

def extract_text_from_pdf(pdf_data, report_id):
    import pdfplumber
    import io
    text_parts = []
    try:
        with pdfplumber.open(io.BytesIO(pdf_data)) as pdf:
            for i, page in enumerate(pdf.pages):
                text = page.extract_text()
                if text:
                    text_parts.append(f"--- Page {i+1} ---\n{text}")
    except Exception as e:
        logger.error(f"pdfplumber failed: {e}")
    full_text = '\n'.join(text_parts)
    if len(full_text) < 1000:
        logger.warning(f"Low extraction, trying OCR for {report_id}")
        full_text = ocr_pdf(pdf_data)
    return full_text

def ocr_pdf(pdf_data):
    try:
        import pytesseract
        from pdf2image import convert_from_bytes
        images = convert_from_bytes(pdf_data)
        text_parts = []
        for i, image in enumerate(images):
            text = pytesseract.image_to_string(image)
            text_parts.append(f"--- Page {i+1} ---\n{text}")
        return '\n'.join(text_parts)
    except Exception as e:
        logger.error(f"OCR failed: {e}")
        return ""

def chunk_document(text, metadata):
    import re
    chunks = []
    pages = re.split(r'--- Page (\d+) ---', text)
    current_page = 0
    for i, segment in enumerate(pages):
        if segment.isdigit():
            current_page = int(segment)
            continue
        paragraphs = [p.strip() for p in segment.split('\n\n') if p.strip()]
        for para in paragraphs:
            chunk_type = classify_chunk(para)
            chunk = {
                'chunk_id': hashlib.md5(f"{metadata['report_id']}:{current_page}:{para[:50]}".encode()).hexdigest(),
                'report_id': metadata['report_id'],
                'text': para,
                'chunk_type': chunk_type,
                'page': current_page,
                'fiscal_year': metadata.get('fiscal_year'),
                'auditee': metadata.get('auditee'),
                'report_type': metadata.get('report_type'),
            }
            chunks.append(chunk)
    return chunks

def classify_chunk(text):
    text_lower = text.lower()
    if any(w in text_lower for w in ['ksh', 'kes', 'million', 'billion', 'budget', 'expenditure']):
        if any(w in text_lower for w in ['table', 'vote', 'head', 'item']):
            return 'table'
    if any(w in text_lower for w in ['finding', 'observed', 'noted that', 'audit revealed']):
        return 'finding'
    if any(w in text_lower for w in ['recommend', 'recommended', 'should', 'ought to']):
        return 'recommendation'
    if any(w in text_lower for w in ['constitution', 'article', 'section', 'public audit act', 'pfma']):
        return 'legal_ref'
    return 'narrative'

def extract_entities(text):
    try:
        import spacy
        nlp = spacy.load('en_core_web_sm')
        doc = nlp(text[:100000])
        return [{'text': ent.text, 'label': ent.label_, 'start': ent.start_char, 'end': ent.end_char} for ent in doc.ents]
    except Exception as e:
        logger.error(f"NER failed: {e}")
        return []

def generate_embeddings(chunks):
    try:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer('BAAI/bge-m3')
        texts = [chunk['text'] for chunk in chunks]
        embeddings = model.encode(texts, batch_size=32, show_progress_bar=False)
        return embeddings.tolist()
    except Exception as e:
        logger.error(f"Embedding failed: {e}")
        return []

def store_in_qdrant(report_id, chunks, embeddings, metadata):
    from qdrant_client import QdrantClient
    from qdrant_client.models import PointStruct, VectorParams, Distance
    from app.config import get_settings
    settings = get_settings()
    
    if settings.qdrant_api_key:
        client = QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key, timeout=30)
    else:
        client = QdrantClient(url=settings.qdrant_url, timeout=30)
        
    collection_name = settings.qdrant_collection_name
    try:
        client.create_collection(collection_name=collection_name, vectors_config=VectorParams(size=1024, distance=Distance.COSINE))
    except Exception:
        pass
        
    points = []
    for chunk, embedding in zip(chunks, embeddings):
        # Classify chunk type to match the SQL enum value exactly
        chunk_type = chunk['chunk_type']
        if chunk_type == 'legal_ref':
            chunk_type = 'legal_reference'
            
        import uuid
        try:
            chunk_uuid = str(uuid.UUID(hex=chunk['chunk_id']))
        except ValueError:
            chunk_uuid = str(uuid.uuid4())

        points.append(PointStruct(
            id=chunk_uuid,
            vector=embedding,
            payload={
                'document_id': report_id,
                'content': chunk['text'],
                'chunk_type': chunk_type,
                'page_range_start': chunk['page'],
                'page_range_end': chunk['page'],
                'fiscal_year': metadata.get('fiscal_year'),
                'audit_period': metadata.get('audit_period'),
                'auditee': metadata.get('auditee'),
                'report_type': metadata.get('report_type'),
                'entities': chunk.get('entities', []),
            }
        ))
    client.upsert(collection_name=collection_name, points=points)
    logger.info(f"Stored {len(points)} vectors in Qdrant")

def store_in_postgresql_chunks(report_id, chunks, metadata):
    from app.config import get_settings
    db_url = get_settings().database_url
    if db_url.startswith("postgresql+asyncpg://"):
        db_url = db_url.replace("postgresql+asyncpg://", "postgresql://")
    elif db_url.startswith("postgresql+async://"):
        db_url = db_url.replace("postgresql+async://", "postgresql://")
        
    import psycopg2
    from psycopg2.extras import execute_values, Json
    import uuid
    
    conn = psycopg2.connect(db_url)
    try:
        with conn.cursor() as cursor:
            # First, clean existing chunks for this report to maintain idempotency
            cursor.execute("DELETE FROM chunks WHERE document_id = %s", (report_id,))
            
            rows = []
            for idx, chunk in enumerate(chunks):
                # Classify chunk type to match the SQL enum value exactly
                chunk_type = chunk['chunk_type']
                if chunk_type == 'legal_ref':
                    chunk_type = 'legal_reference'
                
                try:
                    chunk_uuid = str(uuid.UUID(hex=chunk['chunk_id']))
                except ValueError:
                    chunk_uuid = str(uuid.uuid4())
                
                token_count = len(chunk['text'].split())
                
                rows.append((
                    chunk_uuid,
                    report_id,
                    chunk_type,
                    chunk['text'],
                    chunk['page'],
                    chunk['page'],
                    chunk.get('section_heading', ''),
                    Json(chunk.get('entities', [])),
                    'en',
                    metadata.get('audit_period', ''),
                    metadata.get('auditee', ''),
                    token_count,
                    idx,
                    datetime.utcnow()
                ))
                
            insert_query = """
                INSERT INTO chunks (
                    id, document_id, chunk_type, content, page_range_start, page_range_end,
                    section_heading, entities, language, audit_period, auditee, token_count, chunk_index, created_at
                ) VALUES %s
            """
            
            execute_values(cursor, insert_query, rows)
            conn.commit()
            logger.info(f"Stored {len(chunks)} chunks in PostgreSQL chunks table.")
    except Exception as e:
        logger.error(f"Failed to store chunks in PostgreSQL: {e}")
        conn.rollback()
        raise e
    finally:
        conn.close()

def store_in_neo4j(report_id, entities, metadata, chunks):
    from neo4j import GraphDatabase
    uri = os.getenv('NEO4J_URI', 'bolt://localhost:7687')
    user = os.getenv('NEO4J_USER', 'neo4j')
    password = os.getenv('NEO4J_PASSWORD', 'password')
    try:
        driver = GraphDatabase.driver(uri, auth=(user, password))
        with driver.session() as session:
            session.run("MERGE (r:AuditReport {report_id: $rid}) SET r.title = $title, r.fiscal_year = $fy, r.report_type = $type, r.source_url = $url",
                         rid=report_id, title=metadata.get('title', ''), fy=metadata.get('fiscal_year', ''), type=metadata.get('report_type', ''), url=metadata.get('source_url', ''))
            if metadata.get('auditee'):
                session.run("MERGE (a:Auditee {name: $auditee}) MERGE (r:AuditReport {report_id: $rid}) MERGE (a)<-[:AUDITS]-(r)",
                           auditee=metadata['auditee'], rid=report_id)
            for chunk in chunks:
                if chunk['chunk_type'] == 'finding':
                    session.run("MERGE (f:Finding {finding_id: $cid}) SET f.description = $text, f.page = $page MERGE (r:AuditReport {report_id: $rid}) MERGE (r)-[:CONTAINS]->(f)",
                               cid=chunk['chunk_id'], text=chunk['text'][:500], page=chunk['page'], rid=report_id)
        driver.close()
        logger.info(f"Stored graph data for {report_id}")
    except Exception as e:
        logger.warning(f"Neo4j store failed (optional graph step): {e}")

def update_postgresql_status(report_id, status, chunks_count=0, error=None):
    from app.config import get_settings
    db_url = get_settings().database_url
    if db_url.startswith("postgresql+asyncpg://"):
        db_url = db_url.replace("postgresql+asyncpg://", "postgresql://")
    elif db_url.startswith("postgresql+async://"):
        db_url = db_url.replace("postgresql+async://", "postgresql://")
        
    import psycopg2
    conn = psycopg2.connect(db_url)
    
    # Map embedded status to ready matching database enums
    db_status = 'ready' if status == 'embedded' else status
    
    try:
        with conn.cursor() as cursor:
            cursor.execute(
                "UPDATE documents SET status = %s, page_count = %s, error_message = %s, updated_at = NOW() WHERE id = %s",
                (db_status, chunks_count, error, report_id)
            )
            conn.commit()
    except Exception as e:
        logger.error(f"Failed to update document status in Postgres: {e}")
        conn.rollback()
    finally:
        conn.close()

@app.task
def consume_pdf_stream():
    while True:
        try:
            messages = redis_client.xreadgroup(groupname='workers', consumername='pdf-worker-1', streams={STREAMS['pdf_parse']: '>'}, count=1, block=5000)
            if not messages:
                continue
            for stream_name, entries in messages:
                for entry_id, fields in entries:
                    payload = json.loads(fields.get('payload', '{}'))
                    process_pdf.delay(
                        report_id=payload.get('report_id'),
                        minio_path=payload.get('minio_path', ''),
                        s3_url=payload.get('s3_url', ''),
                        metadata=payload,
                    )
                    redis_client.xack(stream_name, 'workers', entry_id)
                    logger.info(f"Queued task: {payload.get('report_id')}")
        except Exception as e:
            logger.error(f"Stream consumer error: {e}")
            time.sleep(5)

@app.task
def run_scraper(max_pages=None):
    import subprocess
    from pathlib import Path
    
    # Path to app/scraper directory
    scraper_dir = Path(__file__).resolve().parent.parent / "scraper"
    if not scraper_dir.exists():
        # Fallback to local workspace layout
        scraper_dir = Path(__file__).resolve().parent.parent.parent / "app" / "scraper"
        
    cmd = ["scrapy", "crawl", "oag_kenya"]
    if max_pages is not None:
        cmd.extend(["-a", f"max_pages={max_pages}"])
        
    logger.info(f"Triggering Scrapy crawler: {' '.join(cmd)}")
    
    # Execute scraper as subprocess in its directory
    process = subprocess.Popen(
        cmd,
        cwd=str(scraper_dir),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )
    stdout, stderr = process.communicate()
    
    logger.info(f"Scraper finished with exit code {process.returncode}")
    if process.returncode != 0:
        logger.error(f"Scraper error: {stderr}")
    return {"status": "success", "returncode": process.returncode}