MapMil/centers/parsing/workers/worker.py

"""CP worker: Redis jobs via adapter registry + optional Telethon listener."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
from pathlib import Path

# Monorepo local runs / Docker: ensure contracts/ is importable
_HERE = Path(__file__).resolve()
for _candidate in (_HERE.parent, *_HERE.parents):
    if (_candidate / "contracts").is_dir():
        if str(_candidate) not in sys.path:
            sys.path.insert(0, str(_candidate))
        break

import httpx
import redis

from contracts.queues import (
    LEGACY_JOB_QUEUE_KEY,
    family_for_source,
    queue_key_for_family,
    queue_key_for_source,
)
from workers.adapters import build_registry
from workers.adapters.base import WorkerContext

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("cp-worker")

REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
CA_API_URL = os.getenv("CA_API_URL", "http://ca-api:8000")
INTERNAL_TOKEN = os.getenv("INTERNAL_TOKEN", "dev-internal-token")
POLL_TIMEOUT = int(os.getenv("WORKER_POLL_TIMEOUT", "5"))
LISTENER_ENABLED = os.getenv("TELEGRAM_LISTENER_ENABLED", "true").lower() not in (
    "0",
    "false",
    "no",
    "off",
)
# Comma-separated families this process polls (default: derived from adapters)
WORKER_FAMILIES = os.getenv("WORKER_FAMILIES", "").strip()

REGISTRY = build_registry()


def get_redis() -> redis.Redis:
    return redis.from_url(REDIS_URL, decode_responses=True)


def _poll_keys() -> list[str]:
    if WORKER_FAMILIES:
        families = [f.strip() for f in WORKER_FAMILIES.split(",") if f.strip()]
    else:
        families = sorted(
            {
                family_for_source(source_type)
                for source_type in REGISTRY.enabled_types()
            }
        )
    keys = [queue_key_for_family(f) for f in families]
    # Backward compatible: telegram workers also drain legacy cp:jobs
    if "telegram" in families and LEGACY_JOB_QUEUE_KEY not in keys:
        keys.append(LEGACY_JOB_QUEUE_KEY)
    return keys


async def post_ingest(job_id: int, events: list[dict]) -> None:
    if not events:
        return

    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            f"{CA_API_URL}/internal/ingest",
            json={"job_id": job_id, "events": events},
            headers={"X-Internal-Token": INTERNAL_TOKEN},
        )
        response.raise_for_status()
        result = response.json()
        logger.info(
            "Ingested job %s: new=%s skipped=%s map=%s",
            job_id,
            result.get("ingested", 0),
            result.get("skipped", 0),
            result.get("map_objects_synced", 0),
        )


async def patch_job_status(job_id: int, status: str, error: str | None = None) -> None:
    async with httpx.AsyncClient(timeout=30.0) as client:
        params = {"status": status}
        if error:
            params["error"] = error
        response = await client.patch(
            f"{CA_API_URL}/internal/jobs/{job_id}",
            params=params,
            headers={"X-Internal-Token": INTERNAL_TOKEN},
        )
        response.raise_for_status()


async def handle_job(payload: dict, *, ctx: WorkerContext) -> None:
    job_id = payload["job_id"]
    source_type = payload["source_type"]
    source_config = payload.get("source_config", {})

    logger.info("Processing job %s (%s)", job_id, source_type)

    adapter = REGISTRY.get(source_type)
    if adapter is None:
        try:
            target = queue_key_for_source(source_type)
        except ValueError:
            await patch_job_status(
                job_id,
                "failed",
                error=f"Unsupported source_type: {source_type}",
            )
            return
        get_redis().rpush(target, json.dumps(payload))
        logger.warning(
            "Job %s (%s) not enabled here; requeued to %s",
            job_id,
            source_type,
            target,
        )
        return

    await patch_job_status(job_id, "running")
    events, error = await adapter.run(job_id, source_config, ctx=ctx)

    if error:
        logger.error("Job %s failed: %s", job_id, error)
        await patch_job_status(job_id, "failed", error=error)
        return

    try:
        await post_ingest(job_id, events)
        if not events:
            await patch_job_status(job_id, "completed", error="No events found")
        logger.info("Job %s completed with %s events", job_id, len(events))
    except Exception as exc:
        logger.exception("Ingest failed for job %s", job_id)
        await patch_job_status(job_id, "failed", error=str(exc))


async def worker_loop(*, ctx: WorkerContext) -> None:
    r = get_redis()
    keys = _poll_keys()
    logger.info(
        "CP worker started, polling %s (adapters: %s)",
        keys,
        ", ".join(REGISTRY.enabled_types()) or "(none)",
    )

    while True:
        try:
            item = await asyncio.to_thread(r.blpop, keys, POLL_TIMEOUT)
            if not item:
                continue
            _, raw = item
            payload = json.loads(raw)
            await handle_job(payload, ctx=ctx)
        except redis.RedisError as exc:
            logger.error("Redis error: %s", exc)
            await asyncio.sleep(3)
        except Exception:
            logger.exception("Unexpected worker error")
            await asyncio.sleep(1)


async def run_with_listener() -> None:
    if "telegram" not in REGISTRY:
        logger.warning("Listener requested but telegram adapter is not enabled")
        await worker_loop(ctx=WorkerContext())
        return

    from workers.sources.telegram_client import TelegramAuthError, TelegramConfigError
    from workers.sources.telegram_listener import TelegramListener
    from workers.sources.telegram_session import close_shared_client, get_shared_client

    while True:
        try:
            client = await get_shared_client()
            break
        except (TelegramAuthError, TelegramConfigError, ConnectionError, OSError) as exc:
            logger.error(
                "Telegram connect failed (%s); retry in 60s (batch without TG)",
                exc,
            )
            try:
                await asyncio.wait_for(
                    worker_loop(ctx=WorkerContext(tg_client=None)),
                    timeout=60,
                )
            except asyncio.TimeoutError:
                pass
            continue
        except Exception:
            logger.exception("Unexpected Telegram connect error; retry in 60s")
            await asyncio.sleep(60)

    listener = TelegramListener(client)
    ctx = WorkerContext(tg_client=client)
    worker_task = asyncio.create_task(worker_loop(ctx=ctx))
    listener_task = asyncio.create_task(listener.run())

    logger.info("Telegram listener enabled (shared session with batch worker)")
    try:
        await client.run_until_disconnected()
    finally:
        worker_task.cancel()
        listener_task.cancel()
        await close_shared_client()


async def run_batch_only() -> None:
    logger.info("Telegram listener disabled")
    await worker_loop(ctx=WorkerContext(tg_client=None))


def main() -> None:
    try:
        if LISTENER_ENABLED and "telegram" in REGISTRY:
            asyncio.run(run_with_listener())
        else:
            asyncio.run(run_batch_only())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()