# Copyright (C) 2020  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

import datetime
import json
import logging
import sys

import pika

from swh.core.statsd import statsd
from swh.scheduler import get_scheduler

logger = logging.getLogger(__name__)


def utcnow():
    return datetime.datetime.now(tz=datetime.timezone.utc)


def get_listener(broker_url, queue_name, scheduler_backend):
    connection = pika.BlockingConnection(pika.URLParameters(broker_url))
    channel = connection.channel()

    channel.queue_declare(queue=queue_name, durable=True)

    exchange = "celeryev"
    routing_key = "#"
    channel.queue_bind(queue=queue_name, exchange=exchange, routing_key=routing_key)

    channel.basic_qos(prefetch_count=1000)

    channel.basic_consume(
        queue=queue_name, on_message_callback=get_on_message(scheduler_backend),
    )

    return channel


def get_on_message(scheduler_backend):
    def on_message(channel, method_frame, properties, body):
        try:
            events = json.loads(body)
        except Exception:
            logger.warning("Could not parse body %r", body)
            events = []

        if not isinstance(events, list):
            events = [events]

        for event in events:
            logger.debug("Received event %r", event)
            process_event(event, scheduler_backend)

        channel.basic_ack(delivery_tag=method_frame.delivery_tag)

    return on_message


def process_event(event, scheduler_backend):
    uuid = event.get("uuid")
    if not uuid:
        return

    event_type = event["type"]
    statsd.increment(
        "swh_scheduler_listener_handled_event_total", tags={"event_type": event_type}
    )

    if event_type == "task-started":
        scheduler_backend.start_task_run(
            uuid, timestamp=utcnow(), metadata={"worker": event.get("hostname")},
        )
    elif event_type == "task-result":
        result = event["result"]

        status = None

        if isinstance(result, dict) and "status" in result:
            status = result["status"]
            if status == "success":
                status = "eventful" if result.get("eventful") else "uneventful"

        if status is None:
            status = "eventful" if result else "uneventful"

        scheduler_backend.end_task_run(
            uuid, timestamp=utcnow(), status=status, result=result
        )
    elif event_type == "task-failed":
        scheduler_backend.end_task_run(uuid, timestamp=utcnow(), status="failed")


if __name__ == "__main__":
    url = sys.argv[1]
    logging.basicConfig(level=logging.DEBUG)
    scheduler_backend = get_scheduler("local", args={"db": "service=swh-scheduler"})
    channel = get_listener(url, "celeryev.test", scheduler_backend)
    logger.info("Start consuming")
    channel.start_consuming()
