import argparse
from sdsvkie.utils import read_json, yaml_load, write_json
from sdsvkie.utils.post_processing.invoice_post_processing import *
from sdsvkie.utils.post_processing.common_post_processing import normalize_number



def postprocess_invoice(invoice_data):
    if 'date' in invoice_data:
        invoice_data['date'] = post_processing_datetime(invoice_data['date'])

    #### normalize number
    number_fields = ['total_value', 'VAT_amount_value']
    for number_field in number_fields:
        if number_field not in invoice_data:
            continue
        invoice_data[number_field] = normalize_number(invoice_data[number_field])
    if 'buyer_tax_code_value' in invoice_data:
        invoice_data['buyer_tax_code_value'] = normalize_number(invoice_data['buyer_tax_code_value'], rerserve_minus=True)
    if 'seller_tax_code_value' in invoice_data:
        invoice_data['seller_tax_code_value'] = normalize_number(invoice_data['seller_tax_code_value'], rerserve_minus=True)
    if "seller_mobile_value" in invoice_data:
        invoice_data['seller_mobile_value'] = normalize_number(invoice_data['seller_mobile_value'], rerserve_minus=False, reserve_plus=True)


    for field_name in invoice_data.keys():
        field_value = invoice_data[field_name]
        field_value = field_value.replace("✪", " ")
        field_value = field_value.replace("\t", " ")
        field_value = re.sub(r"\s+", " ", field_value)
        invoice_data[field_name] = field_value

    return invoice_data




def format_e2e_data(input_json, output_json, cfg):
    cfg = yaml_load(cfg)
    classes = cfg['classes']
    value_classes = [cls_name for cls_name in classes if "_key" not in cls_name and "other" not in cls_name]
    print(value_classes)
    in_data = read_json(input_json)

    out_data = {}
    for img_id, img_data in in_data.items():

        new_img_data = postprocess_invoice(img_data)
        


        for cls_value in value_classes:
            if cls_value not in new_img_data:
                new_img_data[cls_value] = ""


        out_data[img_id] = new_img_data
    

    write_json(data=out_data, json_path=output_json)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=str)
    parser.add_argument("--input", type=str, help="e2e label file path")
    parser.add_argument("--out", type=str, help='postprocess e2e label')

    args = parser.parse_args()

    format_e2e_data(args.input, args.out, args.cfg)


