import os
import pandas as pd
from transformers import CLIPTokenizer
import tkinter as tk
from tkinter import filedialog, scrolledtext
import threading

# CLIPTokenizer 초기화
s_token = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")

def rcs(text):
    if text is None:
        return None
    token_ids = s_token.encode(text)
    return len(token_ids)

def process_files(dflist, text_box, select_button, process_button):
    # /tags 하위 폴더 만들기
    output_dir = "processed"
    os.makedirs(output_dir, exist_ok=True)

    for i, _df in enumerate(dflist):
        # parquet 파일 읽기
        df = pd.read_parquet(_df, engine="pyarrow")

        # 'tokens' 열 추가
        tokens = []
        total = len(df)
        for idx, text in enumerate(df['general']):
            if text is not None:
                tokens.append(rcs(text))

                # 진행 상황 출력
                if (idx + 1) % 100 == 0 or idx + 1 == total:
                    progress = f"Processing file {_df}: {idx + 1}/{total} ({(idx + 1) / total * 100:.2f}%)\n"
                    text_box.insert(tk.END, progress)
                    text_box.see(tk.END)
            else:
                tokens.append(None)

        df['tokens'] = tokens

        # 처리된 파일 저장 경로
        output_path = os.path.join(output_dir, os.path.basename(_df))

        # parquet 파일로 저장
        df.to_parquet(output_path, engine="pyarrow")
        text_box.insert(tk.END, f"Finished processing {_df}\n")
        text_box.see(tk.END)

    text_box.insert(tk.END, "모든 파일이 성공적으로 처리되었습니다.\n")
    text_box.see(tk.END)
    
    # 작업이 모두 종료되면 output_dir 윈도우 폴더가 열린다.
    os.startfile(output_dir)

    # dflist 초기화
    dflist = []

    # 버튼 다시 활성화
    select_button.config(state=tk.NORMAL)
    process_button.config(state=tk.NORMAL)

def select_files():
    file_paths = filedialog.askopenfilenames(filetypes=[("Parquet files", "*.parquet")])
    if file_paths:
        dflist.extend(file_paths)
        text_box.insert(tk.END, f"Selected files:\n{file_paths}\n")
        text_box.see(tk.END)

def start_processing():
    if not dflist:
        return
    select_button.config(state=tk.DISABLED)
    process_button.config(state=tk.DISABLED)

    # 파일 처리 스레드 시작
    threading.Thread(target=process_files, args=(dflist, text_box, select_button, process_button)).start()

# Tkinter UI 설정
root = tk.Tk()
root.title("프롬프트 스태커용 parquet 파일 토큰 업데이트 도구")

frame = tk.Frame(root)
frame.pack(padx=10, pady=10)

select_button = tk.Button(frame, text="Parquet 파일 선택", command=select_files)
select_button.pack(side=tk.LEFT, padx=5, pady=5)

process_button = tk.Button(frame, text="토큰 계산 시작", command=start_processing)
process_button.pack(side=tk.LEFT, padx=5, pady=5)

text_box = scrolledtext.ScrolledText(root, width=80, height=20)
text_box.pack(padx=10, pady=10)

# dflist 초기화
dflist = []

root.mainloop()