#!/usr/bin/python3 -I
"""
PDF 文字移除：解析內容流，移除字元中心落在指定矩形內的文字（含 Form XObject 內的文字）。
被移除的字元以 TJ 位移補償，同一行其餘文字位置不變。

用法： pdftool_textedit.py --in in.pdf --out out.pdf --rects rects.json
  rects.json = [{"page":1,"box":[x1,y1,x2,y2]}, ...]   box 為未旋轉、原點為頁面左下的使用者座標
輸出： {"ok":true,"counts":[每個矩形移除的字元數,...]}
"""
import argparse
import json
import sys

import pikepdf
from pikepdf import Array, Dictionary, Name, Operator, Stream, String

IDENT = (1.0, 0.0, 0.0, 1.0, 0.0, 0.0)


def mmul(a, b):
    """先套用 a 再套用 b（PDF 列向量慣例：a × b）"""
    return (
        a[0] * b[0] + a[1] * b[2], a[0] * b[1] + a[1] * b[3],
        a[2] * b[0] + a[3] * b[2], a[2] * b[1] + a[3] * b[3],
        a[4] * b[0] + a[5] * b[2] + b[4], a[4] * b[1] + a[5] * b[3] + b[5],
    )


def apply_pt(m, x, y):
    return (m[0] * x + m[2] * y + m[4], m[1] * x + m[3] * y + m[5])


def num(o, default=0.0):
    try:
        return float(o)
    except Exception:
        return default


class FontInfo:
    def __init__(self, fd):
        self.two = False
        self.scale = 1.0           # 字寬單位 → 1/1000 em
        self.dw = 500.0
        self.widths = {}
        self.ranges = []
        self.first = 0
        try:
            sub = str(fd.get('/Subtype', ''))
            if sub == '/Type0':
                self.two = True
                desc = fd.get('/DescendantFonts')[0]
                self.dw = num(desc.get('/DW', 1000), 1000.0)
                w = desc.get('/W')
                if w is not None:
                    self._parse_w(w)
            else:
                self.first = int(num(fd.get('/FirstChar', 0)))
                ws = fd.get('/Widths')
                if ws is not None:
                    self.widths = {self.first + i: num(v) for i, v in enumerate(ws)}
                    self.dw = 0.0
                    fdesc = fd.get('/FontDescriptor')
                    if fdesc is not None:
                        self.dw = num(fdesc.get('/MissingWidth', 0))
                if sub == '/Type3':
                    fm = fd.get('/FontMatrix')
                    self.scale = num(fm[0], 0.001) * 1000.0 if fm is not None else 1.0
        except Exception:
            pass

    def _parse_w(self, w):
        i, n = 0, len(w)
        while i < n:
            c = int(num(w[i]))
            nxt = w[i + 1] if i + 1 < n else None
            if isinstance(nxt, Array):
                for k, v in enumerate(nxt):
                    self.widths[c + k] = num(v)
                i += 2
            elif nxt is not None and i + 2 < n:
                self.ranges.append((c, int(num(nxt)), num(w[i + 2])))
                i += 3
            else:
                break

    def width(self, code):
        v = self.widths.get(code)
        if v is None:
            for a, b, ww in self.ranges:
                if a <= code <= b:
                    v = ww
                    break
        if v is None:
            v = self.dw
        return v * self.scale

    def codes(self, data):
        if self.two:
            i = 0
            while i < len(data):
                if i + 1 < len(data):
                    yield data[i:i + 2], (data[i] << 8) | data[i + 1]
                    i += 2
                else:
                    yield data[i:i + 1], data[i]
                    i += 1
        else:
            for b in data:
                yield bytes([b]), b


class State:
    def __init__(self, ctm):
        self.ctm = ctm
        self.stack = []
        self.tm = IDENT
        self.tlm = IDENT
        self.font = None
        self.fs = 0.0
        self.tc = 0.0
        self.tw = 0.0
        self.th = 1.0
        self.tl = 0.0
        self.rise = 0.0


class Editor:
    def __init__(self, pdf, rects):
        self.pdf = pdf
        self.rects = rects          # [(x1,y1,x2,y2)] 已含頁面偏移
        self.counts = [0] * len(rects)
        self.font_cache = {}

    def hit(self, x, y):
        for i, (x1, y1, x2, y2) in enumerate(self.rects):
            if x1 - 0.5 <= x <= x2 + 0.5 and y1 - 0.5 <= y <= y2 + 0.5:
                return i
        return -1

    def font_for(self, resources, name):
        try:
            fd = resources.get('/Font').get(name)
        except Exception:
            fd = None
        if fd is None:
            return FontInfo(Dictionary())
        key = id(fd) if not fd.is_indirect else fd.objgen
        if key not in self.font_cache:
            self.font_cache[key] = FontInfo(fd)
        return self.font_cache[key]

    # 處理一串指令；回傳 (新指令, 是否有變更)
    def process(self, ops, resources, ctm, holder, depth=0):
        st = State(ctm)
        out = []
        changed = False
        for operands, op in ops:
            o = str(op)
            if o == 'q':
                st.stack.append((st.ctm, st.tm, st.tlm, st.font, st.fs, st.tc, st.tw, st.th, st.tl, st.rise))
            elif o == 'Q':
                if st.stack:
                    st.ctm, st.tm, st.tlm, st.font, st.fs, st.tc, st.tw, st.th, st.tl, st.rise = st.stack.pop()
            elif o == 'cm' and len(operands) == 6:
                st.ctm = mmul(tuple(num(v) for v in operands), st.ctm)
            elif o == 'BT':
                st.tm = IDENT
                st.tlm = IDENT
            elif o == 'Tf' and len(operands) == 2:
                st.font = self.font_for(resources, str(operands[0]))
                st.fs = num(operands[1])
            elif o == 'Tc':
                st.tc = num(operands[0])
            elif o == 'Tw':
                st.tw = num(operands[0])
            elif o == 'Tz':
                st.th = num(operands[0], 100.0) / 100.0
            elif o == 'TL':
                st.tl = num(operands[0])
            elif o == 'Ts':
                st.rise = num(operands[0])
            elif o in ('Td', 'TD') and len(operands) == 2:
                tx, ty = num(operands[0]), num(operands[1])
                if o == 'TD':
                    st.tl = -ty
                st.tlm = mmul((1, 0, 0, 1, tx, ty), st.tlm)
                st.tm = st.tlm
            elif o == 'T*':
                st.tlm = mmul((1, 0, 0, 1, 0, -st.tl), st.tlm)
                st.tm = st.tlm
            elif o == 'Tm' and len(operands) == 6:
                st.tm = tuple(num(v) for v in operands)
                st.tlm = st.tm
            elif o in ('Tj', 'TJ', "'", '"'):
                pre = []
                if o == "'":
                    st.tlm = mmul((1, 0, 0, 1, 0, -st.tl), st.tlm)
                    st.tm = st.tlm
                    pre = [([], Operator('T*'))]
                    items = [operands[0]]
                elif o == '"':
                    st.tw, st.tc = num(operands[0]), num(operands[1])
                    st.tlm = mmul((1, 0, 0, 1, 0, -st.tl), st.tlm)
                    st.tm = st.tlm
                    pre = [([operands[0]], Operator('Tw')), ([operands[1]], Operator('Tc')), ([], Operator('T*'))]
                    items = [operands[2]]
                elif o == 'Tj':
                    items = [operands[0]]
                else:
                    items = list(operands[0])
                new_items, removed = self.show(st, items)
                if removed:
                    changed = True
                    out.extend(pre)
                    out.append(([Array(new_items)], Operator('TJ')))
                else:
                    if o in ("'", '"'):
                        out.extend(pre)
                        out.append(([items[0]], Operator('Tj')))
                    else:
                        out.append((operands, op))
                continue
            elif o == 'Do' and depth < 6 and len(operands) == 1:
                done = self.do_form(operands[0], resources, st.ctm, holder, depth)
                if done:
                    changed = True
            out.append((operands, op))
        return out, changed

    def show(self, st, items):
        """回傳 (新的 TJ 陣列元素, 是否移除了字元)，並更新 st.tm"""
        font = st.font or FontInfo(Dictionary())
        fs, th = st.fs, st.th
        new = []
        buf = b''
        removed_any = False
        for it in items:
            if isinstance(it, String):
                data = bytes(it)
                for chunk, code in font.codes(data):
                    w = font.width(code)
                    tx = (w / 1000.0 * fs + st.tc + (st.tw if (len(chunk) == 1 and code == 32) else 0.0)) * th
                    cx, cy = w / 2000.0 * fs * th, 0.3 * fs + st.rise
                    px, py = apply_pt(mmul(st.tm, st.ctm), cx, cy)
                    idx = self.hit(px, py) if fs else -1
                    if idx >= 0 and fs * th != 0:
                        self.counts[idx] += 1
                        removed_any = True
                        if buf:
                            new.append(String(buf))
                            buf = b''
                        new.append(-tx * 1000.0 / (fs * th))
                    else:
                        buf += chunk
                    st.tm = mmul((1, 0, 0, 1, tx, 0), st.tm)
            else:
                n = num(it)
                if buf:
                    new.append(String(buf))
                    buf = b''
                new.append(it)
                st.tm = mmul((1, 0, 0, 1, -n / 1000.0 * fs * th, 0), st.tm)
        if buf:
            new.append(String(buf))
        return new, removed_any

    def do_form(self, name, resources, ctm, holder, depth):
        try:
            xo = resources.get('/XObject')
            form = xo.get(name)
        except Exception:
            return False
        if form is None or str(form.get('/Subtype', '')) != '/Form':
            return False
        m = form.get('/Matrix')
        fm = tuple(num(v) for v in m) if m is not None and len(m) == 6 else IDENT
        sub_res = form.get('/Resources') or resources
        try:
            sub_ops = list(pikepdf.parse_content_stream(form))
        except Exception:
            return False
        own = Dictionary({k: v for k, v in sub_res.items()})
        sx = sub_res.get('/XObject')
        own_xo = Dictionary({k: v for k, v in sx.items()}) if sx is not None else Dictionary()
        own['/XObject'] = own_xo
        new_ops, changed = self.process(sub_ops, own, mmul(fm, ctm), {'xobj': own_xo}, depth + 1)
        if not changed:
            return False
        # 以複本取代，避免影響共用此 Form 的其他頁面
        new_form = self.pdf.make_stream(pikepdf.unparse_content_stream(new_ops))
        for k, v in form.items():
            if k not in ('/Length', '/Filter', '/DecodeParms'):
                new_form[k] = v
        new_form['/Resources'] = self.pdf.make_indirect(own)
        holder['xobj'][name] = self.pdf.make_indirect(new_form)
        return True


def find_resources(page):
    node = page.obj
    while node is not None:
        r = node.get('/Resources')
        if r is not None:
            return r
        node = node.get('/Parent')
    return Dictionary()


def edit_page(pdf, page, rects_user, ed):
    ox, oy = 0.0, 0.0
    box = page.obj.get('/CropBox') or page.obj.get('/MediaBox')
    if box is not None:
        ox, oy = num(box[0]), num(box[1])
    ed.rects = [(r[0] + ox, r[1] + oy, r[2] + ox, r[3] + oy) for r in rects_user]
    res = find_resources(page)
    # 為此頁建立自己的 Resources / XObject 複本，之後替換 Form 才不會影響其他頁
    own_res = Dictionary({k: v for k, v in res.items()})
    xo = res.get('/XObject')
    own_xo = Dictionary({k: v for k, v in xo.items()}) if xo is not None else Dictionary()
    own_res['/XObject'] = own_xo
    ops = list(pikepdf.parse_content_stream(page))
    holder = {'xobj': own_xo}
    new_ops, changed = ed.process(ops, own_res, IDENT, holder, 0)
    if changed:
        page.obj['/Resources'] = pdf.make_indirect(own_res)
        page.obj['/Contents'] = pdf.make_stream(pikepdf.unparse_content_stream(new_ops))
    return changed


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--in', dest='inp', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--rects', required=True)
    a = ap.parse_args()
    try:
        with open(a.rects, encoding='utf-8') as f:
            rects = json.load(f)
        pdf = pikepdf.open(a.inp)
        ed = Editor(pdf, [])
        counts = [0] * len(rects)
        by_page = {}
        for i, r in enumerate(rects):
            by_page.setdefault(int(r['page']), []).append((i, [float(v) for v in r['box']]))
        changed_pages = 0
        for pno, lst in by_page.items():
            if pno < 1 or pno > len(pdf.pages):
                continue
            ed.counts = [0] * len(lst)
            if edit_page(pdf, pdf.pages[pno - 1], [b for _, b in lst], ed):
                changed_pages += 1
            for k, (i, _) in enumerate(lst):
                counts[i] = ed.counts[k]
        pdf.save(a.out)
        print(json.dumps({'ok': True, 'counts': counts, 'pages_changed': changed_pages}))
    except Exception as e:
        print(json.dumps({'ok': False, 'error': str(e)[:300]}, ensure_ascii=False))
        sys.exit(2)


if __name__ == '__main__':
    main()