from datetime import datetime
import glob
import os
import soundfile as sf
import re
import tqdm
import pysubs2

### 부동산 - 전사 srt -> split ogg + 텍스트
root_path = ""

dir_list = glob.glob(root_path + "*/") ## root path 밑에 있는 sub dir 모두 가져오기

pattern_1 = r'\(\(\)\)' ## (())
pattern_2 = r'\{(\w*)\}' ## {@@@} -> @@
pattern_2_1 = r'\[(\w*)\]' ## [@@@] -> @@
pattern_2_2 = r'\((\w*)\)' ## (@@@) 찾기 -> @@
pattern_2_1 = r'[^/]\(([\w|\s]*)\)[^/]' ## (@@)/(@@)이 아닌 (@@) 찾기 repl=' \1 '

pattern_3_1 = r'\(([\w|\s|.]*)\)/\(([\w|\s|.]*)\)' ## (@@)/(@@) -> {@@}/{@@} 임시 변환, repl=r'{\1}/{\2}'
pattern_3_2 = r'\(([\w|\s]*)\)' ## (@@) 찾기, repl=r'\1'
pattern_3_3 = r'\{([\w|\s|.]*)\}/\{([\w|\s|.]*)\}' ## {@@}/{@@} -> (@@)/(@@) 임시변환 해제, r'(\1)/(\2)'

pattern_4 = r'[\(\)\{\}]' ## (, ), {, } 찾기
pattern_4_1 = r'&([\w|-]*)&' ## &(id-id123)& 찾기

pattern_5 = r"[^\uAC00-\uD7A30-9a-zA-Z\s]" ## 특문
pattern_6 = r" {2,}" ## 빈 문자열 2개 이상

pattern_7 = r"\\N"

now = datetime.now()
time_str = now.astimezone().strftime('%Y_%m_%d_%H_%M_%S')

def postprocess(origin):
    _a = re.sub(pattern=pattern_7, repl=' ', string = origin)
    _a = re.sub(pattern=pattern_1, repl='', string = _a)
    _a = re.sub(pattern=pattern_2, repl=r'\1', string = _a)
    _a = re.sub(pattern=pattern_2_1, repl=r'\1', string = _a)
    # _a = re.sub(pattern=pattern_3, repl=r'\1', string = _a)

    _a = re.sub(pattern=pattern_5, repl = "", string = _a)
    _a = re.sub(pattern=pattern_6, repl = " ", string = _a)
    _a = _a.strip()

    if len(_a) == 0:
        return None
    else:
        return _a


def time_convert(time):
    split = re.split(r',|:', time)

    mi = int(split[0])
    mi = mi * 60 + int(split[1])
    mi = mi * 60 + int(split[2]) + float('0.' + split[3])
    
    return mi

 
for _dir in dir_list:
    post_none_list = []
    not_exist_list = []

    ## _dir이 directory일 때만 실행
    if not os.path.isdir(_dir):
        continue

    ## 저장용 split directory 생성
    output_dir = os.path.join(_dir, "split")
    os.makedirs(output_dir, exist_ok=True)

    ## _dir 폴더 하위의 srt 파일 모두 가져오기
    srt_list = glob.glob(os.path.join(_dir,'*.srt'))

    for srt in tqdm.tqdm(srt_list):
        ## 파일명 전처리 -> 폴더 구조별 동작이 다를 수 있으니, 확인 필요
        base_name = srt.rsplit('_', 2)[0]
        folder_name, file_name = base_name.rsplit('/', 1)

        ## srt와 wav가 함께 존재하지 않으면 fail list에 추가 후 continue
        if not os.path.isfile(base_name + '.wav'):
            not_exist_list.append(file_name)
            continue

        y, sr = sf.read(base_name + '.wav')

        subs = pysubs2.load(srt, encoding="utf-8")
        for idx, sub in enumerate(subs):
            start = int(sub.start / 1000 * sr)
            end = int(sub.end/1000*sr)

            ## 전사 내용 후처리 -> 데이터셋별 튜닝 필요.
            post = postprocess(sub.text)

            ## split file 이름 처리
            str_num = str(idx)
            str_num = '_' + ('0' * (3 - len(str_num))) + str_num
            tmp_file_name = file_name + str_num + ".wav"

            ## 전사 내용이 이상값이면 None, None이면 fail list에 추가
            if post is None:
                post_none_list.append([tmp_file_name, sub.text])
                continue
        
            sp_wav = y[start : end+1]
            sf.write(os.path.join(output_dir, tmp_file_name), sp_wav, sr)

            with open(os.path.join(output_dir, tmp_file_name.replace('wav', 'txt')), 'w+') as f:
                f.write(post + '\n')



        ## 직접 srt static하게 자르기
        # with open(srt, 'r') as f:
        #     text = ''
        #     start = 0
        #     end = 0

        #     b = False

        #     for idx, line in enumerate(f):
        #         if str.isdigit(line):
        #             b = True
        #             continue
                
        #         if idx % 4 == 1: ## time
        #             try:
        #                 start, end = map(time_convert, line.split(' --> '))
        #             except:
        #                 print(srt, '\t', line)
        #                 exit()
        #         elif idx % 4 == 2: ## transcribe
        #             str_num = str(idx//4)
        #             str_num = '_' + ('0' * (3 - len(str_num))) + str_num
        #             file_name = file_name + str_num + ".wav"

        #             post = postprocess(line.strip())
        #             if post is None:
        #                 post_none_list.append([file_name, line])
        #                 continue

        #             sp_wav = y[int(sr * start) : int(sr * end)+1]
        #             sf.write(os.path.join(output_dir, file_name), sp_wav, sr)

        #             with open(os.path.join(output_dir, file_name.replace('wav', 'txt')), 'w+') as f:
        #                 f.write(post + '\n')
    

    if post_none_list:
        print('postprocess none = ', len(post_none_list))
        txt = f'log/post_none_{time_str}.txt'

        with open(txt, 'w+') as f:
            for item in post_none_list:
                f.write("\t".join(item) + "\n")
    
    if not_exist_list:
        print('not exist = ', len(not_exist_list))
        txt = f'log/not_exist_{time_str}.txt'

        for item in post_none_list:
            with open(txt, 'w+') as f:
                f.write("\t".join(item) + "\n")