#!/usr/bin/env python3
"""MinerU RDNA 适配补丁 — 一键应用脚本"""

import re
import os
import sys


def find_infer_dir():
    import mineru.model.utils.tools.infer
    return os.path.dirname(mineru.model.utils.tools.infer.__file__)


def patch_a_predict_rec_imgw(infer_dir):
    """predict_rec.py: imgW 对齐到 32"""
    f = os.path.join(infer_dir, 'predict_rec.py')
    c = open(f).read()
    old = r'(^ *)(imgW = max\(min\(imgW, self\.limited_max_width\), self\.limited_min_width\)\n)'
    new = r'\1\2\1imgW = math.ceil(imgW / 32) * 32\n'
    c2 = re.sub(old, new, c, flags=re.MULTILINE)
    if c2 == c:
        if 'math.ceil(imgW / 32)' not in c:
            raise RuntimeError('Patch A: cannot find imgW line')
        print('Patch A: already applied')
    else:
        open(f, 'w').write(c2)
        print('Patch A: imgW 32-align inserted')


def patch_b_predict_rec_batch(infer_dir):
    """predict_rec.py: 批次填充"""
    f = os.path.join(infer_dir, 'predict_rec.py')
    c = open(f).read()
    # 兼容 mineru 3.4.0：赋值行可能是
    #   norm_img_batch = np.concatenate(norm_img_batch)
    # 或
    #   norm_img_batch = np.ascontiguousarray(np.concatenate(norm_img_batch))
    # 用宽松正则匹配「norm_img_batch = ...np.concatenate(norm_img_batch)...」整行
    old = (
        r'(^ *)(norm_img_batch\s*=\s*'
        r'(?:np\.ascontiguousarray\()?'           # 可选 ascontiguousarray 包裹
        r'np\.concatenate\(norm_img_batch\)'
        r'(?:\))?'                                # 可选闭合括号
        r')'
    )
    new = (
        r'\1actual_batch_size = len(norm_img_batch)\n'
        r'\1if actual_batch_size < batch_num:\n'
        r'\1    pad_size = batch_num - actual_batch_size\n'
        r'\1    pad_img = np.zeros_like(norm_img_batch[0])\n'
        r'\1    for _ in range(pad_size):\n'
        r'\1        norm_img_batch.append(pad_img)\n'
        r'\1\2'
    )
    c2 = re.sub(old, new, c, flags=re.MULTILINE)
    if c2 == c:
        if 'actual_batch_size' not in c:
            raise RuntimeError('Patch B: cannot find norm_img_batch concatenation')
        print('Patch B: already applied')
    else:
        open(f, 'w').write(c2)
        print('Patch B: batch padding inserted')
    # 修改 range(len(rec_result)) → range(actual_batch_size)
    c3 = open(f).read()
    c4 = re.sub(
        r'( +)for rno in range\(len\(rec_result\)\):',
        r'\1for rno in range(actual_batch_size):',
        c3
    )
    open(f, 'w').write(c4)


def patch_c_predict_det_contiguous(infer_dir):
    """predict_det.py: contiguous 检查"""
    f = os.path.join(infer_dir, 'predict_det.py')
    c = open(f).read()
    old = r'(^ *)(inp = inp\.to\(self\.device\)\n)'
    new = r'\1\2\1if not inp.is_contiguous():\n\1    inp = inp.contiguous()\n'
    c2 = re.sub(old, new, c, flags=re.MULTILINE)
    if c2 == c:
        if 'is_contiguous' not in c:
            raise RuntimeError('Patch C: cannot find inp.to(device) line')
        print('Patch C: already applied')
    else:
        open(f, 'w').write(c2)
        print('Patch C: contiguous check inserted')


def main():
    infer_dir = find_infer_dir()
    print(f'MinerU infer dir: {infer_dir}')
    patch_a_predict_rec_imgw(infer_dir)
    patch_b_predict_rec_batch(infer_dir)
    patch_c_predict_det_contiguous(infer_dir)
    print('All MinerU RDNA patches applied.')


if __name__ == '__main__':
    main()