#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""Query subtrees of an org-mode outline.

Extracts the subtrees from the desired headers.

For example:

    $ cat tmp.org
    * foo
    ** bar
    baz
    bazz
    ** quux
    * snorf
    ** borf
    ** snorf
    ** horg
    $ ./remember.py bar snorf < tmp.org
    ** bar
    baz
    bazz
    * snorf
    ** borf
    ** snorf
    ** horg

Arguably the above should return `* foo` first, to give context to
`** bar`, but it doesn’t.

There’s a hacky option to use, rather than a fixed list, headers that
are dates in ISO-8601 format whose distance in days from the present
date is an integer power of some base, like 3:

    $ ./remember.py -n 3 x --since 2018-05-30 -i somefile
    ** 2018-05-27
    stuff
    ** 2018-05-29
    stuff
    ** 2018-05-30
    *** stuff
    ** 2018-05-31
    *** stuff
    *** stuff

The fact that it needs “x” is a bug.

"""
from __future__ import print_function

import argparse
import datetime
import math
import re
import sys


try:
    from itertools import imap as map
except ImportError:
    pass


header_re = re.compile(r'(\*+)\s+(.*)')
def parse(line):
    'Return level of org-mode header (or None) and the other text on the line.'
    mo = header_re.match(line)
    if not mo:
        return None, line
    return len(mo.group(1)), mo.group(2)


def deparse(level, text):
    'Inverse of parse, as far as possible.'
    return text if level is None else '{} {}\n'.format('*' * level, text)


def beneath(lines, wants):
    'Yield (level, text) pairs at or beneath wanted headers.'
    lines = map(parse, lines)
    while True:
        for level, text in lines:
            if level is not None and wants(text.strip()):
                cutoff_level = level
                yield level, text
                break
        else:
            return

        for level, text in lines:
            if (level is not None and level <= cutoff_level
                    and not wants(text.strip())):
                break
            yield level, text
        else:
            return    


def date(s):
    'Idempotently convert an ISO-8601 date string to a datetime.date.'
    return (s if isinstance(s, datetime.date) else
            datetime.datetime.strptime(s, '%Y-%m-%d').date())


def wants_days(since, base):
    '''Return a predicate to identify dates at powers of `base` from `since`.

    That is, ``wants_days('2018-03-19', 5)(d)`` returns true iff ``d``
    is a date that is 1, 5, 25, 125, or some other power of 5 away
    from 2018-03-19, or also (as a special case) if it is 2018-03-19
    itself.

    Both `since` and the date passed in to the returned predicate can
    be either datetime.date objects or ISO-8601 date strings.

    '''
    since = date(since)
    def wants(header):
        try:
            candidate = date(header)
        except ValueError:
            return False
        diff = abs(candidate - since).days
        if not diff:
            return True

        log = math.log(diff) / math.log(base)
        return int(log) == log

    return wants


def main():
    'Command-line interface.'
    formatter = argparse.RawDescriptionHelpFormatter
    p = argparse.ArgumentParser(description=__doc__, formatter_class=formatter)
    p.add_argument('-i', '--input', type=open, default=sys.stdin,
                   help='input file to read (default stdin)')
    p.add_argument('header', nargs='+',
                   help='one or more headers to match literally')
    p.add_argument('--since', default=datetime.date.today().isoformat(),
                   type=date,
                   help='date to use as reference (default today, %(default)r)')
    p.add_argument('-n', '--days-base', type=int,
                   help='ignore header and display dates whose time to SINCE'
                   + ' has an integer logarithm to the base DAYS_BASE')
    args = p.parse_args()

    wants = (set(args.header).__contains__ if args.days_base is None else
             wants_days(args.since, args.days_base))
    for level, text in beneath(args.input, wants):
        print(deparse(level, text), end='')


if __name__ == '__main__':
    main()