##
# Copyright (c) 2001-2004 Twisted Matrix Laboratories.
# Copyright (c) 2010-2017 Apple Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
##
from __future__ import print_function

import re
from zope.interface import implements
import urllib
import tempfile

from twisted.internet import defer
from txweb2.stream import IStream, FileStream, BufferedStream, readStream
from txweb2.stream import generatorToStream, readAndDiscard
from txweb2 import http_headers
from cStringIO import StringIO

#
# Multipart MIME Reader
#


class MimeFormatError(Exception):
    pass

# parseContentDispositionFormData is absolutely horrible, but as
# browsers don't seem to believe in sensible quoting rules, it's
# really the only way to handle the header.  (Quotes can be in the
# filename, unescaped)
cd_regexp = re.compile(
    ' *form-data; *name="([^"]*)"(?:; *filename="(.*)")?$',
    re.IGNORECASE)


def parseContentDispositionFormData(value):
    match = cd_regexp.match(value)
    if not match:
        # Error parsing.
        raise ValueError("Unknown content-disposition format.")
    name = match.group(1)
    filename = match.group(2)
    return name, filename


# @defer.deferredGenerator
def _readHeaders(stream):
    """Read the MIME headers. Assumes we've just finished reading in the
    boundary string."""

    ctype = fieldname = filename = None
    headers = []

    # Now read headers
    while 1:
        line = stream.readline(size=1024)
        if isinstance(line, defer.Deferred):
            line = defer.waitForDeferred(line)
            yield line
            line = line.getResult()
        # print("GOT", line)
        if not line.endswith('\r\n'):
            if line == "":
                raise MimeFormatError("Unexpected end of stream.")
            else:
                raise MimeFormatError("Header line too long")

        line = line[:-2]  # strip \r\n
        if line == "":
            break  # End of headers

        parts = line.split(':', 1)
        if len(parts) != 2:
            raise MimeFormatError("Header did not have a :")
        name, value = parts
        name = name.lower()
        headers.append((name, value))

        if name == "content-type":
            ctype = http_headers.parseContentType(http_headers.tokenize((value,), foldCase=False))
        elif name == "content-disposition":
            fieldname, filename = parseContentDispositionFormData(value)

    if ctype is None:
        ctype = http_headers.MimeType('application', 'octet-stream')
    if fieldname is None:
        raise MimeFormatError('Content-disposition invalid or omitted.')

    # End of headers, return (field name, content-type, filename)
    yield fieldname, filename, ctype
    return
_readHeaders = defer.deferredGenerator(_readHeaders)


class _BoundaryWatchingStream(object):

    def __init__(self, stream, boundary):
        self.stream = stream
        self.boundary = boundary
        self.data = ''
        self.deferred = defer.Deferred()

    length = None  # unknown

    def read(self):
        if self.stream is None:
            if self.deferred is not None:
                deferred = self.deferred
                self.deferred = None
                deferred.callback(None)
            return None
        newdata = self.stream.read()
        if isinstance(newdata, defer.Deferred):
            return newdata.addCallbacks(self._gotRead, self._gotError)
        return self._gotRead(newdata)

    def _gotRead(self, newdata):
        if not newdata:
            raise MimeFormatError("Unexpected EOF")
        # BLECH, converting buffer back into string.
        self.data += str(newdata)
        data = self.data
        boundary = self.boundary
        off = data.find(boundary)

        if off == -1:
            # No full boundary, check for the first character
            off = data.rfind(boundary[0], max(0, len(data) - len(boundary)))
            if off != -1:
                # We could have a partial boundary, store it for next time
                self.data = data[off:]
                return data[:off]
            else:
                self.data = ''
                return data
        else:
            self.stream.pushback(data[off + len(boundary):])
            self.stream = None
            return data[:off]

    def _gotError(self, err):
        # Propogate error back to MultipartMimeStream also
        if self.deferred is not None:
            deferred = self.deferred
            self.deferred = None
            deferred.errback(err)
        return err

    def close(self):
        # Assume error will be raised again and handled by MMS?
        readAndDiscard(self).addErrback(lambda _: None)


class MultipartMimeStream(object):
    implements(IStream)

    def __init__(self, stream, boundary):
        self.stream = BufferedStream(stream)
        self.boundary = "--" + boundary
        self.first = True

    def read(self):
        """
        Return a deferred which will fire with a tuple of:
        (fieldname, filename, ctype, dataStream)
        or None when all done.

        Format errors will be sent to the errback.

        Returns None when all done.

        IMPORTANT: you *must* exhaust dataStream returned by this call
        before calling .read() again!
        """
        if self.first:
            self.first = False
            d = self._readFirstBoundary()
        else:
            d = self._readBoundaryLine()
        d.addCallback(self._doReadHeaders)
        d.addCallback(self._gotHeaders)
        return d

    def _readFirstBoundary(self):
        # print("_readFirstBoundary")
        line = self.stream.readline(size=1024)
        if isinstance(line, defer.Deferred):
            line = defer.waitForDeferred(line)
            yield line
            line = line.getResult()
        if line != self.boundary + '\r\n':
            raise MimeFormatError("Extra data before first boundary: %r looking for: %r" % (line, self.boundary + '\r\n'))

        self.boundary = "\r\n" + self.boundary
        yield True
        return
    _readFirstBoundary = defer.deferredGenerator(_readFirstBoundary)

    def _readBoundaryLine(self):
        # print("_readBoundaryLine")
        line = self.stream.readline(size=1024)
        if isinstance(line, defer.Deferred):
            line = defer.waitForDeferred(line)
            yield line
            line = line.getResult()

        if line == "--\r\n":
            # THE END!
            yield False
            return
        elif line != "\r\n":
            raise MimeFormatError("Unexpected data on same line as boundary: %r" % (line,))
        yield True
        return
    _readBoundaryLine = defer.deferredGenerator(_readBoundaryLine)

    def _doReadHeaders(self, morefields):
        # print("_doReadHeaders", morefields)
        if not morefields:
            return None
        return _readHeaders(self.stream)

    def _gotHeaders(self, headers):
        if headers is None:
            return None
        bws = _BoundaryWatchingStream(self.stream, self.boundary)
        self.deferred = bws.deferred
        ret = list(headers)
        ret.append(bws)
        return tuple(ret)


def readIntoFile(stream, outFile, maxlen):
    """Read the stream into a file, but not if it's longer than maxlen.
    Returns Deferred which will be triggered on finish.
    """
    curlen = [0]

    def done(_):
        return _

    def write(data):
        curlen[0] += len(data)
        if curlen[0] > maxlen:
            raise MimeFormatError("Maximum length of %d bytes exceeded." %
                                  maxlen)

        outFile.write(data)
    return readStream(stream, write).addBoth(done)


# @defer.deferredGenerator
def parseMultipartFormData(stream, boundary,
                           maxMem=100 * 1024, maxFields=1024, maxSize=10 * 1024 * 1024):
    # If the stream length is known to be too large upfront, abort immediately

    if stream.length is not None and stream.length > maxSize:
        raise MimeFormatError("Maximum length of %d bytes exceeded." % maxSize)

    mms = MultipartMimeStream(stream, boundary)
    numFields = 0
    args = {}
    files = {}

    while 1:
        datas = mms.read()
        if isinstance(datas, defer.Deferred):
            datas = defer.waitForDeferred(datas)
            yield datas
            datas = datas.getResult()
        if datas is None:
            break

        numFields += 1
        if numFields == maxFields:
            raise MimeFormatError("Maximum number of fields %d exceeded" % maxFields)

        # Parse data
        fieldname, filename, ctype, stream = datas
        if filename is None:
            # Not a file
            outfile = StringIO()
            maxBuf = min(maxSize, maxMem)
        else:
            outfile = tempfile.NamedTemporaryFile()
            maxBuf = maxSize
        x = readIntoFile(stream, outfile, maxBuf)
        if isinstance(x, defer.Deferred):
            x = defer.waitForDeferred(x)
            yield x
            x = x.getResult()
        if filename is None:
            # Is a normal form field
            outfile.seek(0)
            data = outfile.read()
            args.setdefault(fieldname, []).append(data)
            maxMem -= len(data)
            maxSize -= len(data)
        else:
            # Is a file upload
            maxSize -= outfile.tell()
            outfile.seek(0)
            files.setdefault(fieldname, []).append((filename, ctype, outfile))

    yield args, files
    return
parseMultipartFormData = defer.deferredGenerator(parseMultipartFormData)

#
# x-www-urlencoded reader
#


def parse_urlencoded_stream(
    input, maxMem=100 * 1024,
    keep_blank_values=False, strict_parsing=False
):

    lastdata = ''
    still_going = 1

    while still_going:
        try:
            yield input.wait
            data = input.next()
        except StopIteration:
            pairs = [lastdata]
            still_going = 0
        else:
            maxMem -= len(data)
            if maxMem < 0:
                raise MimeFormatError("Maximum length of %d bytes exceeded." %
                                      maxMem)
            pairs = str(data).split('&')
            pairs[0] = lastdata + pairs[0]
            lastdata = pairs.pop()

        for name_value in pairs:
            nv = name_value.split('=', 1)
            if len(nv) != 2:
                if strict_parsing:
                    raise MimeFormatError("bad query field: %s") % repr(name_value)
                continue
            if len(nv[1]) or keep_blank_values:
                name = urllib.unquote(nv[0].replace('+', ' '))
                value = urllib.unquote(nv[1].replace('+', ' '))
                yield name, value
parse_urlencoded_stream = generatorToStream(parse_urlencoded_stream)


def parse_urlencoded(stream, maxMem=100 * 1024, maxFields=1024,
                     keep_blank_values=False, strict_parsing=False):
    d = {}
    numFields = 0

    s = parse_urlencoded_stream(stream, maxMem, keep_blank_values, strict_parsing)

    while 1:
        datas = s.read()
        if isinstance(datas, defer.Deferred):
            datas = defer.waitForDeferred(datas)
            yield datas
            datas = datas.getResult()
        if datas is None:
            break
        name, value = datas

        numFields += 1
        if numFields == maxFields:
            raise MimeFormatError("Maximum number of fields %d exceeded" % maxFields)

        if name in d:
            d[name].append(value)
        else:
            d[name] = [value]
    yield d
    return
parse_urlencoded = defer.deferredGenerator(parse_urlencoded)


if __name__ == '__main__':
    d = parseMultipartFormData(
        FileStream(open("upload.txt")), "----------0xKhTmLbOuNdArY")
    from twext.python.log import Logger
    log = Logger()
    d.addErrback(log.err)

    def pr(s):
        print(s)
    d.addCallback(pr)


__all__ = ['parseMultipartFormData', 'parse_urlencoded', 'parse_urlencoded_stream', 'MultipartMimeStream', 'MimeFormatError']