from celery_worker.worker import app
import numpy as np
import cv2

@app.task(name="process_id")
def process_id(rq_id, sub_id, folder_name, list_url, user_id):
    from common.serve_model import predict
    from celery_worker.client_connector import CeleryConnector
    
    c_connector = CeleryConnector()
    try:
        result = predict(rq_id, sub_id, folder_name, list_url, user_id, infer_name="id_card")
        print(result)
        result = {
                    "status": 200,
                    "content": result,
                    "message": "Success",
                }
        c_connector.process_id_result((rq_id, result))
        return {"rq_id": rq_id}
        # if image_croped is not None:
        #     if result["data"] == []:
        #         result = {
        #         "status": 404,
        #         "content": {},
        #         }
        #         c_connector.process_id_result((rq_id, result, None))
        #         return {"rq_id": rq_id}
        #     else:
        #         result = {
        #             "status": 200,
        #             "content": result,
        #             "message": "Success",
        #         }
        #         c_connector.process_id_result((rq_id, result))
        #         return {"rq_id": rq_id}
        # elif image_croped is None:
        #     result = {
        #         "status": 404,
        #         "content": {},
        #     }
        #     c_connector.process_id_result((rq_id, result, None))
        #     return {"rq_id": rq_id}
            
    except Exception as e:
        print(e)
        result = {
                "status": 404,
                "content": {},
        }
        c_connector.process_id_result((rq_id, result, None))
        return {"rq_id": rq_id}


@app.task(name="process_driver_license")
def process_driver_license(rq_id, sub_id, folder_name, list_url, user_id):
    from common.serve_model import predict
    from celery_worker.client_connector import CeleryConnector
    
    c_connector = CeleryConnector()
    try:
        result = predict(rq_id, sub_id, folder_name, list_url, user_id, infer_name="driving_license")
        result = {
                    "status": 200,
                    "content": result,
                    "message": "Success",
                }
        c_connector.process_driver_license_result((rq_id, result))
        return {"rq_id": rq_id}
        # result, image_croped = predict(str(url), "driving_license")
        # if image_croped is not None:
        #     if result["data"] == []:
        #         result = {
        #         "status": 404,
        #         "content": {},
        #         }
        #         c_connector.process_driver_license_result((rq_id, result, None))
        #         return {"rq_id": rq_id}
        #     else:
        #         result = {
        #             "status": 200,
        #             "content": result,
        #             "message": "Success",
        #         }
        #         path_image_croped = "/app/media/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id)
        #         cv2.imwrite("/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id), image_croped)
        #         c_connector.process_driver_license_result((rq_id, result, path_image_croped))
        #         return {"rq_id": rq_id}
        # elif image_croped is None:
        #     result = {
        #         "status": 404,
        #         "content": {},
        #     }
        #     c_connector.process_driver_license_result((rq_id, result, None))
        #     return {"rq_id": rq_id}
    except Exception as e:
        print(e)
        result = {
                "status": 404,
                "content": {},
        }
        c_connector.process_driver_license_result((rq_id, result, None))
        return {"rq_id": rq_id}


@app.task(name="process_template_matching")
def process_template_matching(rq_id, sub_id, folder_name, url, tmp_json, user_id):  
    from TemplateMatching.src.ocr_master import Extractor
    from celery_worker.client_connector import CeleryConnector
    import urllib

    c_connector = CeleryConnector()
    extractor = Extractor()
    try:
        req = urllib.request.urlopen(url)
        arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
        img = cv2.imdecode(arr, -1)
        imgs = [img]
        image_aliged = extractor.image_alige(imgs, tmp_json)
        if image_aliged is None:
            result = {
                "status": 401,
                "content": "Image is not match with template",
            }
            c_connector.process_template_matching_result(
                (rq_id, result, None)
            )
            return {"rq_id": rq_id}
        else:
            output = extractor.extract_information(
                image_aliged, tmp_json
            )
            path_image_croped = "/app/media/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id)
            cv2.imwrite("/users/{}/subscriptions/{}/requests/{}/{}/image_croped.jpg".format(user_id,sub_id,folder_name,rq_id), image_aliged)
            if output == {}:
                result = {"status": 404, "content": {}}
                c_connector.process_template_matching_result((rq_id, result, None))
                return {"rq_id": rq_id}
            else:
                result = {
                    "document_type": "template_matching",
                    "fields": []
                }
                print(output)
                for field in tmp_json["fields"]:
                    print(field["label"])
                    field_value = {
                        "label": field["label"],
                        "value": output[field["label"]],
                        "box": [float(num) for num in field["box"]],
                        "confidence": 0.98 #TODO confidence
                    }
                    result["fields"].append(field_value)
                
                print(result)
                result = {"status": 200, "content": result}
                c_connector.process_template_matching_result(
                    (rq_id, result, path_image_croped)
                )
                return {"rq_id": rq_id}
    except Exception as e:
        print(e)
        result = {"status": 404, "content": {}}
        c_connector.process_template_matching_result((rq_id, result, None))
        return {"rq_id": rq_id}


# @app.task(name="process_invoice")
# def process_invoice(rq_id, url):
#     from celery_worker.client_connector import CeleryConnector
#     from Kie_Hoanglv.prediction import predict
    
#     c_connector = CeleryConnector()
#     try:
#         print(url)
#         result = predict(str(url))
#         hoadon = {"status": 200, "content": result, "message": "Success"}
#         c_connector.process_invoice_result((rq_id, hoadon))
#         return {"rq_id": rq_id}
    
#     except Exception as e:
#         print(e)
#         hoadon = {"status": 404, "content": {}}
#         c_connector.process_invoice_result((rq_id, hoadon))
#         return {"rq_id": rq_id}

@app.task(name="process_invoice")
def process_invoice(rq_id, list_url):
    from celery_worker.client_connector import CeleryConnector
    from common.process_pdf import compile_output
    
    c_connector = CeleryConnector()
    try:
        result = compile_output(list_url)
        hoadon = {"status": 200, "content": result, "message": "Success"}
        c_connector.process_invoice_result((rq_id, hoadon))
        return {"rq_id": rq_id}
    except Exception as e:
        print(e)
        hoadon = {"status": 404, "content": {}}
        c_connector.process_invoice_result((rq_id, hoadon))
        return {"rq_id": rq_id}
    

@app.task(name="process_ocr_with_box")
def process_ocr_with_box(rq_id, list_url):
    from celery_worker.client_connector import CeleryConnector
    from common.process_pdf import compile_output_ocr_base

    c_connector = CeleryConnector()
    try:
        result = compile_output_ocr_base(list_url)
        result = {"status": 200, "content": result, "message": "Success"}
        c_connector.process_ocr_with_box_result((rq_id, result))
        return {"rq_id": rq_id}
    except Exception as e:
        print(e)
        result = {"status": 404, "content": {}}
        c_connector.process_ocr_with_box_result((rq_id, result))
        return {"rq_id": rq_id}