from docx import Document
from docx.shared import Pt


from colorama import Fore, Style, init
init(autoreset=True)
from .translator import translate_preserve_paragraphs

from collections import Counter
import sys


def get_paragraph_main_font_size(para):
    sizes = []

    for run in para.runs:
        if run.font.size is not None:
            sizes.append(run.font.size.pt)

    if not sizes:
        return None  # 全部继承样式，拿不到字号

    return Counter(sizes).most_common(1)[0][0]

def classify_role(p_size, body_size):
    if body_size is None or p_size is None:
        return "body"
    return "title" if p_size > body_size * 1.25 else "body"

def extract_and_translate_docx(input_path: str):
    try:
        doc = Document(input_path)
    except Exception as e:
        print(f"\n\n {Fore.RED}{Style.BRIGHT}>{Style.RESET_ALL} [GoodGute] Error: {e}")
        sys.exit(1)
    # 先收集所有段落字号（用于找 body 的基准字号）
    para_infos = []

    for para in doc.paragraphs:
        text = para.text.strip()
        if not text:
            continue

        size = get_paragraph_main_font_size(para)

        para_infos.append({
            "text": text,
            "size": size,
        })

    # 统计 body 的主字号（出现次数最多的那个）
    sizes = [p["size"] for p in para_infos if p["size"] is not None]
    body_size = Counter(sizes).most_common(1)[0][0] if sizes else None

    items = []

    total_API_CALLING_TIMES = 0
    total_num_cache_hit = 0
    total_num_cache_miss = 0
    total_para_num = 0
    total_words_before = 0
    total_words_after = 0
    for p in para_infos:
        # return results, para_num, words_count_before, words_count_after, HF_API_CALLING_TIMES, num_cache_hit，num_cache_miss
        translated, para_num, words_count_before, words_count_after, HF_API_CALLING_TIMES, num_cache_hit, num_cache_miss = translate_preserve_paragraphs(p["text"])

        total_API_CALLING_TIMES += HF_API_CALLING_TIMES
        total_num_cache_hit += num_cache_hit
        total_num_cache_miss += num_cache_miss
        total_para_num += para_num
        total_words_before += words_count_before
        total_words_after += words_count_after

        role = classify_role(
            p_size=p["size"],
            body_size=body_size
        )

        items.append({
            "text": translated,
            "role": role,
            "size": p["size"]
        })

    return items, total_para_num, total_words_before, total_words_after, total_API_CALLING_TIMES, total_num_cache_hit, total_num_cache_miss



def render_docx(items, output_path: str):
    new_doc = Document()

    for item in items:
        p = new_doc.add_paragraph()
        run = p.add_run(item["text"])

        if item["role"] == "title":
            run.bold = True
            run.font.size = Pt(14)
        else:
            run.font.size = Pt(11)
    try:
        new_doc.save(output_path)
    except Exception as e:
        print(f"\n\n {Fore.RED}{Style.BRIGHT}>{Style.RESET_ALL} [GoodGute] Error: {e}")
        sys.exit(1)

def extract_docx_text(input_path: str):
    try:
        doc = Document(input_path)
    except Exception as e:
        print(f"\n\n {Fore.RED}{Style.BRIGHT}>{Style.RESET_ALL} [GoodGute] Error: {e}")
        sys.exit(1)
    
    paragraphs = []

    for para in doc.paragraphs:
        text = para.text.strip()
        if text:
            paragraphs.append(text)

    full_text = "\n\n".join(paragraphs)

    return full_text

def translate_docx(text):

    translated_full, para_num, words_before, words_after, \
    api_calls, cache_hit, cache_miss = \
        translate_preserve_paragraphs(text)

    return translated_full
if __name__ == "__main__":
    items, para_num, words_count_before, words_count_after, total_API_CALLING_TIMES, total_num_cache_hit, total_num_cache_miss = extract_and_translate_docx("./tests/test1.docx")
    render_docx(items, "./tests/output1.docx")
    print(f"para_num:{para_num}, words_count_before:{words_count_before}, words_count_after:{words_count_after}, total_API_CALLING_TIMES:{total_API_CALLING_TIMES}, total_num_cache_hit:{total_num_cache_hit}, total_num_cache_miss:{total_num_cache_miss}")