services/workers.py

from __future__ import annotations

import asyncio
import tempfile
import httpx
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
from dataclasses import dataclass

from services import job_store
from services.cv_chunker import chunk_cv
from services.cv_converter import CVConverter
from services.job_matcher import JobMatcher


# ---------------------------------------------------------------------------
# Dedicated single-thread executor for ML work.
#
# Why not asyncio's default thread pool?
#   asyncio.to_thread() uses the loop's default ThreadPoolExecutor which is
#   shared with ALL other coroutines in the process (file I/O, HTTP clients,
#   etc.).  When two heavy ML tasks run simultaneously they can saturate that
#   shared pool, starving incoming HTTP requests of threads and making the
#   server appear frozen even though the event loop is technically free.
#
#   A dedicated pool with max_workers=1 means:
#   • ML work is 100% isolated — never competes with HTTP request threads.
#   • Only one Marker/chunk_cv call runs at a time (model is not thread-safe).
#   • The asyncio default pool stays free for file reads, httpx, etc.
# ---------------------------------------------------------------------------

_ML_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ml-worker")


class CvWorker:
    """Handles background tasks for CV processing."""

    def __init__(self, converter: CVConverter, matcher: JobMatcher):
        self.converter = converter
        self.matcher   = matcher

    async def run_cv_processing(
        self,
        job_id: str,
        file_bytes: bytes,
        filename: str,
        callback_url: str,
        callback_secret: str = None,
    ) -> None:
        """Process a CV and POST the result back to the Java callback endpoint."""
        loop     = asyncio.get_running_loop()
        tmp_path: Optional[Path] = None
        try:
            t_start = datetime.now(timezone.utc)

            suffix = Path(filename).suffix.lower() if filename else ".pdf"
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                tmp.write(file_bytes)
                tmp_path = Path(tmp.name)

            # Conversion — CPU-bound, isolated executor
            t_conv_start = datetime.now(timezone.utc)
            conversion = await loop.run_in_executor(
                _ML_EXECUTOR, self.converter.convert, tmp_path
            )
            t_conv_end = datetime.now(timezone.utc)
            conversion_ms = int((t_conv_end - t_conv_start).total_seconds() * 1000)

            if not conversion.success:
                error_msg = conversion.error or "Conversion failed"
                job_store.set_failed(job_id, error_msg)
                await self.post_callback(callback_url, {
                    "message":    error_msg,
                    "statusCode": 422,
                    "payload":    None,
                }, callback_secret)
                return

            # chunk_cv with embedder — CPU-bound, isolated executor
            t_chunk_start = datetime.now(timezone.utc)
            chunks = await loop.run_in_executor(
                _ML_EXECUTOR,
                lambda: chunk_cv(conversion.markdown, embedder=self.matcher._embed),
            )
            t_chunk_end = datetime.now(timezone.utc)
            chunking_ms = int((t_chunk_end - t_chunk_start).total_seconds() * 1000)

            t_end    = datetime.now(timezone.utc)
            total_ms = int((t_end - t_start).total_seconds() * 1000)

            payload = {
                "markdown":           conversion.markdown,
                "cv_title":           chunks["cv_title"],
                "seniority":          chunks["seniority"],
                "years_experience":   chunks["years_experience"],
                "category":           chunks["category"],
                "completeness_score": chunks["completeness_score"],
                "chunks": {
                    "summary":    chunks["chunks"]["summary"],
                    "contact":    chunks["chunks"]["contact"],
                    "links":      chunks["chunks"]["links"],
                    "skills":     chunks["chunks"]["skills"],
                    "experience": chunks["chunks"]["experience"],
                    "education":  chunks["chunks"]["education"],
                    "projects":   chunks["chunks"]["projects"],
                    "awards":     chunks["chunks"]["awards"],
                },
                "section_embeddings": chunks.get("section_embeddings"),
                "file_type":     conversion.file_type,
                "method_used":   conversion.method_used,
                "is_scanned":    conversion.is_scanned,
                "page_count":    conversion.page_count,
                "warnings":      conversion.warnings,
                "timing": {
                    "conversion_ms": conversion_ms,
                    "chunking_ms":   chunking_ms,
                    "total_ms":      total_ms,
                },
                "processed_at": t_end.isoformat(),
            }

            job_store.set_completed(job_id, payload)

            await self.post_callback(callback_url, {
                "message":    "CV processed successfully",
                "statusCode": 200,
                "payload":    payload,
            }, callback_secret)

        except Exception as exc:
            error_msg = f"Processing error: {exc}"
            job_store.set_failed(job_id, error_msg)
            await self.post_callback(callback_url, {
                "message":    error_msg,
                "statusCode": 500,
                "payload":    None,
            }, callback_secret)

        finally:
            if tmp_path is not None:
                try:
                    tmp_path.unlink(missing_ok=True)
                except Exception:
                    pass

    async def post_callback(
        self,
        callback_url: str,
        body: dict,
        callback_secret: str = None,
    ) -> None:
        """POST the processing result back to Java."""
        try:
            headers = {}
            if callback_secret:
                headers["X-Callback-Secret"] = callback_secret
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(callback_url, json=body, headers=headers)
                response.raise_for_status()
        except Exception as exc:
            print(f"[CV-ASYNC] Failed to POST callback to {callback_url}: {exc}")


@dataclass
class CvTask:
    job_id:          str
    file_bytes:      bytes
    filename:        str
    callback_url:    str
    callback_secret: str = None


class QueueManager:
    """
    Manages an asyncio Queue with a single background worker.

    concurrency=1 is intentional — the Marker ML model and the
    sentence-transformer are NOT thread-safe and must not run in
    parallel on the same process.  Jobs queue up and are processed
    one at a time.  The dedicated _ML_EXECUTOR above ensures ML work
    never blocks the asyncio event loop — HTTP endpoints (job-status,
    new submissions) stay responsive even while a CV is being processed.
    """

    def __init__(self, worker: CvWorker, concurrency: int = 1):
        self.worker      = worker
        self.concurrency = concurrency
        self.queue:  asyncio.Queue = asyncio.Queue()
        self.tasks:  list          = []
        self._active_ids: set[str] = set()

    async def start(self) -> None:
        for _ in range(self.concurrency):
            task = asyncio.create_task(self._worker_loop())
            self.tasks.append(task)
        print(f"[QUEUE] Started {self.concurrency} worker(s) — one CV processed at a time.")

    async def stop(self) -> None:
        for task in self.tasks:
            task.cancel()
        await asyncio.gather(*self.tasks, return_exceptions=True)
        _ML_EXECUTOR.shutdown(wait=False)
        print("[QUEUE] Stopped all workers.")

    async def _worker_loop(self) -> None:
        while True:
            try:
                task: CvTask = await self.queue.get()
                job_store.set_processing(task.job_id)
                self._active_ids.discard(task.job_id)
                print(
                    f"[QUEUE] Worker picked up job {task.job_id}. "
                    f"Remaining in queue: {self.queue.qsize()}"
                )
                await self.worker.run_cv_processing(
                    task.job_id,
                    task.file_bytes,
                    task.filename,
                    task.callback_url,
                    task.callback_secret,
                )
                self.queue.task_done()
            except asyncio.CancelledError:
                break
            except Exception as exc:
                print(f"[QUEUE] Unhandled error in worker loop: {exc}")

    async def enqueue(self, task: CvTask) -> None:
        if task.job_id in self._active_ids:
            print(f"[QUEUE] Duplicate submission rejected: job {task.job_id} is already queued.")
            raise ValueError(f"Job {task.job_id} is already in the processing queue.")
        self._active_ids.add(task.job_id)
        job_store.set_queued(task.job_id)
        await self.queue.put(task)
        print(f"[QUEUE] Job {task.job_id} enqueued. Queue depth: {self.queue.qsize()}")

    def get_queue_position(self, job_id: str) -> Optional[int]:
        try:
            for idx, task in enumerate(list(self.queue._queue)):
                if task.job_id == job_id:
                    return idx + 1
        except Exception as exc:
            print(f"[QUEUE] Error checking queue position: {exc}")
        return None