#!/usr/bin/env python3
"""neusatz: setzt .docx- und PDF-Dateien mit Typst neu.

Typst bricht jeden Absatz als Ganzes nach Knuth-Plass um
(par.linebreaks: "optimized"), mit Silbentrennung in der Sprache des Textes.

  neusatz buch.docx                 erzeugt buch-neusatz.typ und buch-neusatz.pdf
  neusatz buch.pdf --seiten 20-40   übernimmt nur die Seiten 20 bis 40
  neusatz buch-neusatz.typ          setzt eine von Hand korrigierte .typ-Datei neu

Benötigt typst und pandoc, für PDF-Dateien zusätzlich pdftohtml (poppler-utils):
  sudo apt install pandoc poppler-utils
  Typst: Programmdatei von https://github.com/typst/typst/releases nach ~/.local/bin
"""

import argparse
import collections
import datetime
import json
import re
import shutil
import statistics
import subprocess
import sys
import zipfile
import xml.etree.ElementTree as ET
from pathlib import Path

W = '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}'
ZOOM = 1.5                 # Pixel je Punkt in der XML-Ausgabe von pdftohtml
TWIP_MM = 25.4 / 1440
PT_MM = 25.4 / 72
FALLBACK_FONT = 'Libertinus Serif'
WS = ' \t\r\n'             # bewusst ohne geschütztes Leerzeichen
SPACE = {'t': 'Space'}

LIST_STYLE = {'Decimal': '1', 'DefaultStyle': '1', 'Example': '1', 'UpperRoman': 'I',
              'LowerRoman': 'i', 'UpperAlpha': 'A', 'LowerAlpha': 'a'}


def info(msg):
    print(msg, flush=True)


def fail(msg):
    print(f'neusatz: {msg}', file=sys.stderr)
    sys.exit(1)


def run(cmd, cwd=None, stdin=None):
    try:
        r = subprocess.run(cmd, cwd=cwd, input=stdin, capture_output=True,
                           text=True, encoding='utf-8', errors='replace')
    except FileNotFoundError:
        fail(f'{cmd[0]} ist nicht installiert.')
    if r.returncode != 0:
        fail(f'{cmd[0]} ist fehlgeschlagen:\n{r.stderr.strip()}')
    return r.stdout


class Layout:
    """Satzspiegel und Schrift, so weit sie sich aus der Vorlage ablesen lassen."""

    def __init__(self):
        self.title = ''
        self.lang, self.region = 'de', 'DE'
        self.page = None           # (Breite, Höhe) in mm
        self.margins = None        # dict in mm
        self.mirror = False
        self.font = None
        self.size = None           # pt
        self.indent = False        # Erstzeileneinzug statt Abstand zwischen Absätzen
        self.chapter_break = False
        self.numbering = {}        # Überschriftenebene -> Pandoc-Listenstil


# ---------------------------------------------------------------- Pandoc-AST

def plain_text(x):
    if isinstance(x, dict):
        t = x.get('t')
        if t == 'Str':
            return x['c']
        if t in ('Space', 'SoftBreak', 'LineBreak'):
            return ' '
        if t == 'Note':
            return ''
        return plain_text(x.get('c', []))
    if isinstance(x, list):
        return ''.join(plain_text(v) for v in x)
    return ''


def header(level, inlines, ident=''):
    return {'t': 'Header', 'c': [level, [ident, [], []], inlines]}


def raw_block(code):
    return {'t': 'RawBlock', 'c': ['typst', code]}


def typst_escape(s, at_start=False):
    """Maskiert Text für Typst-Markup.

    Pandoc 3.1 lässt einiges durch (~, //, -- und Listenzeichen am
    Absatzanfang), darum geht aller Text als fertig maskiertes Typst hinaus.
    """
    out = []
    for i, ch in enumerate(s):
        prev, nxt = s[i - 1:i], s[i + 1:i + 2]
        if (ch in '\\#$*_`<>@[]~'
                or (ch == '/' and ('/' in (prev, nxt) or '*' in (prev, nxt)))
                or (ch == '-' and ('-' in (prev, nxt) or nxt == '?'))
                or (at_start and i == 0 and ch in '-+=/')):
            out.append('\\')
        out.append(ch)
    text = ''.join(out)
    if at_start:
        text = re.sub(r'^(\d+)\.', r'\1\\.', text)
    return text


# Inlines, die Pandoc als #funktion[...] schreibt. Folgt direkt "(", "[" oder
# ".wort", liest Typst das als Fortsetzung des Aufrufs.
CALL_INLINES = {'Emph', 'Strong', 'SmallCaps', 'Superscript', 'Subscript', 'Underline',
                'Strikeout', 'Note', 'Link', 'Span', 'Image', 'Cite'}


def escape_ast(x, keep_shy):
    """Maskiert allen Text; weiche Trennstriche der Vorlage fallen weg, wenn
    keep_shy falsch ist (sie stammen meist aus dem alten Umbruch)."""
    if isinstance(x, list):
        out, prev = [], None
        for v in x:
            e = escape_ast(v, keep_shy)
            if prev in CALL_INLINES and e.get('t') == 'RawInline' and re.match(r'[(\[]|\.\w', e['c'][1]):
                e['c'][1] = '\\' + e['c'][1]
            out.append(e)
            prev = v.get('t') if isinstance(v, dict) else None
        return out
    if isinstance(x, dict) and 't' in x:
        t = x['t']
        if t == 'Str':
            text = x['c'] if keep_shy else x['c'].replace('­', '')
            return {'t': 'RawInline', 'c': ['typst', typst_escape(text)]}
        if 'c' not in x:
            return x
        c = escape_ast(x['c'], keep_shy)
        if t in ('Para', 'Plain') and x['c'] and x['c'][0].get('t') == 'Str':
            text = x['c'][0]['c'] if keep_shy else x['c'][0]['c'].replace('­', '')
            c[0]['c'][1] = typst_escape(text, at_start=True)
        return {'t': t, 'c': c}
    return x


def pandoc_api_version():
    return json.loads(run(['pandoc', '-f', 'markdown', '-t', 'json'], stdin=''))['pandoc-api-version']


def guess_lang(text):
    words = collections.Counter(re.findall(r'[a-zäöüß]+', text.lower()))
    de = sum(words[w] for w in ('der', 'die', 'das', 'und', 'nicht', 'ist', 'ich', 'sie'))
    en = sum(words[w] for w in ('the', 'and', 'of', 'to', 'is', 'that', 'you', 'it'))
    return ('de', 'DE') if de >= en else ('en', None)


# ---------------------------------------------------------------- DOCX

class DocxStyles:
    def __init__(self, path):
        with zipfile.ZipFile(path) as z:
            names = set(z.namelist())

            def load(name):
                return ET.fromstring(z.read(name)) if name in names else None

            styles, self.doc, settings = load('word/styles.xml'), load('word/document.xml'), load('word/settings.xml')
            core = load('docProps/core.xml')
        if self.doc is None:
            fail(f'{path.name} enthält kein word/document.xml.')
        title = core.find('{http://purl.org/dc/elements/1.1/}title') if core is not None else None
        self.title = (title.text or '').strip() if title is not None else ''

        self.by_id, self.by_name, self.default_id = {}, {}, None
        self.default_font = self.default_size = self.default_lang = None
        if styles is not None:
            rpr = styles.find(f'{W}docDefaults/{W}rPrDefault/{W}rPr')
            if rpr is not None:
                self.default_font, self.default_size = self._font(rpr)
                lang = rpr.find(f'{W}lang')
                self.default_lang = lang.get(f'{W}val') if lang is not None else None
            for st in styles.findall(f'{W}style'):
                if st.get(f'{W}type') != 'paragraph':
                    continue
                sid = st.get(f'{W}styleId')
                ppr = st.find(f'{W}pPr')
                ind = ppr.find(f'{W}ind') if ppr is not None else None
                brk = ppr.find(f'{W}pageBreakBefore') if ppr is not None else None
                font, size = self._font(st.find(f'{W}rPr'))
                s = {
                    'name': self._val(st, f'{W}name') or sid,
                    'based': self._val(st, f'{W}basedOn'),
                    'outline': self._val(st, f'{W}pPr/{W}outlineLvl'),
                    'break': brk is not None and brk.get(f'{W}val') not in ('0', 'false'),
                    'first_line': ind.get(f'{W}firstLine') if ind is not None else None,
                    'font': font, 'size': size,
                }
                self.by_id[sid] = s
                self.by_name[s['name']] = sid
                if st.get(f'{W}default') in ('1', 'true'):
                    self.default_id = sid
        self.mirror = settings is not None and any(
            e.get(f'{W}val') not in ('0', 'false') for e in settings.iter(f'{W}mirrorMargins'))

    @staticmethod
    def _val(el, path):
        e = el.find(path)
        return e.get(f'{W}val') if e is not None else None

    @staticmethod
    def _font(rpr):
        if rpr is None:
            return None, None
        fonts, sz = rpr.find(f'{W}rFonts'), rpr.find(f'{W}sz')
        font = fonts.get(f'{W}hAnsi') or fonts.get(f'{W}ascii') if fonts is not None else None
        size = int(sz.get(f'{W}val')) / 2 if sz is not None and sz.get(f'{W}val', '').isdigit() else None
        return font, size

    def resolve(self, sid, key):
        for _ in range(20):
            s = self.by_id.get(sid)
            if s is None:
                return None
            if s[key] not in (None, False):
                return s[key]
            sid = s['based']
        return None

    def heading_level(self, name):
        m = re.fullmatch(r'(?:heading|überschrift)\s*(\d)', name, re.I)
        if m:
            return int(m.group(1))
        lvl = self.resolve(self.by_name.get(name), 'outline')
        return int(lvl) + 1 if lvl and lvl.isdigit() and int(lvl) < 9 else None

    @staticmethod
    def is_toc(name):
        return bool(re.search(r'inhaltsverz|^(toc|verzeichnis|inhaltsverzeichnis|contents)\s*\d', name, re.I))

    @staticmethod
    def is_quote(name):
        return bool(re.search(r'zitat|quot|citation', name, re.I))

    def body_style(self):
        counts = collections.Counter(e.get(f'{W}val') for e in self.doc.iter(f'{W}pStyle'))
        unstyled = sum(1 for p in self.doc.iter(f'{W}p') if p.find(f'{W}pPr/{W}pStyle') is None)
        if self.default_id:
            counts[self.default_id] += unstyled
        return counts.most_common(1)[0][0] if counts else self.default_id

    def page_breaks(self):
        return sum(1 for br in self.doc.iter(f'{W}br') if br.get(f'{W}type') == 'page')

    def apply_layout(self, lay):
        sects = list(self.doc.iter(f'{W}sectPr'))
        if sects:
            sz, mar = sects[-1].find(f'{W}pgSz'), sects[-1].find(f'{W}pgMar')
            if sz is not None and sz.get(f'{W}w'):
                lay.page = (int(sz.get(f'{W}w')) * TWIP_MM, int(sz.get(f'{W}h')) * TWIP_MM)
            if mar is not None:
                get = lambda k: abs(int(mar.get(f'{W}{k}', '1418'))) * TWIP_MM
                lay.margins = {'left': get('left') + int(mar.get(f'{W}gutter', '0')) * TWIP_MM,
                               'right': get('right'), 'top': get('top'), 'bottom': get('bottom')}
        lay.mirror = self.mirror
        body = self.body_style()
        lay.font = self.resolve(body, 'font') or self.default_font
        lay.size = self.resolve(body, 'size') or self.default_size
        first = self.resolve(body, 'first_line')
        lay.indent = bool(first and first.lstrip('-').isdigit() and int(first) > 0)
        if self.default_lang:
            parts = self.default_lang.split('-')
            lay.lang, lay.region = parts[0].lower(), (parts[1].upper() if len(parts) > 1 else None)


class DocxConverter:
    """Bereinigt Pandocs docx+styles-AST: Formatvorlagen werden zu Überschriften,
    Zitaten und einem Typst-Inhaltsverzeichnis."""

    WRAPPERS = ('Emph', 'Strong', 'Underline', 'Strikeout', 'Superscript', 'Subscript', 'SmallCaps')

    def __init__(self, styles, lay):
        self.st = styles
        self.toc_done = False
        self.numbering = {}
        self.page = self.box = None        # Seite und Satzspiegel in Zoll
        if lay.page and lay.margins:
            m = lay.margins
            self.page = (lay.page[0] / 25.4, lay.page[1] / 25.4)
            self.box = ((lay.page[0] - m['left'] - m['right']) / 25.4 * 0.98,
                        (lay.page[1] - m['top'] - max(m['bottom'], 18)) / 25.4 * 0.9)

    def full_page_image(self, inls):
        """Ein Absatz, der nur aus einem seitengroßen Bild besteht (etwa ein Umschlag)."""
        images = [i for i in inls if i['t'] == 'Image']
        if self.page is None or len(images) != 1 or any(
                i['t'] not in ('Image', 'Space', 'SoftBreak', 'LineBreak') for i in inls):
            return None
        w, h = image_size(images[0])
        if w and h and w >= 0.9 * self.page[0] and h >= 0.9 * self.page[1]:
            path = typst_str(images[0]['c'][2][0])
            return raw_block(f'#page(margin: 0pt, numbering: none, '
                             f'image({path}, width: 100%, height: 100%, fit: "cover"))')
        return None

    def fit_image(self, img):
        w, h = image_size(img)
        if not (w and h and self.box):
            return img
        f = min(1, self.box[0] / w, self.box[1] / h)
        if f < 1:
            attr = img['c'][0]
            kv = [kv for kv in attr[2] if kv[0] not in ('width', 'height')]
            kv += [['width', f'{w * f:.3f}in'], ['height', f'{h * f:.3f}in']]
            img = {'t': 'Image', 'c': [[attr[0], attr[1], kv], img['c'][1], img['c'][2]]}
        return img

    def blocks(self, blocks):
        out = []
        for b in blocks:
            style = dict(b['c'][0][2]).get('custom-style', '') if b['t'] == 'Div' else ''
            if style and self.st.is_toc(style):
                # Das alte Inhaltsverzeichnis hat Seitenzahlen der alten Fassung.
                if not self.toc_done:
                    if out and out[-1]['t'] in ('Para', 'Plain'):
                        title = plain_text(out[-1]).strip()
                        if len(title) < 40 and re.search(r'inhalt|contents', title, re.I):
                            out.pop()
                    out.append(raw_block('#outline()'))
                    self.toc_done = True
                continue
            out.extend(self.block(b))
        return out

    def block(self, b):
        t, c = b['t'], b.get('c')
        if t in ('Para', 'Plain'):
            cover = self.full_page_image(c)
            if cover:
                return [cover]
            inl = self.inlines(c)
            return [{'t': t, 'c': inl}] if plain_text(inl).strip() or any(
                i['t'] in ('Image', 'Note') for i in inl) else []
        if t == 'Header':
            return [{'t': t, 'c': [c[0], c[1], self.inlines(c[2])]}]
        if t == 'BlockQuote':
            return [{'t': t, 'c': self.blocks(c)}]
        if t == 'Div':
            style = dict(c[0][2]).get('custom-style', '')
            inner = self.blocks(c[1])
            if style:
                lvl = self.st.heading_level(style)
                if lvl and len(inner) == 1 and inner[0]['t'] in ('Para', 'Plain'):
                    return [header(lvl, inner[0]['c'], 'ohne-nummer')]
                if self.st.is_quote(style) and inner and not (len(inner) == 1 and inner[0]['t'] == 'BlockQuote'):
                    return [{'t': 'BlockQuote', 'c': inner}]
            return inner
        if t in ('OrderedList', 'BulletList'):
            items = [self.blocks(i) for i in (c[1] if t == 'OrderedList' else c)]
            if items and all(len(i) == 1 and i[0]['t'] == 'Header' for i in items):
                # Word nummeriert Überschriften gern über Listen.
                heads = [i[0] for i in items]
                if t == 'OrderedList':
                    for h in heads:
                        self.numbering.setdefault(h['c'][0], c[0][1]['t'])
                        h['c'][1] = ['', [], []]
                return heads
            return [{'t': t, 'c': [c[0], items] if t == 'OrderedList' else items}]
        if t == 'DefinitionList':
            return [{'t': t, 'c': [[self.inlines(term), [self.blocks(d) for d in defs]] for term, defs in c]}]
        if t == 'LineBlock':
            return [{'t': t, 'c': [self.inlines(line) for line in c]}]
        return [b]

    def inlines(self, inls):
        out = []
        for i in inls:
            t, c = i['t'], i.get('c')
            if t == 'Span':
                out.extend(self.inlines(c[1]))      # Zeichenvorlagen und Textmarken
            elif t == 'Note':
                out.append({'t': t, 'c': self.blocks(c)})
            elif t in self.WRAPPERS:
                out.append({'t': t, 'c': self.inlines(c)})
            elif t == 'Quoted':
                out.append({'t': t, 'c': [c[0], self.inlines(c[1])]})
            elif t == 'Link':
                if c[2][0].startswith('#'):
                    out.extend(self.inlines(c[1]))  # Sprungziele gibt es im Neusatz nicht
                else:
                    out.append({'t': t, 'c': [c[0], self.inlines(c[1]), c[2]]})
            elif t == 'Image':
                out.append(self.fit_image(i))
            else:
                out.append(i)
        return out


def image_size(img):
    """Breite und Höhe eines Pandoc-Bildes in Zoll, sofern angegeben."""
    unit = {'in': 1, 'cm': 1 / 2.54, 'mm': 1 / 25.4, 'pt': 1 / 72, 'px': 1 / 96, '': 1 / 96}
    kv = dict(img['c'][0][2])
    out = []
    for key in ('width', 'height'):
        m = re.fullmatch(r'([\d.]+)\s*(in|cm|mm|pt|px|)', kv.get(key, ''))
        out.append(float(m.group(1)) * unit[m.group(2)] if m else None)
    return out


def docx_to_ast(src, workdir, media, lay):
    st = DocxStyles(src)
    st.apply_layout(lay)
    doc = json.loads(run(['pandoc', '-f', 'docx+styles', f'--extract-media={media}',
                          '-t', 'json', str(src.resolve())], cwd=workdir))
    conv = DocxConverter(st, lay)
    doc['blocks'] = conv.blocks(doc['blocks'])
    lay.numbering = conv.numbering

    heads = [b for b in doc['blocks'] if b['t'] == 'Header']
    if not lay.numbering:
        for h in heads:
            h['c'][1][0] = ''
    h1 = [h for h in heads if h['c'][0] == 1]
    h1_break = any(st.resolve(sid, 'break') for name, sid in st.by_name.items()
                   if st.heading_level(name) == 1)
    lay.chapter_break = len(h1) >= 2 and (h1_break or st.page_breaks() >= len(h1) / 2)
    lay.title = st.title
    if not st.default_lang:
        lay.lang, lay.region = guess_lang(plain_text(doc['blocks']))
    info(f'Word-Datei gelesen: {len(heads)} Überschriften, {count_notes(doc["blocks"])} Fußnoten.')
    return doc


def count_notes(x):
    if isinstance(x, list):
        return sum(count_notes(v) for v in x)
    if isinstance(x, dict):
        return (x.get('t') == 'Note') + count_notes(x.get('c', []))
    return 0


# ---------------------------------------------------------------- PDF

class Frag:
    """Ein Textstück aus pdftohtml, mit Lage auf der Seite."""
    __slots__ = ('top', 'left', 'w', 'h', 'size', 'family', 'segs', 'text')


class Run:
    __slots__ = ('text', 'it', 'bf', 'kind', 'note')

    def __init__(self, text='', it=False, bf=False, kind='', note=None):
        self.text, self.it, self.bf, self.kind, self.note = text, it, bf, kind, note

    def copy(self):
        return Run(self.text, self.it, self.bf, self.kind, self.note)


class Line:
    def __init__(self, frag):
        self.frags = [frag]
        self.top, self.h = frag.top, frag.h

    def add(self, frag):
        self.frags.append(frag)
        if frag.h > self.h:
            self.top, self.h = frag.top, frag.h

    def finish(self, page):
        self.page = page
        self.frags.sort(key=lambda f: f.left)
        weight = collections.Counter()
        for f in self.frags:
            weight[f.size] += len(f.text.strip(WS))
        self.size = weight.most_common(1)[0][0]
        self.base = max(f.top + f.h for f in self.frags if f.size == self.size)
        self.left = min(f.left for f in self.frags)
        self.right = max(f.left + f.w for f in self.frags)
        self.runs = self._runs()
        self.text = ''.join(r.text for r in self.runs)

    def _runs(self):
        runs, prev = [], None
        for f in self.frags:
            kind = ''
            if f.size <= 0.8 * self.size:
                if f.top + f.h < self.base - 0.2 * self.h:
                    kind = 'sup'
                elif re.fullmatch(r"[A-ZÄÖÜ\s.,'’-]+", f.text) and re.search('[A-ZÄÖÜ]', f.text):
                    kind = 'sc'
            if prev is not None and kind != 'sup' and runs:
                gap = f.left - (prev.left + prev.w)
                if gap > 0.15 * self.size and runs[-1].text[-1:] not in WS and f.text[:1] not in WS:
                    runs.append(Run(' '))
            runs.extend(Run(text, it, bf, kind) for text, it, bf in f.segs if text)
            prev = f
        # Unechte Kapitälchen: großer Anfangsbuchstabe, Rest kleiner gesetzte Versalien
        for i, r in enumerate(runs):
            if r.kind != 'sc':
                continue
            r.text = r.text.lower()
            if i and runs[i - 1].kind == '':
                m = re.search(r'(?:^|[\s(„»«"‚\'])([A-ZÄÖÜ])$', runs[i - 1].text)
                if m:
                    runs[i - 1].text = runs[i - 1].text[:-1]
                    r.text = m.group(1) + r.text
        return [r for r in runs if r.text]


def frag_segments(el):
    segs = []

    def walk(node, it, bf):
        it = it or node.tag == 'i'
        bf = bf or node.tag == 'b'
        if node.text:
            segs.append((node.text, it, bf))
        for child in node:
            walk(child, it, bf)
            if child.tail:
                segs.append((child.tail, it, bf))

    walk(el, False, False)
    return segs


def read_pdf_pages(src, pages):
    cmd = ['pdftohtml', '-xml', '-i', '-q', '-stdout', '-zoom', str(ZOOM)]
    if pages:
        cmd += ['-f', str(pages[0]), '-l', str(pages[1])]
    xml = run(cmd + [str(src)])
    xml = re.sub(r'<!DOCTYPE[^>]*>', '', xml)
    xml = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f]', '', xml)
    try:
        root = ET.fromstring(xml)
    except ET.ParseError as e:
        fail(f'Die Ausgabe von pdftohtml ließ sich nicht lesen ({e}).')

    fonts = {fs.get('id'): (float(fs.get('size')), fs.get('family', ''))
             for fs in root.iter('fontspec')}
    result = []
    for page in root.findall('page'):
        frags = []
        for el in page.findall('text'):
            f = Frag()
            f.top, f.left = float(el.get('top')), float(el.get('left'))
            f.w, f.h = float(el.get('width')), float(el.get('height'))
            f.size, f.family = fonts.get(el.get('font'), (f.h, ''))
            f.segs = frag_segments(el)
            f.text = ''.join(s[0] for s in f.segs)
            if not f.text.strip(WS) or f.h <= 0:
                continue
            if any(g.text == f.text and abs(g.top - f.top) <= 0.15 * f.h
                   and abs(g.left - f.left) <= 0.15 * f.h for g in frags[-8:]):
                continue                    # Schatten- oder Doppeldruck
            frags.append(f)
        lines = []
        for f in sorted(frags, key=lambda f: (f.top, f.left)):
            for ln in reversed(lines[-3:]):
                if min(f.top + f.h, ln.top + ln.h) - max(f.top, ln.top) >= 0.5 * min(f.h, ln.h):
                    ln.add(f)
                    break
            else:
                lines.append(Line(f))
        number = int(page.get('number'))
        for ln in lines:
            ln.finish(number)
        lines.sort(key=lambda ln: ln.top)
        result.append({'number': number, 'w': float(page.get('width')),
                       'h': float(page.get('height')), 'lines': lines})
    return result


def norm_running(text):
    return re.sub(r'\d+', '#', re.sub(r'\s+', ' ', text)).strip().lower()


PAGENUM_RE = re.compile(r'[-–—\s]*((seite|page|s\.)\s*)?(\d+|[ivxlcdm]+)(\s*(von|of|/)\s*\d+)?[-–—\s]*', re.I)
ENDS_RE = re.compile(r'[.!?:;»«“”"\')\]]\s*$')
BULLET_RE = re.compile(r'\s*[•◦▪▫‣●○■□]\s')
ENUM_RE = re.compile(r'\s*\(?(\d{1,3}|[a-zA-Z])[.)]\s')
TOC_TITLE_RE = re.compile(r'\s*(inhalt|inhaltsverzeichnis|contents|table of contents)\s*', re.I)
CONJ = {'und', 'oder', 'bis', 'sowie', 'als', 'bzw', 'beziehungsweise', 'respektive', 'wie',
        'noch', 'statt', 'anstatt', 'zu', 'and', 'or', 'nor', 'to'}


def remove_running_heads(pages):
    def in_band(ln, page):
        return ln.top < 0.12 * page['h'] or ln.base > 0.88 * page['h']

    def edges(page):
        return set(page['lines'][:2] + page['lines'][-2:])

    keys = collections.Counter()
    for page in pages:
        for ln in edges(page):
            if in_band(ln, page):
                keys[norm_running(ln.text)] += 1
    limit = max(2, 0.2 * len(pages))
    for page in pages:
        drop = {ln for ln in edges(page) if in_band(ln, page) and (
            keys[norm_running(ln.text)] >= limit or PAGENUM_RE.fullmatch(ln.text))}
        page['lines'] = [ln for ln in page['lines'] if ln not in drop]


def join_lines(run_lists):
    """Fügt Zeilen zu einem Absatz zusammen und nimmt Trennstriche heraus."""
    out = []
    for runs in run_lists:
        runs = [r.copy() for r in runs]
        if not runs:
            continue
        if out:
            if out[-1].note is not None or runs[0].note is not None:
                out.append(Run(' '))
            else:
                prev = out[-1].text.rstrip(WS)
                nxt = runs[0].text.lstrip(WS)
                if prev[-1:] in ('-', '­', '‐') and len(prev) > 1 and prev[-2].isalpha() and nxt[:1].isalpha():
                    word = re.match(r'\S+', nxt).group(0).rstrip('.,;:!?').lower()
                    if nxt[0].islower() and word not in CONJ:
                        out[-1].text = prev[:-1]
                    elif nxt[0].islower():
                        out[-1].text = prev + ' '
                    else:
                        out[-1].text = prev[:-1] if prev[-1] == '­' else prev
                elif prev.endswith('-') and prev[-2:-1].isdigit() and nxt[:1].isdigit():
                    out[-1].text = prev
                else:
                    out[-1].text = prev + ' '
                runs[0].text = nxt
        out.extend(runs)
    return out


def split_ws(text):
    core = text.strip(WS)
    if not core:
        return bool(text), '', False
    return text[0] in WS, core, text[-1] in WS


def strs(text):
    out = []
    for i, part in enumerate(re.split(r'[ \t\r\n]+', text)):
        if i:
            out.append(SPACE)
        if part:
            out.append({'t': 'Str', 'c': part})
    return out


def to_inlines(runs):
    merged = []
    for r in runs:
        last = merged[-1] if merged else None
        if (last is not None and r.note is None and last.note is None
                and (r.it, r.bf, r.kind) == (last.it, last.bf, last.kind)):
            last.text += r.text
        else:
            merged.append(r.copy())
    out = []
    for r in merged:
        if r.note is not None:
            out.append({'t': 'Note', 'c': r.note.blocks()})
            continue
        lead, core, trail = split_ws(r.text)
        if lead:
            out.append(SPACE)
        if core:
            inl = strs(core)
            if r.kind == 'sup':
                inl = [{'t': 'Superscript', 'c': inl}]
            elif r.kind == 'sc':
                inl = [{'t': 'SmallCaps', 'c': inl}]
            if r.it:
                inl = [{'t': 'Emph', 'c': inl}]
            if r.bf:
                inl = [{'t': 'Strong', 'c': inl}]
            out.extend(inl)
        if trail:
            out.append(SPACE)
    clean = []
    for i in out:
        if i['t'] == 'Space' and (not clean or clean[-1]['t'] == 'Space'):
            continue
        clean.append(i)
    while clean and clean[-1]['t'] == 'Space':
        clean.pop()
    return clean


class Footnote:
    def __init__(self, num):
        self.num, self.parts, self.used = num, [], False

    def blocks(self):
        return [{'t': 'Para', 'c': to_inlines(join_lines(self.parts))}]


def footnote_number(ln):
    first = ln.runs[0] if ln.runs else None
    if first is None:
        return None, []
    runs = [r.copy() for r in ln.runs]
    if first.kind == 'sup' and first.text.strip(WS).isdigit():
        return int(first.text.strip(WS)), runs[1:]
    m = re.match(r'[ \t]*(\d{1,3})(?=[ \t]|[^\d.,])[ \t]*', first.text)
    if m:
        runs[0].text = first.text[m.end():]
        return int(m.group(1)), runs
    return None, runs


class PdfReader:
    def __init__(self, pages):
        self.pages = pages
        lines = [ln for p in pages for ln in p['lines']]
        if not lines:
            fail('Im PDF steht kein Text. Ist es ein Scan? Dann zuerst Texterkennung, '
                 'zum Beispiel mit ocrmypdf.')
        weight = collections.Counter()
        for ln in lines:
            weight[ln.size] += len(ln.text)
        self.body = weight.most_common(1)[0][0]
        body_lines = [ln for ln in lines if abs(ln.size - self.body) < 0.5]
        self.body_h = statistics.median(ln.h for ln in body_lines)
        lefts = sorted(ln.left for ln in body_lines)
        rights = sorted(ln.right for ln in body_lines)
        self.col_left = lefts[len(lefts) // 10]
        self.col_right = rights[min(len(rights) - 1, len(rights) * 9 // 10)]
        self.width = max(1.0, self.col_right - self.col_left)
        diffs = [b.top - a.top for p in pages for a, b in zip(p['lines'], p['lines'][1:])
                 if abs(a.size - self.body) < 0.5 and abs(b.size - self.body) < 0.5
                 and 0 < b.top - a.top < 2 * a.h]
        self.pitch = statistics.median(diffs) if diffs else 1.2 * self.body_h
        self.heading_size = 1.12 * self.body
        self.orphans = 0

    # Fußnoten stehen klein am Seitenende und werden an ihre Ziffer im Text gehängt.
    def take_footnotes(self):
        notes, last = {}, None
        for page in self.pages:
            lines = page['lines']
            k = len(lines)
            while k > 0 and lines[k - 1].size < 0.92 * self.body:
                k -= 1
            zone = lines[k:]
            if not zone or k == 0:
                last = None
                continue
            numbered = [footnote_number(ln)[0] for ln in zone]
            markers = {int(r.text.strip(WS)) for ln in lines[:k] for r in ln.runs
                       if r.kind == 'sup' and r.text.strip(WS).isdigit()}
            wide_gap = zone[0].top - lines[k - 1].top > 2.2 * self.pitch
            if not (wide_gap or markers & {n for n in numbered if n is not None}):
                last = None
                continue
            entries = []
            min_left = min(ln.left for ln in zone)
            for ln in zone:
                num, runs = footnote_number(ln)
                if num is not None and ln.left <= min_left + 0.3 * ln.size:
                    entries.append(Footnote(num))
                    entries[-1].parts.append(runs)
                elif entries:
                    entries[-1].parts.append(ln.runs)
                elif last is not None:
                    last.parts.append(ln.runs)     # Fortsetzung von der Vorseite
                else:
                    entries.append(Footnote(None))
                    entries[-1].parts.append(ln.runs)
            notes[page['number']] = entries
            last = entries[-1] if entries else last
            page['lines'] = lines[:k]

        for page in self.pages:
            n = page['number']
            for ln in page['lines']:
                for i, r in enumerate(ln.runs):
                    if r.kind != 'sup' or not r.text.strip(WS).isdigit():
                        continue
                    num = int(r.text.strip(WS))
                    for cand in (notes.get(n, []), notes.get(n + 1, []), notes.get(n - 1, [])):
                        hit = next((e for e in cand if e.num == num and not e.used), None)
                        if hit:
                            hit.used = True
                            ln.runs[i] = Run(note=hit)
                            break
        for page in self.pages:
            rest = [e for e in notes.get(page['number'], []) if not e.used]
            if rest and page['lines']:
                page['lines'][-1].runs.extend(Run(note=e) for e in rest)
                self.orphans += len(rest)
        return sum(len(v) for v in notes.values())

    def short(self, ln):
        return ln.right < self.col_left + 0.8 * self.width

    def new_paragraph(self, para, prev, ln):
        is_h = ln.size >= self.heading_size
        if (prev.size >= self.heading_size) != is_h or abs(ln.size - prev.size) > 0.08 * prev.size:
            return True
        pitch = self.pitch * prev.h / self.body_h
        if is_h:
            return ln.page != prev.page or ln.top - prev.top > 1.6 * pitch
        if BULLET_RE.match(ln.text):
            return True
        if ENUM_RE.match(ln.text) and (self.short(prev) or ENDS_RE.search(prev.text)):
            return True
        if ln.page != prev.page:
            if self.short(prev):
                return True
            return bool(ENDS_RE.search(prev.text)) and prev.right < self.col_left + 0.93 * self.width
        gap = ln.top - prev.top
        if gap <= 0 or gap > 1.4 * pitch:
            return True
        if self.short(prev) and abs(ln.left - prev.left) > 0.5 * prev.size:
            return True
        first = para['lines'][0]
        return (ln.left - prev.left > 0.8 * prev.size and ln.left - self.col_left < 4 * prev.size
                and abs(prev.left - self.col_left) < 0.3 * prev.size
                and not BULLET_RE.match(first.text) and not ENUM_RE.match(first.text))

    def paragraphs(self):
        paras = []
        for page in self.pages:
            for ln in page['lines']:
                if paras and not self.new_paragraph(paras[-1], paras[-1]['lines'][-1], ln):
                    paras[-1]['lines'].append(ln)
                else:
                    paras.append({'lines': [ln], 'heading': ln.size >= self.heading_size})
        return paras

    def to_blocks(self, paras, lay):
        counts = collections.Counter(round(p['lines'][0].size) for p in paras if p['heading'])
        if sum(counts.values()) >= 3:
            for p in paras:
                if p['heading'] and counts[round(p['lines'][0].size)] == 1:
                    p['heading'] = False    # Einzelstück wie der Buchtitel, keine Gliederungsebene
        sizes = sorted({round(p['lines'][0].size) for p in paras if p['heading']}, reverse=True)
        level = {s: min(i + 1, 6) for i, s in enumerate(sizes)}
        blocks, toc, h1_tops = [], False, []
        page_h = {p['number']: p['h'] for p in self.pages}
        for p in paras:
            lines = p['lines']
            text = ' '.join(ln.text for ln in lines)
            if toc:
                ends = sum(bool(re.search(r'\d+\s*$', ln.text)) for ln in lines)
                if ends >= len(lines) / 2:
                    continue
                toc = False
            if p['heading'] and TOC_TITLE_RE.fullmatch(text):
                blocks.append(raw_block('#outline()'))
                toc = True
                continue
            runs = join_lines([ln.runs for ln in lines])
            if p['heading'] and all(r.bf for r in runs if r.text.strip(WS)):
                for r in runs:
                    r.bf = False            # Überschriften sind ohnehin fett
            inl = to_inlines(runs)
            if not inl:
                continue
            if p['heading']:
                lvl = level[round(lines[0].size)]
                if lvl == 1:
                    h1_tops.append(lines[0].top / page_h[lines[0].page])
                blocks.append(header(lvl, inl))
                continue
            para = {'t': 'Para', 'c': inl}
            if (len(lines) >= 2 and min(ln.left for ln in lines) > self.col_left + 0.04 * self.width
                    and not BULLET_RE.match(lines[0].text) and not ENUM_RE.match(lines[0].text)):
                para = {'t': 'BlockQuote', 'c': [para]}
            blocks.append(para)
        lay.chapter_break = len(h1_tops) >= 2 and sum(t < 0.35 for t in h1_tops) >= 0.7 * len(h1_tops)
        return blocks

    def apply_layout(self, paras, lay):
        mm = lambda px: px / ZOOM * PT_MM
        lay.page = (mm(statistics.median(p['w'] for p in self.pages)),
                    mm(statistics.median(p['h'] for p in self.pages)))
        tops = [p['lines'][0].top for p in self.pages if p['lines']]
        # Volle Seiten zählen, nicht Kapitelenden
        bottoms = sorted(p['h'] - p['bottom'] for p in self.pages if 'bottom' in p)
        page_w = statistics.median(p['w'] for p in self.pages)
        lay.margins = {'left': mm(self.col_left), 'right': mm(page_w - self.col_right),
                       'top': mm(statistics.median(tops)), 'bottom': mm(bottoms[len(bottoms) // 4])}
        lay.size = round(self.body / ZOOM * 2) / 2
        families = collections.Counter()
        for page in self.pages:
            for ln in page['lines']:
                for f in ln.frags:
                    if abs(f.size - self.body) < 0.5:
                        families[f.family] += len(f.text)
        if families:
            name = re.sub(r'^[A-Z]{6}\+', '', families.most_common(1)[0][0])
            name = re.split(r'[-,]', name)[0]
            name = re.sub(r'(PS)?MT$', '', name)
            lay.font = re.sub(r'(?<=[a-z])(?=[A-Z])', ' ', name)
        starts = [p for p in paras if not p['heading'] and len(p['lines']) >= 2]
        indented = sum(1 for p in starts
                       if p['lines'][0].left - p['lines'][1].left > 0.8 * p['lines'][0].size)
        lay.indent = bool(starts) and indented > len(starts) / 2


def pdf_to_ast(src, pages, lay):
    reader = PdfReader(read_pdf_pages(src, pages))
    remove_running_heads(reader.pages)
    for page in reader.pages:
        if page['lines']:
            page['bottom'] = max(ln.base for ln in page['lines'])
    notes = reader.take_footnotes()
    paras = reader.paragraphs()
    blocks = reader.to_blocks(paras, lay)
    reader.apply_layout(paras, lay)
    lay.lang, lay.region = guess_lang(' '.join(ln.text for p in reader.pages for ln in p['lines']))
    heads = sum(1 for b in blocks if b['t'] == 'Header')
    info(f'PDF gelesen: {len(reader.pages)} Seiten, {len(blocks)} Absätze, '
         f'{heads} Überschriften, {notes} Fußnoten.')
    if shutil.which('pdfinfo'):
        m = re.search(r'^Title:\s*(.+)$', run(['pdfinfo', str(src)]), re.M)
        lay.title = m.group(1).strip() if m else ''
    if reader.orphans:
        info(f'Hinweis: {reader.orphans} Fußnote(n) ohne erkennbare Ziffer im Text '
             'hängen jetzt am Ende ihrer Seite.')
    if shutil.which('pdfimages'):
        cmd = ['pdfimages', '-list'] + (['-f', str(pages[0]), '-l', str(pages[1])] if pages else [])
        images = len(run(cmd + [str(src)]).splitlines()) - 2
        if images > 0:
            info(f'Hinweis: Bilder übernimmt neusatz aus PDFs nicht (hier: {images}).')
    return {'pandoc-api-version': pandoc_api_version(), 'meta': {}, 'blocks': blocks}


# ---------------------------------------------------------------- Typst

def typst_fonts():
    return set(run(['typst', 'fonts']).splitlines())


def typst_str(s):
    return '"' + s.replace('\\', '\\\\').replace('"', '\\"') + '"'


def heading_numbering(numbering):
    levels = sorted(numbering)
    skip = levels[0] - 1
    top = LIST_STYLE.get(numbering[levels[0]], '1')
    lines = []
    if top == '1':
        lines.append(f'#set heading(numbering: (..n) => numbering("1.1", ..n.pos().slice({skip})))')
    else:
        lines.append(f'#set heading(numbering: (..n) => {{\n'
                     f'  let n = n.pos().slice({skip})\n'
                     f'  if n.len() == 1 {{ numbering("{top}.", ..n) }} else {{ numbering("1.1", ..n) }}\n'
                     f'}})')
    for lvl in range(1, levels[0]):
        lines.append(f'#show heading.where(level: {lvl}): set heading(numbering: none)')
    lines.append('#show <ohne-nummer>: set heading(numbering: none)')
    return lines


def preamble(lay, source, typ_name, args):
    fonts = typst_fonts()
    font = args.schrift or lay.font
    if font and font not in fonts:
        info(f'Hinweis: Die Schrift »{font}« ist nicht installiert, Typst nimmt {FALLBACK_FONT}.')
    font_list = [f for f in (font, FALLBACK_FONT) if f and f in fonts] or [FALLBACK_FONT]
    size = args.groesse or lay.size or 11

    if args.papier:
        page = f'paper: {typst_str(args.papier)}'
    elif lay.page:
        page = f'width: {lay.page[0]:.1f}mm, height: {lay.page[1]:.1f}mm'
    else:
        page = 'paper: "a4"'
    if lay.margins and not args.papier:
        m = dict(lay.margins)
        m = {k: max(12.0, v) for k, v in m.items()}
        m['bottom'] = max(m['bottom'], 18.0)      # Platz für die Seitenzahl
        side = ('inside', 'outside') if lay.mirror else ('left', 'right')
        page += (f',\n  margin: ({side[0]}: {m["left"]:.1f}mm, {side[1]}: {m["right"]:.1f}mm, '
                 f'top: {m["top"]:.1f}mm, bottom: {m["bottom"]:.1f}mm)')

    region = f', region: "{lay.region}"' if lay.region else ''
    par = ('first-line-indent: 1.5em, spacing: 0.65em' if lay.indent else 'spacing: 1.2em')
    lines = [
        f'// Neusatz von: {source.name}',
        f'// Erzeugt von neusatz am {datetime.date.today():%d.%m.%Y}. Text und Einstellungen',
        f'// dürfen bearbeitet werden, danach neu setzen mit:  neusatz "{typ_name}"',
        '',
        f'#set document(title: {typst_str(lay.title or source.stem)})',
        f'#set page({page},\n  numbering: "1")',
        f'#set text(lang: "{lay.lang}"{region}, font: ({", ".join(typst_str(f) for f in font_list)},), '
        f'size: {size:g}pt, hyphenate: true)',
        '',
        '// Knuth-Plass: "optimized" bricht jeden Absatz als Ganzes um.',
        f'#set par(justify: {"false" if args.flatterrand else "true"}, linebreaks: "optimized", {par})',
        '',
        '#show heading: set block(above: 1.8em, below: 1em)',
        '#show heading: set text(hyphenate: false)',
        '#show heading: set par(justify: false)',
    ]
    if lay.chapter_break:
        lines.append('#show heading.where(level: 1): it => { pagebreak(weak: true); it }')
    if lay.numbering:
        lines += heading_numbering(lay.numbering)
    lines += [
        '',
        '// Hilfsfunktionen für Pandocs Typst-Ausgabe',
        '#let blockquote(body) = quote(block: true, body)',
        '#let horizontalrule = align(center, line(length: 30%, stroke: 0.5pt))',
        '#show quote.where(block: true): set text(size: 0.95em)',
        '',
        '// ---------------------------------------------------------------- Text',
        '',
    ]
    return '\n'.join(lines)


def compile_typ(typ, pdf):
    r = subprocess.run(['typst', 'compile', '--root', str(typ.parent), str(typ), str(pdf)],
                       capture_output=True, text=True)
    if r.stderr.strip():
        print(r.stderr.strip(), file=sys.stderr)
    if r.returncode != 0:
        fail(f'Typst konnte {typ.name} nicht setzen. Die Fehlermeldung steht oben.')
    pages = ''
    if shutil.which('pdfinfo'):
        m = re.search(r'^Pages:\s+(\d+)', run(['pdfinfo', str(pdf)]), re.M)
        pages = f' ({m.group(1)} Seiten)' if m else ''
    info(f'Fertig: {pdf}{pages}')


def parse_pages(text):
    m = re.fullmatch(r'(\d+)(?:-(\d+))?', text or '')
    if not m:
        raise argparse.ArgumentTypeError('Seiten bitte als 20-40 oder 7 angeben.')
    a, b = int(m.group(1)), int(m.group(2) or m.group(1))
    if b < a:
        raise argparse.ArgumentTypeError('Die erste Seite muss vor der letzten liegen.')
    return a, b


def main():
    ap = argparse.ArgumentParser(
        prog='neusatz',
        description='Setzt .docx- und PDF-Dateien mit Typst neu, Zeilenumbruch nach Knuth-Plass.')
    ap.add_argument('eingabe', type=Path, help='.docx, .pdf oder eine schon erzeugte .typ-Datei')
    ap.add_argument('-o', '--ausgabe', type=Path, help='Ziel-PDF (Standard: <name>-neusatz.pdf neben der Eingabe)')
    ap.add_argument('--papier', help='Papierformat statt dem der Vorlage, z. B. a5, a4, us-letter')
    ap.add_argument('--schrift', help='Schriftfamilie statt der aus der Vorlage')
    ap.add_argument('--groesse', type=float, help='Schriftgröße in Punkt')
    ap.add_argument('--sprache', help='Sprache für Silbentrennung, z. B. de oder en (Standard: erkannt)')
    ap.add_argument('--flatterrand', action='store_true', help='Flattersatz statt Blocksatz')
    ap.add_argument('--trennstellen', action='store_true',
                    help='weiche Trennstriche der Vorlage behalten (Standard: Typst trennt selbst)')
    ap.add_argument('--seiten', type=parse_pages, help='nur diese PDF-Seiten übernehmen, z. B. 20-40')
    ap.add_argument('--nur-typ', action='store_true', help='nur die .typ-Datei schreiben, nicht setzen')
    args = ap.parse_args()

    src = args.eingabe
    if not src.is_file():
        fail(f'{src} gibt es nicht.')
    kind = src.suffix.lower()
    if kind == '.doc':
        fail('Bitte die .doc-Datei zuerst als .docx speichern.')
    if kind not in ('.docx', '.pdf', '.typ'):
        fail('Ich kann nur .docx-, .pdf- und .typ-Dateien lesen.')
    for tool in ['typst'] + (['pandoc'] if kind != '.typ' else []) + (['pdftohtml'] if kind == '.pdf' else []):
        if not shutil.which(tool):
            fail(f'{tool} fehlt. Siehe Kopf des Skripts für die Installation.')

    if kind == '.typ':
        compile_typ(src, args.ausgabe or src.with_suffix('.pdf'))
        return

    pdf = (args.ausgabe or src.with_name(f'{src.stem}-neusatz.pdf')).resolve()
    typ = pdf.with_suffix('.typ')
    if pdf.resolve() == src.resolve():
        fail('Die Ausgabe würde die Eingabe überschreiben.')
    media = f'{pdf.stem}-medien'
    lay = Layout()
    if kind == '.docx':
        doc = docx_to_ast(src, typ.parent, media, lay)
    else:
        doc = pdf_to_ast(src, args.seiten, lay)
    if args.sprache:
        lay.lang, lay.region = args.sprache.lower(), None

    doc['blocks'] = escape_ast(doc['blocks'], args.trennstellen)
    body = run(['pandoc', '-f', 'json', '-t', 'typst', '--wrap=none'],
               stdin=json.dumps(doc), cwd=typ.parent)
    if typ.exists():
        backup = typ.with_suffix('.typ.bak')
        typ.replace(backup)
        info(f'Die alte {typ.name} liegt jetzt als {backup.name} daneben.')
    typ.write_text(preamble(lay, src, typ.name, args) + body, encoding='utf-8')
    info(f'Typst-Quelltext: {typ}')
    if not args.nur_typ:
        compile_typ(typ, pdf)


if __name__ == '__main__':
    main()