import fitz  # PyMuPDF
import os
import random
import tkinter as tk
from tkinter import ttk, filedialog, messagebox

# 🔧 주석 및 하이라이트 회색 처리 (확률 기반)
def create_modified_page(original_page, probability):
    temp_doc = fitz.open()
    temp_doc.insert_pdf(original_page.parent, from_page=original_page.number, to_page=original_page.number)
    page = temp_doc[0]

    annot = page.first_annot
    while annot:
        if annot.type[1] in ("Square", "Highlight"):
            # 확률 계산
            should_modify = False
            if probability == 100:
                should_modify = True
            elif probability == 0:
                should_modify = False
            else:
                rand_value = random.randint(1, 100)
                should_modify = rand_value <= probability

            if should_modify:
                if annot.type[1] == "Highlight":
                    # ✅ 하이라이트 색: 검정색, 불투명도 100%
                    annot.set_colors(stroke=(0, 0, 0), fill=(0, 0, 0))
                    annot.set_opacity(1.0)
                    annot.update()

                    # ✅ 해당 영역 글자 회색으로 덮어쓰기
                    quad_points = annot.vertices
                    quads = [fitz.Quad(quad_points[i:i+4]).rect for i in range(0, len(quad_points), 4)]

                    for rect in quads:
                        words = page.get_text("words", clip=rect)
                        for w in words:
                            word = w[4]
                            word_rect = fitz.Rect(w[0], w[1], w[2], w[3])
                            page.insert_textbox(
                                word_rect,
                                word,
                                fontsize=11,
                                fontname="helv",
                                color=(0.6, 0.6, 0.6),  # 회색 글씨
                                overlay=True
                            )
                else:  # Square 주석 처리
                    annot.set_colors(stroke=(0.5, 0.5, 0.5), fill=(0.7, 0.7, 0.7))  # 회색
                    annot.set_opacity(1.0)
                    annot.update()
        annot = annot.next

    return temp_doc

# 📄 PDF 처리 함수
def process_pdf(file_path, include_original=False, number_value=100):
    original_doc = fitz.open(file_path)
    result_doc = fitz.open()

    for page_index in range(len(original_doc)):
        page = original_doc[page_index]

        # 수정본 추가 (확률 반영)
        modified_doc = create_modified_page(page, number_value)
        result_doc.insert_pdf(modified_doc)
        print(f"{page_index + 1}페이지: 수정본 추가됨")

        # 원본 포함 여부에 따라 추가
        if include_original:
            result_doc.insert_pdf(original_doc, from_page=page_index, to_page=page_index)
            print(f"{page_index + 1}페이지: 원본 추가됨")

    # 저장 경로 생성
    base, ext = os.path.splitext(file_path)
    suffix = "-with-original" if include_original else "-blank-only"
    new_path = f"{base}{suffix}{ext}"

    try:
        result_doc.save(new_path)
        result_doc.close()
        print(f"\n✅ PDF 저장 완료: {new_path}")
        messagebox.showinfo("완료", f"PDF 저장 완료:\n{new_path}")
    except Exception as e:
        messagebox.showerror("에러", f"저장 중 오류 발생:\n{e}")

# 📂 PDF 파일 선택
def select_pdf_file():
    return filedialog.askopenfilename(
        title="PDF 파일 선택",
        filetypes=[("PDF Files", "*.pdf")]
    )

# 🔢 숫자 유효성 검사
def validate_number(number_str):
    try:
        value = int(number_str)
        if 0 <= value <= 100:
            return value
        else:
            raise ValueError
    except ValueError:
        messagebox.showerror("입력 오류", "0부터 100 사이의 확률을 입력하세요.")
        return None

# 🖥️ GUI 구성 및 실행
def start_gui():
    root = tk.Tk()
    root.title("PDF 문제집 생성기")
    root.geometry("450x320")

    selected_option = tk.StringVar(value="blank_only")
    number_input = tk.StringVar(value="100")  # ✅ 초기 확률 100%

    # ▶ 버튼 클릭 시 실행
    def run_process():
        file_path = select_pdf_file()
        if not file_path:
            print("❌ 파일이 선택되지 않았습니다.")
            return

        number_value = validate_number(number_input.get())
        if number_value is None:
            return

        include_original = selected_option.get() == "with_original"
        print(f"\n📂 선택된 파일: {file_path}")
        print(f"🧾 선택 옵션: {'원본 포함' if include_original else '빈칸만'}")
        print(f"🔢 입력한 확률: {number_value}%")
        process_pdf(file_path, include_original=include_original, number_value=number_value)

    # 📋 UI 요소 구성
    ttk.Label(root, text="PDF 생성 옵션을 선택하세요:").pack(pady=10)
    ttk.Radiobutton(root, text="빈칸 문제집만", variable=selected_option, value="blank_only").pack(anchor="w", padx=20)
    ttk.Radiobutton(root, text="원본 포함", variable=selected_option, value="with_original").pack(anchor="w", padx=20)

    ttk.Label(root, text="하이라이트 처리 확률을 입력하세요 (0 ~ 100):").pack(pady=10)
    ttk.Entry(root, textvariable=number_input).pack(pady=5)

    ttk.Button(root, text="📄 PDF 파일 선택 및 생성", command=run_process).pack(pady=20)

    root.mainloop()

# ▶ 메인 실행
if __name__ == "__main__":
    start_gui()