#!/usr/bin/env python3
"""
stegano_extract.py -- read the message hidden by stegano_embed.py from a PDF.

Collects, in content-stream order, every glyph whose fill colour is
"black 99 %" (DeviceGray 0.01 g) and decodes the period code:
    "."      -> space
    ".."     -> "."
    ".CR."   -> carriage return (newline)

Usage
-----
    python3 stegano_extract.py book.pdf
    python3 stegano_extract.py book.pdf --raw          # undecoded stream
    python3 stegano_extract.py book.pdf --colors       # colour histogram (debugging)

Requires:  pip install pdfminer.six
"""
import argparse
import sys
import unicodedata
from collections import Counter

from pdfminer.high_level import extract_pages
from pdfminer.layout import LTChar, LTContainer

# Spacing accents that pdfTeX draws separately in OT1 (\"u = diaeresis + u).
ACCENTS = {
    '¨': '̈', '´': '́', '`': '̀', '`': '̀',
    '^': '̂', 'ˆ': '̂', '~': '̃', '˜': '̃',
    'ˇ': '̌', '˘': '̆', '˙': '̇',
    '˚': '̊', '¸': '̧', '˝': '̋',
    '¯': '̄',
}
LIGATURES = {'ff': 'ff', 'fi': 'fi', 'fl': 'fl',
             'ffi': 'ffi', 'ffl': 'ffl'}


def color_tuple(ncolor):
    if ncolor is None:
        return ()
    if isinstance(ncolor, (int, float)):
        return (float(ncolor),)
    try:
        return tuple(float(x) for x in ncolor if isinstance(x, (int, float)))
    except TypeError:
        return ()


def is_target(color, gray, tol):
    """Gray 0.01 (or an RGB triple of three equal values 0.01)."""
    if len(color) == 1:
        return abs(color[0] - gray) <= tol
    if len(color) == 3:
        return all(abs(c - gray) <= tol for c in color)
    return False


def walk(obj):
    """Yield LTChar objects in content-stream order (no layout analysis)."""
    if isinstance(obj, LTChar):
        yield obj
    elif isinstance(obj, LTContainer):
        for child in obj:
            yield from walk(child)


def decode_stream(s):
    out, i = [], 0
    while i < len(s):
        if s.startswith('.CR.', i):
            out.append('\n'); i += 4
        elif s[i] == '.':
            if s.startswith('..', i):
                out.append('.'); i += 2
            else:
                out.append(' '); i += 1
        else:
            out.append(s[i]); i += 1
    return ''.join(out)


def clean_glyphs(texts, warn):
    """Ligatures, separately drawn accents (OT1), (cid:N) placeholders."""
    out = []
    pending_accent = None
    for t in texts:
        if t.startswith('(cid:'):
            warn.append(t)
            t = '?'
        t = LIGATURES.get(t, t)
        if t in ACCENTS and len(t) == 1:
            if pending_accent:
                out.append(pending_accent)
            pending_accent = ACCENTS[t]
            continue
        if pending_accent:
            if len(t) == 1 and t.isalpha():
                t = unicodedata.normalize('NFC', t + pending_accent)
            else:
                out.append(pending_accent)
            pending_accent = None
        out.append(t)
    if pending_accent:
        out.append(pending_accent)
    return unicodedata.normalize('NFC', ''.join(out))


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('pdf')
    ap.add_argument('--gray', type=float, default=0.01,
                    help='gray level of the marked glyphs (default 0.01)')
    ap.add_argument('--tol', type=float, default=0.0015,
                    help='tolerance on the gray level (default 0.0015)')
    ap.add_argument('--raw', action='store_true',
                    help='print the marked characters without period decoding')
    ap.add_argument('--colors', action='store_true',
                    help='list the fill colours found (debugging) and exit')
    ap.add_argument('--keep-curly', action='store_true',
                    help="keep U+2019 as is (default: output it as ASCII ', because TeX fonts "
                         "draw an ASCII apostrophe with the right-quote glyph)")
    ap.add_argument('-o', '--output', help='write the message to this UTF-8 file')
    a = ap.parse_args()

    hist = Counter()
    marked, warn = [], []
    for page in extract_pages(a.pdf, laparams=None):
        for ch in walk(page):
            t = ch.get_text()
            if not t or t.isspace():
                continue
            col = color_tuple(getattr(ch.graphicstate, 'ncolor', None))
            if a.colors:
                hist[(getattr(ch, 'ncs', None) and ch.ncs.name, col)] += 1
            if is_target(col, a.gray, a.tol):
                marked.append(t)

    if a.colors:
        for (cs, col), cnt in hist.most_common():
            print(f"{cnt:8d} glyphs  colourspace={cs}  value={col}")
        return

    stream = clean_glyphs(marked, warn)
    if not a.keep_curly:
        stream = stream.replace('\u2019', "'")
    if warn:
        print(f"warning: {len(warn)} glyph(s) without Unicode mapping, shown as '?'",
              file=sys.stderr)
    if not stream:
        print("no glyphs in black 99 % found (try --colors)", file=sys.stderr)
        sys.exit(1)
    text = stream if a.raw else decode_stream(stream)
    if a.output:
        with open(a.output, 'w', encoding='utf-8') as f:
            f.write(text)
    sys.stdout.reconfigure(encoding='utf-8')
    print(text)


if __name__ == '__main__':
    main()