from celery import Celery
import base64
import environ
env = environ.Env(
    DEBUG=(bool, False)
)

class CeleryConnector:
    task_routes = {
        "process_id_result": {"queue": "id_card_rs"},
        "process_driver_license_result": {"queue": "driver_license_rs"},
        "process_invoice_result": {"queue": "invoice_rs"},
        "process_ocr_with_box_result": {"queue": "ocr_with_box_rs"},
        "process_template_matching_result": {"queue": "template_matching_rs"},
        # mock task
        "process_id": {"queue": "id_card"},
        "process_driver_license": {"queue": "driver_license"},
        "process_invoice": {"queue": "invoice"},
        "process_ocr_with_box": {"queue": "ocr_with_box"},
        "process_template_matching": {"queue": "template_matching"},
    }
    app = Celery(
        "postman",
        broker=env.str("CELERY_BROKER", "amqp://test:test@rabbitmq:5672"),
        broker_transport_options={'confirm_publish': False},
    )

    def process_id_result(self, args):
        return self.send_task("process_id_result", args)

    def process_driver_license_result(self, args):
        return self.send_task("process_driver_license_result", args)

    def process_invoice_result(self, args):
        return self.send_task("process_invoice_result", args)

    def process_ocr_with_box_result(self, args):
        return self.send_task("process_ocr_with_box_result", args)

    def process_template_matching_result(self, args):
        return self.send_task("process_template_matching_result", args)

    def process_id(self, args):
        return self.send_task("process_id", args)

    def process_driver_license(self, args):
        return self.send_task("process_driver_license", args)

    def process_invoice(self, args):
        return self.send_task("process_invoice", args)

    def process_ocr_with_box(self, args):
        return self.send_task("process_ocr_with_box", args)

    def process_template_matching(self, args):
        return self.send_task("process_template_matching", args)

    def send_task(self, name=None, args=None):
        if name not in self.task_routes or "queue" not in self.task_routes[name]:
            return self.app.send_task(name, args)

        return self.app.send_task(name, args, queue=self.task_routes[name]["queue"])


def main():
    rq_id = 345
    file_names = "abc.jpg"
    list_data = []

    with open("/home/sds/thucpd/aicr-2022/abc.jpg", "rb") as fs:
        encoded_string = base64.b64encode(fs.read()).decode("utf-8")
        list_data.append(encoded_string)

    c_connector = CeleryConnector()
    a = c_connector.process_id(args=(rq_id, list_data, file_names))

    print(a)


if __name__ == "__main__":
    main()