import base64
import re

from twisted.protocols import basic
from twisted.internet import defer, interfaces

PORT = 2000

class NoQuote:
    def __init__(self, s):
        self.s = s

    def __str__(self):
        return self.s

def _parseItem(s):
    r"""
    >>> _parseItem(r'foo bar')
    ('foo', ' bar')
    >>> _parseItem(r'"foo" bar')
    ('foo', ' bar')
    >>> _parseItem(r'"\"foo\"" bar')
    ('"foo"', ' bar')
    >>> _parseItem(r'"foo\\" bar')
    ('foo\\', ' bar')
    """

    match = re.match(r'([^\\" ]+)', s)

    if match:
        return match.group(1), s[len(match.group(0)):]

    match = re.match(r'"((?:\\\\|\\"|[^\\"])*)"', s)

    if match:
        return (match.group(1).replace('\\"', '"').replace('\\\\', '\\'),
            s[len(match.group(0)):])

    return None, s

def _parseItems(line):
    r"""
    >>> _parseItems(r'OK "foo \"bar\""')
    ['OK', 'foo "bar"']
    >>> _parseItems(r'"foo')
    Traceback (most recent call last):
    ...
    AssertionError: "foo
    """

    items = []

    while line:
        item, line = _parseItem(line)

        if item is None:
            assert False, line

        items.append(item)

        if line.startswith(' '):
            line = line[1:]

    return items

class ManageSieveError(Exception):
    pass

class ManageSieveClient(basic.LineReceiver):
    def __init__(self):
        self.gotCapabilities = False
        self.capabilities = {}
        self.startedTLS = False
        self.pendingResponseMethod = None
        self.pendingResponseArgs = None
        self.pendingResponseBytes = None
        self.pendingResponseLiteral = None
        self.receivedLines = []
        self.contextFactory = None
        # Callback for implicit CAPABILITY command.
        self.pendingCommands = [defer.Deferred()]
        self.pendingCommands[0].addCallback(self._gotInitialCaps)
        self.commandQueue = []

    def _pushQueue(self):
        if self.commandQueue and not self.pendingCommands:
            command, d, data = self.commandQueue.pop(0)
            #print 'send:', command
            self.sendLine(command)

            if data:
                for datum in data:
                    #print 'send:', datum
                    self.sendLine(datum)

            self.pendingCommands.append(d)

    def sendCommand(self, atom, args, data=None):
        def quote(s):
            return '"%s"' % s.replace('\\', '\\\\'.replace('"', '\\"'))

        def formatArg(arg):
            if isinstance(arg, NoQuote):
                return str(arg)
            else:
                return quote(arg)

        command = atom + ' ' + ' '.join(map(formatArg, args))
        d = defer.Deferred()
        self.commandQueue.append((command, d, data))
        self._pushQueue()
        return d

    def _gotCapabilities(self, (code, caps)):
        for items in map(_parseItems, caps):
            if len(items) == 1:
                self.capabilities[items[0]] = True
            elif len(items) == 2:
                self.capabilities[items[0]] = items[1]

        #print 'capabilities:', self.capabilities

    def getCapabilities(self):
        d = self.sendCommand('CAPABILITY', [])
        d.addCallback(self._gotCapabilities)
        return d

    def gotInitialCaps(self):
        # For overriding.
        pass

    def _gotInitialCaps(self, (code, caps)):
        self._gotCapabilities((code, caps))
        self.gotInitialCaps()

    def okReceived(self, code):
        assert self.pendingCommands
        command = self.pendingCommands.pop(0)
        command.callback((code, self.receivedLines))
        self.receivedLines = []
        self._pushQueue()

    def noReceived(self, code):
        assert self.pendingCommands
        self.pendingCommands.pop(0).errback(ManageSieveError(code[0]))

    def byeReceived(self, code):
        pass

    def lineReceived(self, line):
        #print 'received:', repr(line)

        if self.pendingResponseMethod:
            self.pendingResponseBytes -= len(line)
            self.pendingResponseLiteral += line

            if self.pendingResponseBytes < 0:
                raise RuntimeError("ManageSieve syntax error")

            if self.pendingResponseBytes == 0:
                self.pendingResponseMethod(
                    self.pendingResponseArgs + [self.pendingResponseLiteral])
                self.pendingResponseMethod = None
            else:
                self.pendingResponseBytes -= 2
                self.pendingResponseLiteral += '\r\n'

            return

        match = re.match('^(ok|no|bye)(?:| (.+))$', line, re.I)

        if match:
            method = {
                'ok': self.okReceived,
                'no': self.noReceived,
                'bye': self.byeReceived
                }[match.group(1).lower()]

            if match.group(2):
                items = _parseItems(match.group(2))
                match = re.match('{(\d+)}', items[-1])

                if match:
                    self.pendingResponseMethod = method
                    self.pendingResponseArgs = items[:-1]
                    self.pendingResponseBytes = int(match.group(1))
                    self.pendingResponseLiteral = ''
                else:
                    method(items)
            else:
                method([])
        else:
            self.receivedLines.append(line)

    def _getContextFactory(self):
        # Based on IMAP4Client code.

        if self.contextFactory is not None:
            return self.contextFactory
        try:
            from twisted.internet import ssl
        except ImportError:
            return None
        else:
            contextFactory = ssl.ClientContextFactory()
            contextFactory.method = ssl.SSL.TLSv1_METHOD
            return contextFactory

    def startTLS(self):
        if self.startedTLS:
            return defer.fail(RuntimeError("already using TLS"))

        if 'STARTTLS' not in self.capabilities:
            return defer.fail(RuntimeError("server doesn't support TLS"))

        contextFactory = self._getContextFactory()

        if contextFactory is None:
            return defer.fail(RuntimeError("failed to create a TLS context"))

        tls = interfaces.ITLSTransport(self.transport, None)

        if tls is None:
            return defer.fail(RuntimeError(
                "transport does not implement ITLSTransport"))

        # XXX: Yuck.
        d = defer.Deferred()
        d2 = self.sendCommand('STARTTLS', [])
        d2.addCallback(self._okStartTLS, contextFactory, d)
        return d

    def _okStartTLS(self, (code, result), contextFactory, d):
        self.transport.startTLS(contextFactory)
        self.startedTLS = True

        ## Callback for implicit CAPABILITY response.
        #self.pendingCommands.insert(0, defer.Deferred())
        #self.pendingCommands[0].addCallback(self._gotCapabilities)

        # XXX: WTF. We don't see the capabilities the server sends until we
        # send an explicit CAPABILITY command, at which point we see the
        # implicit capabilities and the response to the explicit command.

        d3 = self.getCapabilities()
        self.pendingCommands.insert(0, defer.Deferred())
        d3.chainDeferred(d)

    def authenticate(self, uid, secret):
        if 'SASL' not in self.capabilities:
            raise RuntimeError("server doesn't support SASL")

        methods = self.capabilities['SASL'].split(' ')

        if 'PLAIN' not in methods:
            raise RuntimeError("no supported authentication methods")

        return self.sendCommand('AUTHENTICATE',
            ['PLAIN', base64.b64encode('%s\0%s\0%s' % (uid, uid, secret))])

    def listScripts(self):
        def ok((code, result)):
            scripts = []
            active = None

            for items in map(_parseItems, result):
                if items:
                    script = items[0]

                    if len(items) == 2 and items[1] == 'ACTIVE':
                        active = script

                    scripts.append(script)

            return scripts, active

        d = self.sendCommand('LISTSCRIPTS', [])
        d.addCallback(ok)
        return d

    def putScript(self, name, script):
        def ok((code, result)):
            return code[0]

        lines = script.split('\r\n')
        d = self.sendCommand('PUTSCRIPT',
            [name, NoQuote('{%d+}' % len(script))], lines)
        d.addCallback(ok)
        return d

    def getScript(self, name):
        def ok((code, result)):
            if result and re.match('{\d+}', result[0]):
                return '\r\n'.join(result[1:])
            else:
                return '\r\n'.join(result)

        d = self.sendCommand('GETSCRIPT', [name])
        d.addCallback(ok)
        return d

    def deleteScript(self, name):
        def ok((code, result)):
            return code[0]

        d = self.sendCommand('DELETESCRIPT', [name])
        d.addCallback(ok)
        return d

    def setActive(self, name):
        def ok((code, result)):
            return code[0]

        if name is None:
            name = ''

        d = self.sendCommand('SETACTIVE', [name])
        d.addCallback(ok)
        return d