from tokenize import Number

import unicodedata
from difflib import SequenceMatcher

import sys
if len(sys.argv) < 2:
    print("python my_ocr.py <img path>")
    exit()

img_path = sys.argv[1]
llama_cli = "/home/xxx/data/llama.cpp/build/bin/llama-cli"

# https://github.com/datalab-to/surya/blob/11d18847e8b8c17fd55e254e3e5c44b3fed5c87e/surya/inference/prompts.py#L82
prompt1 = """OCR this image to HTML. Each block is a div with data-label and data-bbox (x0 y0 x1 y1, normalized 0-1000)."""
#prompt2 = "Extract all information from the main body of the document image and represent it in HTML format."
prompt2 = "Extract all."

model1_gguf = "/media/xxx/data/ocr/surya-2.gguf"
mmproj1_gguf = "/media/xxx/data/ocr/surya-2-mmproj.gguf"
model2_gguf = "/media/xxx/data/ocr/HunyuanOCR-bf16.gguf"
mmproj2_gguf = "/media/xxx/data/ocr/mmproj-HunyuanOCR-bf16.gguf"
import subprocess

# ImageMagickで左右にぶった切る
subprocess.run(["convert", img_path, "-crop","2x1@", "+repage", img_path + "_%d.png"])

for n in range(2):
    img_path_n = img_path + f"_{n}.png"
    subprocess.run(["python3", "../deskew_HT/run_deskew.py", img_path_n,"-o", img_path_n])
#    subprocess.run(["python3", "../deskew_HT/run_deskew.py", img_path_n,"-o", img_path_n]) # second deskew

    result1 = subprocess.run([llama_cli, "-p", prompt1, "--image", img_path_n, "-m", model1_gguf,
                                                    "--mmproj", mmproj1_gguf, "--temp", "0", "--single-turn", "--log-disable", "--no-display-prompt",
                                                    "--simple-io", "-c", "8192"],
                                                    capture_output=True, text=True)

    t = result1.stdout.find("Loaded media")
    t = result1.stdout.find(">", t)
    t = result1.stdout.find("\n", t)
    r1 = result1.stdout[t:]
    print(r1)
    result2 = subprocess.run([llama_cli, "-p", prompt2, "--image", img_path_n, "-m", model2_gguf,
                                                    "--mmproj", mmproj2_gguf, "--temp", "0", "--single-turn", "--log-disable", "--no-display-prompt",
                                                    "--simple-io", "-c", "8192"],
                                                    capture_output=True, text=True)

    t = result2.stdout.find("Loaded media")
    t = result2.stdout.find(">", t)
    t = result2.stdout.find("\n", t)
    r2 = result2.stdout[t:]
    r2 = r2.replace("\n", "")
    print(r2)

    def is_kana(c):
        name = unicodedata.name(c, "")
        return "HIRAGANA" in name or "KATAKANA" in name

    def is_kanji(c):
        name = unicodedata.name(c, "")
        return "CJK UNIFIED IDEOGRAPH" in name

    def merge_ocr(ocr_kanji, ocr_kana):
        sm = SequenceMatcher(None, ocr_kanji, ocr_kana, autojunk=False)

        result = []

        for tag, i1, i2, j1, j2 in sm.get_opcodes():

            if tag == "equal":
                result.extend(ocr_kanji[i1:i2])
                continue

            if tag == "replace":

                a = ocr_kanji[i1:i2]
                b = ocr_kana[j1:j2]

                # かなブロック単位で扱う
                # まず「かな側を優先的に採用」
                if any(is_kana(c) for c in b):
                    result.extend(b)
                else:
                    result.extend(a)

                continue

            if tag == "delete":
                result.extend(ocr_kanji[i1:i2])
                continue

            if tag == "insert":
                result.extend(ocr_kana[j1:j2])
                continue

        # 後処理：二重ひらがな吸収
        cleaned = []
        prev = ""

        for c in result:
            if is_kana(c) and is_kana(prev) and c == prev:
                continue
            cleaned.append(c)
            prev = c

        return "".join(cleaned)

    print(merge_ocr(r2, r1))