import logging
import json
import time

from app.tasks import celery_app
from app.database import SessionLocal
from app.models.pdf_document import PDFDocument
from app.services import pdf_service, vector_service
from app.config import settings

logger = logging.getLogger(__name__)


def _redis():
    import redis
    return redis.from_url(settings.REDIS_URL, decode_responses=True)


def _push_step(r, doc_id: int, step: str, message: str):
    r.rpush(f"pdf:steps:{doc_id}", json.dumps({"step": step, "message": message, "ts": time.time()}))
    r.expire(f"pdf:steps:{doc_id}", 3600)


@celery_app.task(name="process_pdf")
def process_pdf(document_id: int, file_path: str):
    """Background task: extract PDF text and store in ChromaDB."""
    db = SessionLocal()
    r = _redis()
    try:
        doc = db.query(PDFDocument).filter(PDFDocument.id == document_id).first()
        if not doc:
            logger.error(f"Document {document_id} not found")
            return

        _push_step(r, document_id, "start", "Starting PDF processing…")

        # Get page count
        total_pages = pdf_service.get_page_count(file_path)
        doc.total_pages = total_pages
        db.commit()
        _push_step(r, document_id, "pages", f"Found {total_pages} pages")

        # Extract text from all pages
        _push_step(r, document_id, "text", "Extracting text from pages…")
        pages = pdf_service.extract_text_by_page(file_path)
        if not pages:
            doc.status = "error"
            doc.error_message = "No text could be extracted from the PDF"
            db.commit()
            _push_step(r, document_id, "error", "No text could be extracted")
            return

        _push_step(r, document_id, "text", f"Extracted text from {len(pages)} pages")

        # Store in ChromaDB
        total_chunks = 0
        for page_num, text in pages.items():
            chunks = vector_service.chunk_text(text)
            total_chunks += len(chunks)
        _push_step(r, document_id, "vector", f"Vectorizing {total_chunks} chunks across {len(pages)} pages…")
        vector_service.store_pages(document_id, pages)
        _push_step(r, document_id, "vector", "Vectorization complete")

        # Mark as ready
        doc.status = "ready"
        db.commit()
        _push_step(r, document_id, "done", f"Ready — {total_pages} pages, {len(pages)} with text")
        logger.info(f"Document {document_id} processed: {total_pages} pages, {len(pages)} with text")

    except Exception as e:
        logger.exception(f"Error processing document {document_id}")
        _push_step(r, document_id, "error", f"Failed: {str(e)[:200]}")
        try:
            doc = db.query(PDFDocument).filter(PDFDocument.id == document_id).first()
            if doc:
                doc.status = "error"
                doc.error_message = str(e)[:500]
                db.commit()
        except Exception:
            pass
    finally:
        db.close()