import crochet
import pappyproxy
import shlex

from pappyproxy.plugin import main_context
from pappyproxy.console import load_reqlist
from pappyproxy.util import PappyException
from twisted.internet import defer

@crochet.wait_for(timeout=None)
@defer.inlineCallbacks
def tag(line):
    """
    Add a tag to requests.
    Usage: tag <tag> [request ids]
    You can tag as many requests as you want at the same time. If no
    ids are given, the tag will be applied to all in-context requests.
    """
    args = shlex.split(line)
    if len(args) == 0:
        raise PappyException('Tag name is required')
    tag = args[0]

    if len(args) > 1:
        reqs = yield load_reqlist(args[1], False)
        ids = [r.reqid for r in reqs]
        print 'Tagging %s with %s' % (', '.join(ids), tag)
    else:
        print "Tagging all in-context requests with %s" % tag
        reqs = main_context().active_requests

    for req in reqs:
        if tag not in req.tags:
            req.tags.append(tag)
            if req.saved:
                yield req.async_save()
            add_req(req)
        else:
            print 'Request %s already has tag %s' % (req.reqid, tag)

@crochet.wait_for(timeout=None)
@defer.inlineCallbacks
def untag(line):
    """
    Remove a tag from requests
    Usage: untag <tag> <request ids>
    You can provide as many request ids as you want and the tag will
    be removed from all of them. If no ids are given, the tag will 
    be removed from all in-context requests.
    """
    args = shlex.split(line)
    if len(args) == 0:
        raise PappyException("Tag and request ids are required")
    tag = args[0]

    ids = []
    if len(args) > 1:
        reqs = yield load_reqlist(args[1], False)
        ids = [r.reqid for r in reqs]
    else:
        print "Untagging all in-context requests with tag %s" % tag
        reqs = main_context().active_requests

    for req in reqs:
        if tag in req.tags:
            req.tags.remove(tag)
            if req.saved:
                yield req.async_save()
    if ids:
        print 'Tag %s removed from %s' % (tag, ', '.join(ids))
    
@crochet.wait_for(timeout=None)
@defer.inlineCallbacks
def clrtag(line):
    """
    Clear all the tags from requests
    Usage: clrtag <request ids>
    """
    args = shlex.split(line)
    if len(args) == 0:
        raise PappyException('No request IDs given')
    reqs = yield load_reqlist(args[0], False)

    for req in reqs:
        if req.tags:
            req.tags = []
            print 'Tags cleared from request %s' % (req.reqid)
            if req.saved:
                yield req.async_save()

###############
## Plugin hooks

def load_cmds(cmd):
    cmd.set_cmds({
        'clrtag': (clrtag, None),
        'untag': (untag, None),
        'tag': (tag, None),
    })
    cmd.add_aliases([
        #('rpy', ''),
    ])