"""
Use for deploy
"""
import argparse
import os
from glob import glob
from pathlib import Path

import cv2
import tqdm

from sdsvkie.cfg import load_cfg
from sdsvkie.engine.predictor import Predictor
from sdsvkie.utils.io_file import  write_json
from sdsvkie.utils import visualize_kie, IMG_EXT, PDF_EXT
import random
"""
python sdsvkie/tools/infer_e2e.py \
    --cfg /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/invoice/06062023/config.yaml \
    --weights /mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/invoice/06062023/best  \
    --device "cuda:0" \
    --img /mnt/hdd2T/AICR/Projects/2023/Vietinbank_POC/Invoice_JPG/ \
    --e2e  /mnt/hdd2T/AICR/Projects/2023/Vietinbank_POC/Invoice_KIE_Results/result.json


"""

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=str, default="sdsvkie/cfg/default.yaml")
    parser.add_argument("--img", type=str, default="img.jpg", help="image path or directory")
    parser.add_argument("--weights", type=str, default=None)
    parser.add_argument("--text_det", type=str, default=None)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--e2e", type=str, default=None)
    parser.add_argument("--vis", type=str, default=None)


    args = parser.parse_args()

    # predictor = Predictor(cfg=cfg)
    predictor = Predictor(**vars(args))

    if args.e2e:
        outdir_e2e = os.path.dirname(args.e2e)
        if not os.path.exists(outdir_e2e):
            os.makedirs(outdir_e2e, exist_ok=True)

    if os.path.isdir(args.img):
        img_paths = sorted(glob(args.img + "/*"))
        print("Infence image dir, total imgs: {}".format(len(img_paths)))
    else:
        img_paths = [args.img]

    out_dict = {}

    if "Others" in predictor.classes:
        colors = {
            "Store_name_value": (30,97,235),
            "id": (28,175,6),
            "Date_value": (241,26,242),
            "Total_value": (255,0,0),
        }
    else:
        colors = [
            (
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
            )
                for _ in range(len(predictor.classes))
        ]


    import time 
    
    for img_path in tqdm.tqdm(img_paths):
        print(img_path)
        if Path(img_path).suffix.lower() in IMG_EXT:
            img = cv2.imread(img_path)
            if img is None:
                print("img is None: ", img_path)
                continue
        elif Path(img_path).suffix.lower() in PDF_EXT:
            img = img_path  #pdf
        else:
            continue
        # try:
        out = predictor(img)
        # except Exception as err:
        #     print(err, img_path)
        #     continue 
        
        out_api = out['end2end_results']

        if not args.e2e:
            print(out_api)
        else:

            img_id = os.path.splitext(os.path.basename(img_path))[0]
            out_dict[img_id] = {
                field_name: field_item['value'] for field_name, field_item in out_api.items()
            }

        if args.vis:

            visualize_kie(
                img,
                boxes=[field_item['box'] for field_name, field_item in out_api.items() if len(field_item['box']) > 0],
                pred_labels=[field_name for field_name, field_item in out_api.items() if len(field_item['box']) > 0],
                image_name=os.path.basename(img_path),
                outdir=args.vis,
                colors=colors,
                texts = [field_item['value'] for field_name, field_item in out_api.items() if len(field_item['box']) > 0]
            )
    
    if args.e2e:
        write_json(os.path.join(args.e2e), out_dict)

        

    