#!/usr/bin/env python
# License: http://creativecommons.org/licenses/by-sa/2.0/de/
# http://guido.vonrudorff.de/store-okular-annotations-in-pdf

import subprocess, argparse, os, poppler, cairo, math
import xml.etree.ElementTree as ET

def get_okular_xml_path(filename, kdepath):
    if kdepath is None:
        kdepath = subprocess.check_output(['kde4-config', '--localprefix'])
    kdepath = os.path.abspath(kdepath.strip())
    kdepath += '/share/apps/okular/docdata/'
    filepath = os.path.basename(filename)
    xmlpath = '%s%d.%s.xml' % (kdepath, os.path.getsize(filename), filepath)
    return xmlpath

def draw_annotation(po, ctx, annotation):
    # highlighter
    width, height = po.get_size()
    if annotation.attrib['type'] == '4':
        for hl in annotation.iter('hl'):
            underline = False
            if 'type' in hl.attrib:
                ctx.set_source_rgba(0, 0, 0, .5)
                underline = True
            else:
                ctx.set_source_rgba(1, 1, 0, .5)
            for q in hl.iter('quad'):
                x = float(q.attrib['dx'])*width
                y = float(q.attrib['dy'])*height
                y2 = float(q.attrib['ay'])*height
                x2 = float(q.attrib['bx'])*width
                w = abs(float(q.attrib['bx'])-float(q.attrib['ax']))*width
                h = abs(float(q.attrib['dy'])-float(q.attrib['ay']))*height
                if underline:
                    ctx.move_to(x, y2)
                    ctx.line_to(x2, y2)
                    ctx.stroke()
                else:
                    ctx.rectangle(x, y, w, h)
                    ctx.fill()
    elif annotation.attrib['type'] == '6':
        ctx.set_source_rgba(0, 1, 0, .5)
        ctx.set_line_width(4)
        first = True
        for p in annotation.iter('path'):
            for pp in p.iter('point'):
                x = float(pp.attrib['x'])*width
                y = float(pp.attrib['y'])*height
                if first:
                    ctx.move_to(x, y)
                    first = False
                else:
                    ctx.line_to(x, y)
            ctx.close_path()
            ctx.stroke()
    elif annotation.attrib['type'] == '2':
        ctx.set_line_width(4)
        for l in annotation.iter('line'):
            first = True
            for p in l.iter('point'):
                x = float(p.attrib['x'])*width
                y = float(p.attrib['y'])*height
                if first:
                    ctx.move_to(x, y)
                    first = False
                else:
                    ctx.line_to(x, y)
            if 'closed' in l.attrib:
                ctx.set_source_rgba(0, 0, 1, .5)
                ctx.close_path()
            else:
                ctx.set_source_rgba(1, 1, 0, .5)
            ctx.stroke()
    elif annotation.attrib['type'] == '3':
        ctx.set_source_rgba(0, 1, 1, .5)
        ctx.set_line_width(4)
        for b in annotation.iter('boundary'):
            x1 = float(b.attrib['l'])*width
            y2 = float(b.attrib['b'])*height
            x2 = float(b.attrib['r'])*width
            y1 = float(b.attrib['t'])*height
            h = y2-y1
            w = x2-x1
            ctx.save()
            ctx.translate(x1+w/2, y1+h/2)
            ctx.scale(w/2, h/2)
            ctx.arc(0., 0., 1., 0., 2 * math.pi)
            ctx.restore()
            ctx.stroke()
    else:
        print 'Unknown annotation type %s.' % annotation.attrib['type']

# arguments
parser = argparse.ArgumentParser()
parser.add_argument('INPUTPDF', help='the file to work on')
parser.add_argument('OUTPUTPDF', help='the file to work on')
parser.add_argument('--kdepath', help='your $HOME/.kde folder. If not specified, kde4-config is called.')
args = parser.parse_args()

# read file
try:
    xmlpath = get_okular_xml_path(args.INPUTPDF, args.kdepath)
except OSError:
    print 'File %s not found or kde4-config not found. Try --kdepath. Aborting.' % args.INPUTPDF
    exit(1)
if not os.path.exists(xmlpath):
    print 'No okular XML annotation file found. Try --kdepath. Aborting.'
    print 'Hint: currently, the file is expected to be %s.' % xmlpath
    exit(1)

# search annotations
f = poppler.document_new_from_file('file://%s' % os.path.abspath(args.INPUTPDF), None)
numpages = f.get_n_pages()

xml = ET.parse(xmlpath)
r = xml.getroot()
po = f.get_page(0)
surface = cairo.PDFSurface(args.OUTPUTPDF, int(po.get_size()[0]), int(po.get_size()[1]))
refdict = {}
for page in r.iter('page'):
    number = int(page.attrib['number'])
    refdict[number] = page
    if number >= numpages:
        print 'Skipping annotation: wrong page number.'
        continue

for number in range(numpages):
    print 'Reading page %d of %d' % (number+1, numpages)
    po = f.get_page(number)
    
    # render page content
    ctx = cairo.Context(surface)

    # add annotations
    if number in refdict:
        for annotation in refdict[number].iter('annotation'):
            draw_annotation(po, ctx, annotation)

    po.render(ctx)
    # background and output
    ctx.set_operator(cairo.OPERATOR_DEST_OVER)
    ctx.set_source_rgb(1, 1, 1)
    ctx.paint()
    surface.show_page()
surface.finish()