"""Rebuild the Level 2 booklet with Ukrainian text.

Strategy: for every translated block, redact the original glyphs (text only -
line art, boxes and images are left untouched) and re-typeset the Ukrainian
text at the original baselines using the same Arial faces, sizes and colours.
"""
import json
import os
import re
import sys
from collections import defaultdict

import fitz

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, HERE)

from lessons import CONTENT_RIGHT, FONTFILES, PAGE_RIGHT, Lesson  # noqa: E402

ALIAS = {'ArialMT': 'uaR', 'Arial-BoldMT': 'uaB', 'Arial-ItalicMT': 'uaI',
         'Arial-BoldItalicMT': 'uaBI', 'Arial-Black': 'uaBLK',
         'HelveticaNeue-Bold': 'uaB', 'ZapfDingbatsITC': 'uaU',
         'SassoonInfantDtB': 'uaDOT'}
_FONTS = {k: fitz.Font(fontfile=v) for k, v in FONTFILES.items()}

DEFAULT_RIGHT = CONTENT_RIGHT   # колонка тексту закінчується перед боксами балів
MIN_SCALE = 0.72        # never shrink a line below this fraction of its size
FIT_MIN = 0.55          # speech bubbles may shrink further - they must fit


def rgb(c):
    return ((c >> 16 & 255) / 255, (c >> 8 & 255) / 255, (c & 255) / 255)


def measure(font, text, size):
    return font.text_length(text, fontsize=size)


TAG = re.compile(r'\{s(\d+)\}')


def parse_rich(text, default_style):
    """'plain {s1}red{s0} plain' -> [(word, style), ...] keeping \n as a token."""
    words, style = [], default_style
    pos = 0
    chunks = []
    for m in TAG.finditer(text):
        chunks.append((text[pos:m.start()], style))
        style = int(m.group(1))
        pos = m.end()
    chunks.append((text[pos:], style))
    for chunk, st in chunks:
        for i, part in enumerate(chunk.split('\n')):
            if i:
                words.append(('\n', st))
            for w in part.split(' '):
                if w:
                    words.append((w, st))
    return words


def wrap_rich(words, styles, width, scale):
    """Greedy wrap of styled words; returns list of lines, each [(word, style)]."""
    lines, line, w_used = [], [], 0.0
    for word, st in words:
        if word == '\n':
            lines.append(line)
            line, w_used = [], 0.0
            continue
        f, size = styles[st]
        ww = measure(f, word, size * scale)
        sp = measure(f, ' ', size * scale) if line else 0.0
        if line and w_used + sp + ww > width:
            lines.append(line)
            line, w_used = [(word, st)], ww
        else:
            line.append((word, st))
            w_used += sp + ww
    lines.append(line)
    return lines


def line_width(line, styles, scale):
    w = 0.0
    for i, (word, st) in enumerate(line):
        f, size = styles[st]
        if i:
            w += measure(f, ' ', size * scale)
        w += measure(f, word, size * scale)
    return w


class Builder:
    def __init__(self, translations, lesson, covers=None):
        self.tr = translations
        self.lesson = lesson
        self.covers = covers
        self.doc = fitz.open(lesson.src)
        self.tpl = json.load(open(lesson.template, encoding='utf-8'))
        self.warnings = []
        self.drawn = defaultdict(list)   # block key -> rects of emitted text
        self.cur_key = None
        self.cur_page = 0
        self.cur_render = 0
        self.cur_border = 0.9

    # ---------------------------------------------------------------- erase
    def erase(self, page, blocks):
        for b in blocks:
            for ln in b['lines']:
                for sp in ln['spans']:
                    r = fitz.Rect(sp['x'] - 0.6, ln['bbox'][1] - 0.4,
                                  sp['x'] + sp['w'] + 0.6, ln['bbox'][3] + 0.4)
                    page.add_redact_annot(r)
        if page.first_annot:
            page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_NONE,
                                  graphics=fitz.PDF_REDACT_LINE_ART_NONE,
                                  text=fitz.PDF_REDACT_TEXT_REMOVE)

    # --------------------------------------------------------------- render
    def style(self, blk, idx):
        st = blk['styles'][min(idx, len(blk['styles']) - 1)]
        return _FONTS[st['font']], ALIAS[st['font']], FONTFILES[st['font']], st['size'], rgb(st['color'])

    def draw(self, page, x, y, text, font_alias, fontfile, size, color):
        if self.cur_key and text.strip():
            fnt = _FONTS[[k for k, v in ALIAS.items() if v == font_alias][0]]
            w = fnt.text_length(text, fontsize=size)
            self.drawn[self.cur_key].append(
                [round(x, 1), round(y - size * 0.78, 1), round(x + w, 1), round(y + size * 0.22, 1)])
        kw = {}
        if self.cur_render:
            kw = {'render_mode': self.cur_render, 'border_width': self.cur_border / max(size, 1e-6),
                  'stroke_opacity': 1}
            page.insert_text((x, y), text, fontname=font_alias, fontfile=fontfile,
                             fontsize=size, color=color, render_mode=self.cur_render,
                             border_width=self.cur_border / max(size, 1e-6))
            return
        page.insert_text((x, y), text, fontname=font_alias, fontfile=fontfile,
                         fontsize=size, color=color, render_mode=0)

    def render_block(self, page, blk, spec, key):
        self.cur_key = key
        self.cur_render = spec.get('render', 0)
        self.cur_border = spec.get('border', 0.9)
        right = spec.get('right', max(blk['bbox'][2], DEFAULT_RIGHT))
        right = min(right, PAGE_RIGHT)
        if spec.get('fit'):
            self.render_fit(page, blk, spec, key)
        elif 'flow' in spec:
            self.render_flow(page, blk, spec, right, key)
        elif 'bullets' in spec:
            self.render_bullets(page, blk, spec, right, key)
        else:
            self.render_lines(page, blk, spec, right, key)

    def render_fit(self, page, blk, spec, key):
        """Centre the text inside a fixed area, shrinking until it fits.

        Used for the characters' speech bubbles: the source text bbox is a good
        proxy for the usable area inside the balloon, so the translation is
        centred on it and scaled down until it stays inside.
        """
        si = spec.get('style', 0)
        styles, meta = {}, {}
        for idx in range(len(blk['styles'])):
            font, alias, ff, size, color = self.style(blk, idx)
            size = spec.get('size', size)
            styles[idx] = (font, size)
            meta[idx] = (alias, ff, color)
        box = fitz.Rect(spec.get('box', blk['bbox']))
        # balloons are elliptical: the usable width at the top and bottom lines
        # is smaller than the bounding box, so inset before fitting
        iw = box.width * spec.get('inset_w', 0.16) / 2
        ih = box.height * spec.get('inset_h', 0.08) / 2
        box.x0 += iw; box.x1 -= iw
        box.y0 += ih; box.y1 -= ih
        box.y0 += spec.get('dy', 0)
        box.y1 += spec.get('dy', 0)
        box.x0 += spec.get('dx', 0)
        box.x1 += spec.get('dx', 0)
        base = styles[si][1]
        lead0 = spec.get('leading', blk['leading'] or base * 1.2)
        words = parse_rich(spec['flow'], si)
        scale = 1.0
        while True:
            lines = wrap_rich(words, styles, box.width, scale)
            widest = max(line_width(l, styles, scale) for l in lines)
            height = (len(lines) - 1) * lead0 * scale + base * scale
            if (widest <= box.width + 0.5 and height <= box.height + 0.5) \
                    or scale <= FIT_MIN + 1e-9:
                break
            scale = round(scale - 0.02, 4)
        if widest > box.width + 1 or height > box.height + 1:
            self.warnings.append(f'{key}: не вписався у {box.width:.0f}x{box.height:.0f}')
        lead = lead0 * scale
        top = box.y0 + (box.height - ((len(lines) - 1) * lead + base * scale)) / 2
        for i, ln in enumerate(lines):
            lx = box.x0 + (box.width - line_width(ln, styles, scale)) / 2
            y = top + base * scale * 0.80 + i * lead
            for j, (word, st) in enumerate(ln):
                font, size = styles[st]
                alias, ff, color = meta[st]
                if j:
                    lx += measure(font, ' ', size * scale)
                self.draw(page, lx, y, word, alias, ff, size * scale, color)
                lx += measure(font, word, size * scale)

    def render_bullets(self, page, blk, spec, right, key):
        bs, ts = spec.get('bs', 0), spec.get('ts', 1)
        bfont, balias, bff, bsize, bcolor = self.style(blk, bs)
        tfont, talias, tff, tsize, tcolor = self.style(blk, ts)
        first = blk['lines'][0]['spans']
        bx = spec.get('bx', first[0]['x'])
        tx = spec.get('tx', first[1]['x'] if len(first) > 1 else bx + 12)
        y0 = spec.get('y0', first[0]['y'])
        lead = spec.get('lead', 14.4)
        gap = spec.get('gap', lead + 6.8)
        budget = len(blk['lines'])
        width = right - tx
        scale = 1.0
        while True:
            groups = [wrap_rich(parse_rich(t, ts), {ts: (tfont, tsize)}, width, scale)
                      for t in spec['bullets']]
            if sum(len(g) for g in groups) <= budget or scale <= MIN_SCALE + 1e-9:
                break
            scale = round(scale - 0.02, 4)
        if sum(len(g) for g in groups) > budget:
            self.warnings.append(f'{key}: bullets overflow')
        y = y0
        for gi, g in enumerate(groups):
            if gi:
                y += gap
            self.draw(page, bx, y, '\u2022', balias, bff, bsize * scale, bcolor)
            for li, ln in enumerate(g):
                x = tx
                for j, (word, st) in enumerate(ln):
                    if j:
                        x += measure(tfont, ' ', tsize * scale)
                    self.draw(page, x, y, word, talias, tff, tsize * scale, tcolor)
                    x += measure(tfont, word, tsize * scale)
                if li < len(g) - 1:
                    y += lead

    def render_flow(self, page, blk, spec, right, key):
        si = spec.get('style', 0)
        styles, meta = {}, {}
        for idx in range(len(blk['styles'])):
            font, alias, ff, size, color = self.style(blk, idx)
            size = spec.get('size', size)
            styles[idx] = (font, size)
            meta[idx] = (alias, ff, color)
        x0 = spec.get('x', blk['bbox'][0])
        width = spec.get('width', right - x0)
        base_size = styles[si][1]
        lead = spec.get('leading', blk['leading'] or base_size * 1.2)
        maxlines = len(blk['lines']) + spec.get('extra', 0)
        align = spec.get('align', 'left')
        words = parse_rich(spec['flow'], si)
        scale = 1.0
        while True:
            lines = wrap_rich(words, styles, width, scale)
            if len(lines) <= maxlines or scale <= MIN_SCALE + 1e-9:
                break
            scale = round(scale - 0.02, 4)
        if len(lines) > maxlines:
            self.warnings.append(f'{key}: {len(lines)} lines > {maxlines}')
        y = spec.get('y0', blk['lines'][0]['spans'][0]['y'])
        for i, ln in enumerate(lines):
            lx = x0
            if align == 'center':
                lx = x0 + (width - line_width(ln, styles, scale)) / 2
            elif align == 'right':
                lx = x0 + width - line_width(ln, styles, scale)
            for j, (word, st) in enumerate(ln):
                font, size = styles[st]
                alias, ff, color = meta[st]
                if j:
                    lx += measure(font, ' ', size * scale)
                self.draw(page, lx, y + i * lead, word, alias, ff, size * scale, color)
                lx += measure(font, word, size * scale)

    def render_lines(self, page, blk, spec, right, key):
        olines = blk['lines']
        lead = spec.get('leading', blk['leading'] or 0)
        base_y = olines[0]['spans'][0]['y']
        for i, entry in enumerate(spec['lines']):
            if entry is None:
                continue
            if isinstance(entry, dict):
                segs, opts = entry['segs'], entry
            else:
                segs, opts = entry, {}
            oline = olines[i] if i < len(olines) else None
            y = oline['spans'][0]['y'] if oline else base_y + lead * i
            y = opts.get('y', y)
            lright = opts.get('right', right)
            lsize = opts.get('size')

            centered = opts.get('center', spec.get('center'))

            def layout(scale, emit=False):
                pen = None
                for j, seg in enumerate(segs):
                    text, sidx = seg[0], seg[1]
                    font, alias, ff, size, color = self.style(blk, sidx)
                    size = (lsize or size) * scale
                    w = measure(font, text, size)
                    ax = seg[2] if len(seg) > 2 else self.anchor(oline, j)
                    if centered and oline is not None and j < len(oline['spans']):
                        osp = oline['spans'][j]
                        ax = osp['x'] + osp['w'] / 2 - w / 2
                    if pen is None:
                        pen = ax if ax is not None else (
                            oline['spans'][0]['x'] if oline else blk['bbox'][0])
                        pen = opts.get('x', pen) if j == 0 else pen
                    elif ax is not None:
                        pen = ax if centered else max(ax, pen + measure(font, ' ', size))
                    if emit:
                        save = self.cur_render
                        if len(seg) > 3:
                            self.cur_render = seg[3]
                        self.draw(page, pen, y, text, alias, ff, size, color)
                        self.cur_render = save
                    pen += w
                return pen if pen is not None else 0.0

            scale = 1.0
            while layout(scale) > lright and scale > MIN_SCALE:
                scale = round(scale - 0.02, 4)
            end = layout(scale, emit=True)
            if end > lright + 1:
                self.warnings.append(f'{key} L{i}: width {end:.0f} > {lright:.0f}')

    @staticmethod
    def anchor(oline, j):
        """Original x for segment j when the source had a real gap there."""
        if oline is None or j >= len(oline['spans']):
            return None
        sp = oline['spans'][j]
        if j == 0 or (sp['gap'] is not None and sp['gap'] > 3):
            return sp['x']
        return None

    # ------------------------------------------------------------------ run
    def run(self, only=None):
        for pg in self.tpl:
            pno = pg['page']
            if only is not None and pno not in only:
                continue
            page = self.doc[pno]
            todo = []
            for blk in pg['blocks']:
                key = f"p{pno}b{blk['i']}"
                if key in self.tr:
                    todo.append((blk, self.tr[key], key))
            # Закривати стару графіку треба ДО набору тексту, окремою фазою
            # (§2.1): інакше білий прямокутник не можна провести там, де ляже
            # новий текст, і від старих рамок лишаються хвостики.
            if self.covers:
                self.covers(page)
            if not todo:
                continue
            self.erase(page, [t[0] for t in todo])
            for blk, spec, key in todo:
                if spec.get('drop'):
                    continue
                self.render_block(page, blk, spec, key)
            self.cur_key = None
        return self.doc

    def save_layout(self, wipes=None):
        """layout.json — прямокутники кожного видрукуваного рядка.

        Це точні дані з моменту малювання, а не реконструкція з готового PDF;
        на них спирається геометричний контроль (§1.5).
        """
        path = os.path.join(self.lesson.work, 'layout.json')
        with open(path, 'w', encoding='utf-8') as fh:
            json.dump({k: v for k, v in self.drawn.items()}, fh)
        with open(os.path.join(self.lesson.work, 'wipes.json'), 'w') as fh:
            json.dump(wipes or [], fh)
        return path

    def save(self, path=None):
        path = path or self.lesson.out('print')
        os.makedirs(os.path.dirname(path), exist_ok=True)
        try:
            self.doc.subset_fonts(verbose=False)
        except Exception as exc:            # older PyMuPDF builds
            print('subset_fonts skipped:', exc)
        lv = self.lesson.level
        self.doc.set_metadata({
            'title': f'bibletime {self.lesson.series}{self.lesson.num} — {lv.subtitle}',
            'author': 'Bible Educational Services',
            'subject': f'bibletime — рівень {lv.n}, серія {self.lesson.series}'})
        self.doc.save(path, garbage=4, deflate=True)
        return path


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('stem', help='напр. Level4_A1')
    ap.add_argument('--pages', help='лише ці сторінки, з нуля: 0,3,7')
    ap.add_argument('--no-brand', action='store_true',
                    help='не застосовувати перебрендування')
    args = ap.parse_args()

    lesson = Lesson(args.stem)
    mod = __import__(f'tr_{lesson.stem}')
    only = {int(x) for x in args.pages.split(',')} if args.pages else None

    b = Builder(mod.TR, lesson, covers=getattr(mod, 'apply_covers', None))
    doc = b.run(only)
    for fn in getattr(mod, 'POST', []):
        fn(doc, only)
    b.save_layout(getattr(mod, 'WIPES', []))
    if not args.no_brand:
        from rebrand import rebrand
        st = rebrand(doc, lesson)
        print(f"  перебрендовано: напис на {st['напис']} стор., "
              f"логотип на {st['логотип']}")
    p = b.save()
    print('зібрано:', p)
    if b.warnings:
        print(f'--- зауважень {len(b.warnings)} ---')
        for w in b.warnings:
            print(' ', w)


if __name__ == '__main__':
    main()