MapMil/centers/parsing/workers/worker.py

"""CP worker: Redis jobs + real-time Telethon listener."""

import asyncio
import json
import logging
import os

import httpx
import redis

from workers.converter import event_record_to_ingest
from workers.parsers.telegram_events import parse_event_posts
from workers.sources.telegram_client import (
    TelegramAuthError,
    TelegramConfigError,
    fetch_channel_posts,
    normalize_channel,
)
from workers.sources.telegram_listener import TelegramListener
from workers.sources.telegram_session import close_shared_client, get_shared_client

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")
JOB_QUEUE_KEY = "cp:jobs"
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",
)


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


async def process_telegram_job(
    job_id: int,
    source_config: dict,
    *,
    client=None,
) -> tuple[list[dict], str | None]:
    channel = source_config.get("channel", "creamy_caprice")
    limit = int(source_config.get("limit", 100))

    try:
        username = normalize_channel(channel)
        posts = await fetch_channel_posts(username, limit=limit, client=client)
    except (TelegramConfigError, TelegramAuthError, ValueError) as exc:
        return [], str(exc)
    except Exception as exc:
        return [], f"Telegram: {exc}"

    records = parse_event_posts(posts)
    events = [event_record_to_ingest(r) for r in records]
    return events, None


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, *, tg_client=None) -> 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)
    await patch_job_status(job_id, "running")

    if source_type == "telegram":
        events, error = await process_telegram_job(job_id, source_config, client=tg_client)
    else:
        events, error = [], f"Unsupported source_type: {source_type}"

    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(*, tg_client=None) -> None:
    r = get_redis()
    logger.info("CP worker started, polling %s", JOB_QUEUE_KEY)

    while True:
        try:
            item = await asyncio.to_thread(r.blpop, JOB_QUEUE_KEY, POLL_TIMEOUT)
            if not item:
                continue
            _, raw = item
            payload = json.loads(raw)
            await handle_job(payload, tg_client=tg_client)
        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:
    client = await get_shared_client()
    listener = TelegramListener(client)
    worker_task = asyncio.create_task(worker_loop(tg_client=client))
    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(tg_client=None)


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


if __name__ == "__main__":
    main()