import argparse
import glob
from pathlib import Path

import cv2
from tqdm import tqdm

from sdsvkie.models.ocr import OCREngine
from sdsvkie.utils.visualize import visualize_ocr
from sdsvkie.utils.io_file import write_txt
from sdsvkie.utils.word_formation import sort_words

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--img", type=str)
    parser.add_argument("--out_dir", type=str, default=None)
    parser.add_argument("--device", type=str, default='cpu')
    parser.add_argument("--reserve_parent_dir", action='store_true')
    parser.add_argument("--only_text", action='store_true')
    parser.add_argument("--out_txt", type=str, default=None)
    parser.add_argument("--text_det", default="yolox-s-general-text-pretrain-20221226")
    parser.add_argument("--text_recog", default="satrn-lite-general-pretrain-20230106")
    
    args = parser.parse_args()
    ocr_engine = OCREngine(text_det=args.text_det, text_recog=args.text_recog, device=args.device)


    if args.reserve_parent_dir:
        paths = glob.glob(args.img + "/*/*")
    else:
        paths = glob.glob(args.img + "/*")
    for path in tqdm(paths):
        img = cv2.imread(path)
        if img is None:
            print(path)
            continue
        ocr_output = ocr_engine(img, extend_ratio=[0.1, 0.3], ratio_thr=5)

        if args.out_dir:
            if args.reserve_parent_dir:
                out_dir_img = Path(args.out_dir) / Path(path).parent.name
            else:
                out_dir_img = Path(args.out_dir)
            if not out_dir_img.exists():
                out_dir_img.mkdir(parents=True)

        
            visualize_ocr(
                img = img,
                boxes=ocr_output['boxes'],
                texts=ocr_output['texts'],
                image_name=str(Path(path).name),
                outdir=str(out_dir_img)
            )
        
        if args.out_txt:
            

            if args.reserve_parent_dir:
                out_dir_txt = Path(args.out_txt) / Path(path).parent.name
            else:
                out_dir_txt = Path(args.out_txt)
            if not out_dir_txt.exists():
                out_dir_txt.mkdir(parents=True)
            
            out_txt_path = out_dir_txt / Path(path).with_suffix(".txt").name
            data = []
            if args.only_text:
                out = sort_words(ocr_output)
                text = " ".join(out['texts'])
                data.append(text)
            else:
                
                for box, text in zip(ocr_output['boxes'], ocr_output['texts']):
                    item = "\t".join([str(int(x)) for x in box])
                    item = "\t".join([item, text])
                    data.append(item)
            write_txt(str(out_txt_path), data)

        



