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

def extract_spans(page):
    blocks = page.get_text("dict")["blocks"]
    spans = []

    sub_para_num = 0
    sub_words_count_before = 0
    sub_words_count_after = 0
    sub_API_CALLING_TIMES = 0
    sub_num_cache_hit = 0
    sub_num_cache_miss = 0
    for b in blocks:
        if b["type"] != 0:
            continue
        for line in b["lines"]:
            for span in line["spans"]:
                text = span["text"].strip()
                if not text:
                    continue

                #results, para_num, words_count_before, words_count_after, HF_API_CALLING_TIMES, num_cache_hit，num_cache_miss
                tranlated_result, para_num, words_count_before, words_count_after, HF_API_CALLING_TIMES, num_cache_hit, num_cache_miss = translate_preserve_paragraphs(text)
                sub_para_num += para_num
                sub_words_count_before += words_count_before
                sub_words_count_after += words_count_after
                sub_API_CALLING_TIMES += HF_API_CALLING_TIMES
                sub_num_cache_hit += num_cache_hit
                sub_num_cache_miss += num_cache_miss

                spans.append({
                    "text": tranlated_result,
                    "size": span["size"],
                })

    return spans, sub_para_num, sub_words_count_before ,sub_words_count_after, sub_API_CALLING_TIMES, sub_num_cache_hit, sub_num_cache_miss


def classify_page_text(spans):
    sizes = [s["size"] for s in spans]
    if not sizes:
        return []

    body_size = max(set(sizes), key=sizes.count)

    result = []
    for s in spans:
        role = "title" if s["size"] > body_size * 1.3 else "body"
        result.append({
            "text": s["text"],
            "role": role
        })

    return result


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

    pages = []

    total_para_num = 0
    total_words_count_before = 0 
    total_words_count_after = 0
    total_HF_API_CALLING_TIMES = 0
    total_num_cache_hit = 0
    total_num_cache_miss = 0
    for i, page in enumerate(doc):
        spans, para_num, words_count_before, words_count_after, HF_API_CALLING_TIMES, num_cache_hit, num_cache_miss = extract_spans(page)
        total_para_num += para_num
        total_words_count_before += words_count_before 
        total_words_count_after += words_count_after
        total_HF_API_CALLING_TIMES += HF_API_CALLING_TIMES
        total_num_cache_hit += num_cache_hit
        total_num_cache_miss += num_cache_miss
        if not spans:
            continue

        classified = classify_page_text(spans)

        pages.append({
            "id": f"page_{i}",
            "items": classified
        })

    return pages, total_para_num, total_words_count_before, total_words_count_after, total_HF_API_CALLING_TIMES, total_num_cache_hit, total_num_cache_miss


def render_pdf(pages, output_path: str):
    doc = fitz.open()

    for page in pages:
        new_page = doc.new_page()
        y = 50  # 当前写入的纵坐标

        for item in page["items"]:
            if item["role"] == "title":
                font = "Helvetica-Bold"
                size = 14
            else:
                font = "Helvetica"
                size = 12

            rect = fitz.Rect(50, y, 545, y + 100)

            new_page.insert_textbox(
                rect,
                item["text"],
                fontname=font,
                fontsize=size,
                lineheight=1.3,
            )

            
            y += size * 2
    try:
        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_pdf_text(input_path: str):
    try:
        doc = fitz.open(input_path)
    except Exception as e:
        print(f"\n\n {Fore.RED}{Style.BRIGHT}>{Style.RESET_ALL} [GoodGute] Error: {e}")
        sys.exit(1)

    all_text = []

    for page in doc:
        all_text.append(page.get_text("text").strip())

    doc.close()

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

    return full_text

def translate_pdf(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__":
    #full_text, page_count = extract_pdf_text("./tests/test1.pdf")

    pages, total_para_num, total_words_count_before, total_words_count_after, total_API_CALLING_TIMES, total_num_cache_hit, total_num_cache_miss = extract_and_process_pdf("./tests/test1.pdf")
    print(f"para_num:{total_para_num}, words_count_before:{total_words_count_before}, words_count_after:{total_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}")
    render_pdf(pages,"./tests/output1.pdf")