#!/usr/bin/python
# vim: set fileencoding=utf-8 : sts=4 : et : sw=4
#    This is a component of mailpie, a full-text search for email
#
#    Copyright © 2008 Jeff Epler <jepler@unpythonic.net>
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import sys, os, subprocess, shutil, getopt, tempfile, re, shlex
import mailpie.swishfilter, mailpie.parsedate, mailpie.log
import bsddb

def usage(result=0):
    print """mailpie-search: Put search results in a mailbox
Usage: %s [-o file|-p|-m] [-b date] [-a date] [-M mailreader] terms...
    -B dir, --base=dir      Choose the base location of the mailpie storage
                            Default: $HOME/.mailpie

Action:
    -c,      --count        Write number of matches
    -o FILE, --out=FILE     Write result to FILE
    -p,      --stdout       Write result to stdout
    -r,      --read:        Open mailbox in mailreader (default)

    -i,      --show-ids     Show the mailpie message ID in a header

    -M prog, --mailreader=PROG
                            Mailreader to invoke (default: "mutt -R -f")

Searching:
    -b DATE, --before=DATE  Only include messages before DATE
    -a DATE, --after=DATE   Only include messages after DATE
                            Date formats are like "date -d=DATE"

    -l NUM, --limit=NUM     Return no more than NUM

    -t, --thread            Include full threads
    -n, --no-thread         Do not include threads (default)

    terms...                See SWISH-RUN(1) for details on search terms
""" % os.path.basename(sys.argv[0])
    raise SystemExit, result

class SwishError(RuntimeError): pass

def run_swish(indexes, args, before=None, after=None):
    if not args:
        return
    swishargs = ['swish-e', '-H0',
        '-x<swishdocpath>\n',
        '-w', " ".join(args)] + ['-f%s' % f for f in indexes]
    if before or after:
        if before is None: before = sys.maxint
        if after is None: after = 0
        swishargs.extend(["-L", "date", str(after), str(before)])

    swish = subprocess.Popen(swishargs, stdout=subprocess.PIPE,
                stdin=open("/dev/null", "r"))
    output = [row.strip("\n") for row in swish.stdout]
    exitcode = swish.wait()
    if exitcode != 0:
        if output and output[-1] == ".": del output[-1]
        if output:
            raise SwishError, (
                "Swish exited with code %d and the following information:\n    "
                    % exitcode
                + "\n    ".join(output)+"\n\nSee 'man SWISH-RUN' for details on"
                "swish search terms")
        else:
            raise SwishError, "Swish exited with code %d" % exitcode
    return output

class Count: pass

def copymail(path, f, show_ids):
    if not os.path.isfile(path): return 0
    if not f: return 1
    i = open(path)
    if show_ids:
        printed = False
        for line in i:
            if (not printed) and (line == "\n" or line == "\r\n"):
                f.write("X-Mailpie-ID: %s\n" % path.replace("/", "")[-40:])
                printed = True
            f.write(line)
    else:
        shutil.copyfileobj(i, f)
    return 1

def maybe_progress(m, a, b):
    a1 = a
    while a1 > 100:
        if a1 % 10: return
        a1 /= 10
    mailpie.log.progress(m, a, b)

def main(args):
    before = after = target = None
    if os.environ.get("MAILPIE_READER"):
        mailreader = shlex.split(os.environ["MAILPIE_READER"])
    else:
        mailreader = ["mutt", "-R", "-f"]
    limit = 2000
    thread = False
    show_ids = False

    base = os.path.expanduser("~/.mailpie")

    try:
        opts, args = getopt.getopt(args, "B:co:prM:a:b:l:tnih?",
            ("base=", "count", "out=", "stdout", "read", "mailreader",
             "help", "after=", "before=", "limit=", "thread", "no-thread",
             "show-ids"))
    except getopt.GetoptError, detail:
        usage(detail)

    for k, v in opts:
        if k in ("-o", "--out"): target = v
        if k in ("-B", "--base"): base = v
        elif k in ("-p", "--stdout"): target = "-"
        elif k in ("-r", "--read"): target = None
        elif k in ("-c", "--count"): target = Count
        elif k in ("-M", "--mailreader"): mailreader = shlex.split(v)
        elif k in ("-l", "--limit"): limit = int(v)
        elif k in ("-a", "--after"): after = mailpie.parsedate.parsedate(v)
        elif k in ("-b", "--before"): before = mailpie.parsedate.parsedate(v)
        elif k in ("-t", "--thread"): thread = True
        elif k in ("-n", "--no-thread"): thread = False
        elif k in ("-l", "--limit"): limit = int(v)
        elif k in ("-i", "--show-ids"): show_ids = True
        elif k in ("--help", "-h", "-?"): usage()

    indexes = mailpie.swishfilter.get_index_files(base)

    if target is None:
        f = tempfile.NamedTemporaryFile(prefix="mailpie", suffix=".mbox")
    elif target == "-":
        f = sys.stdout
    elif target is Count:
        f = None
    else:
        f = open(target, "w")

    count = 0
    msgids = set()

    def q(arg):
        if not " " in arg: return arg

        if "=" in arg:
            a, b = arg.split("=", 1)
            return '%s="%s"' % (a,b)

        return '"%s"' % arg

    if thread:
        thread_db = bsddb.hashopen(base + ".thread.db", "r")

    sargs = []
    for a in args:
        if a.startswith("after="): after = mailpie.parsedate.parsedate(a[6:])
        elif a.startswith("before="): before = mailpie.parsedate.parsedate(a[6:])
        else: sargs.append(q(a))

    mailpie.log.progress("Searching for messages")
    try:
        hashes = run_swish(indexes, sargs, before, after)
    except SwishError, detail:
        raise SystemExit, str(detail)

    if len(hashes) > limit:
        mailpie.log.log("Will only show the first %d of %d matches", limit, len(hashes))
    hashes = hashes[:limit]
    total = len(hashes)
    for i, hash in enumerate(hashes):
        maybe_progress("Retrieving message %d of %d", i, total)
        path = os.path.join(base, hash[:2], hash[2:])
        count += copymail(path, f, show_ids)
        if count == limit: break
        if thread:
            msgids.add(thread_db.get("h:" + hash))

    if thread:
        printed = set(msgids)
        visited = set()
        while msgids and count < limit:
            msgid = msgids.pop()
            mailpie.log.progress("Chasing: %s and %d more", msgid, len(msgids))
            if msgid in visited: continue
            visited.add(msgid)
            th = set(thread_db.get("t:" + msgid, "").split()) - printed
            msgids.update(th)
            if not msgid in printed:
                hash = thread_db.get('m:' + msgid)
                if hash is None: continue
                path = os.path.join(base, hash[:2], hash[2:])
                count += copymail(path, f, show_ids)
    mailpie.log.progress_clear()

    if target is None:
        if count == 0:
            mailpie.log.log("No match")
        else:
            f.flush()
            os.spawnvp(os.P_WAIT, mailreader[0], mailreader + [f.name])

    if target is Count:
        print count

if __name__ == '__main__':
    main(sys.argv[1:])
