# The MIT License (MIT)
#
# Copyright (c) 2014-2017 Susam Pal
#
# 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.

"""Ice - WSGI on the rocks.

Ice is a simple and tiny WSGI microframework meant for developing small
Python web applications.
"""

__version__ = '0.0.2'
__date__ = '25 March 2014'
__author__ = 'Susam Pal <susam@susam.in>'
__credits__ = ('Marcel Hellkamp, for writing bottle, the inspiration '
               'behind ice.')

import collections
import itertools
import re
import cgi
import urllib.parse
import http.server
import http.cookies
import os
import mimetypes

def cube():
    """Return an Ice application with a default home page.

    Create :class:`Ice` object, add a route to return the default page
    when a client requests the server root, i.e. /, using HTTP GET
    method, add an error handler to return HTTP error pages when an
    error occurs and return this object. The returned object can be used
    as a WSGI application.

    Returns:
      Ice: WSGI application.
    """
    app = Ice()

    @app.get('/')
    def default_home_page():
        """Return a default home page."""
        return simple_html('It works!',
                           '
It works!
\n'
                           '
This is the default ice web page.
')

    @app.error()
    def generic_error_page():
        """Return a simple and generic error page."""
        return simple_html(app.response.status_line,
                           '
{title}
\n'
                           '
{description}
\n'
                           '
\n'
                           '<address>Ice/{version}</address>'.format(
                           title=app.response.status_line,
                           description=app.response.status_detail,
                           version=__version__))

    def simple_html(title, body):
        """Return a simple HTML page."""
        return (
            '<!DOCTYPE html>\n'
            '\n<title>{title}</title>\n'
            '\n{body}\n\n\n'
        ).format(title=title, body=body)

    return app

class Ice:

    """A single WSGI application.

    Each instance of this class is a single, distinct callable object
    that functions as WSGI application.
    """

    def __init__(self):
        """Initialize the application."""
        self._router = Router()
        self._server = None
        self._error_handlers = {}

    def run(self, host='127.0.0.1', port=8080):
        """Run the application using a simple WSGI server.

        Arguments:
          host (str, optional): Host on which to listen.
          port (int, optional): Port number on which to listen.
        """
        from wsgiref import simple_server
        self._server = simple_server.make_server(host, port, self)
        self._server.serve_forever()

    def exit(self):
        """Stop the simple WSGI server running the appliation."""
        if self._server is not None:
            self._server.shutdown()
            self._server.server_close()
            self._server = None

    def running(self):
        """Return ``True`` iff simple WSGI server is running.

        Returns:
          bool: ``True`` if simple WSGI server associated with this
          application is running, ``False`` otherwise.
        """
        return self._server is not None

    def get(self, pattern):
        """Decorator to add route for an HTTP GET request.

        Arguments:
          pattern (str): Routing pattern the path must match.

        Returns:
          function: Decorator to add route for HTTP GET request.
        """
        return self.route('GET', pattern)

    def post(self, pattern):
        """Decorator to add route for an HTTP POST request.

        Arguments:
          pattern (str): Routing pattern the path must match.

        Returns:
          function: Decorator to add route for HTTP POST request.
        """
        return self.route('POST', pattern)

    def route(self, method, pattern):
        """Decorator to add route for a request with any HTTP method.

        Arguments:
          method (str): HTTP method name, e.g. GET, POST, etc.
          pattern (str): Routing pattern the path must match.

        Returns:
          function: Decorator function to add route.
        """
        def decorator(callback):
            self._router.add(method, pattern, callback)
            return callback
        return decorator

    def error(self, status=None):
        """Decorator to add a callback that generates error page.

        The *status* parameter specifies the HTTP response status code
        for which the decorated callback should be invoked. If the
        *status* argument is not specified, then the decorated callable
        is considered to be a fallback callback.

        A fallback callback, when defined, is invoked to generate the
        error page for any HTTP response representing an error when
        there is no error handler defined explicitly for the response
        code of the HTTP response.

        Arguments:
          status(int, optional): HTTP response status code.

        Returns:
          function: Decorator function to add error handler.
        """
        def decorator(callback):
            self._error_handlers[status] = callback
            return callback
        return decorator

    def static(self, root, path, media_type=None, charset='UTF-8'):
        """Send content of a static file as response.

        The path to the document root directory should be specified as
        the root argument. This is very important to prevent directory
        traversal attack. This method guarantees that only files within
        the document root directory are served and no files outside this
        directory can be accessed by a client.

        The path to the actual file to be returned should be specified
        as the path argument. This path must be relative to the document
        directory.

        The *media_type* and *charset* arguments are used to set the
        Content-Type header of the HTTP response. If *media_type*
        is not specified or specified as ``None`` (the default), then it
        is guessed from the filename of the file to be returned.

        Arguments:
          root (str): Path to document root directory.
          path (str): Path to file relative to document root directory.
          media_type (str, optional): Media type of file.
          charset (str, optional): Character set of file.

        Returns:
          bytes: Content of file to be returned in the HTTP response.
        """
        root = os.path.abspath(os.path.join(root, ''))
        path = os.path.abspath(os.path.join(root, path.lstrip('/\\')))

        # Save the filename from the path in the response state, so that
        # a following download() call can default to this filename for
        # downloadable file when filename is not explicitly specified.
        self.response.state['filename'] = os.path.basename(path)

        if not path.startswith(root):
            return 403
        elif not os.path.isfile(path):
            return 404

        if media_type is not None:
            self.response.media_type = media_type
        else:
            self.response.media_type = mimetypes.guess_type(path)[0]
        self.response.charset = charset

        with open(path, 'rb') as f:
            return f.read()

    def download(self, content, filename=None,
                 media_type=None, charset='UTF-8'):
        """Send content as attachment (downloadable file).

        The *content* is sent after setting Content-Disposition header
        such that the client prompts the user to save the content
        locally as a file. An HTTP response status code may be specified
        as *content*. If the status code is not ``200``, then this
        method does nothing and returns the status code.

        The filename used for the download is determined according to
        the following rules. The rules are followed in the specified
        order.

          1. If *filename* is specified, then the base name from this
             argument, i.e. ``os.path.basename(filename)``, is used as the
             filename for the download.
          2. If *filename* is not specified or specified as ``None``
             (the default), then the base name from the file path
             specified to a previous :meth:`static` call made while
             handling the current request is used.
          3. If *filename* is not specified and there was no
             :meth:`static` call made previously for the current
             request, then the base name from the current HTTP request
             path is used.
          4. As a result of the above steps, if the resultant *filename*
             turns out to be empty, then :exc:`ice.LogicError` is raised.

        The *media_type* and *charset* arguments are used in the same
        manner as they are used in :meth:`static`.

        Arguments:
          content (str, bytes or int): Content to be sent as download or
            HTTP status code of the response to be returned.
          filename (str): Filename to use for saving the content
          media_type (str, optional): Media type of file.
          charset (str, optional): Character set of file.

        Returns:
          content, i.e. the first argument passed to this method.

        Raises:
          LogicError: When filename cannot be determined.
        """
        if isinstance(content, int) and content != 200:
            return content
        if filename is not None:
            filename = os.path.basename(filename)
        elif 'filename' in self.response.state:
            filename = self.response.state['filename']
        else:
            filename = os.path.basename(self.request.path)

        if filename == '':
            raise LogicError('Cannot determine filename for download')

        if media_type is not None:
            self.response.media_type = media_type
        else:
            self.response.media_type = mimetypes.guess_type(filename)[0]
        self.response.charset = charset
        self.response.add_header('Content-Disposition', 'attachment; '
                                 'filename="{}"'.format(filename))
        return content

    def __call__(self, environ, start_response):
        """Respond to an HTTP request.

        Arguments:
          environ (dict): Dictionary of environment variables
          start_response (callable): Callable to start HTTP response

        Returns:
          list: List containing a single sequence of bytes.
        """
        self.request = Request(environ)
        self.response = Response(start_response)

        route = self._router.resolve(self.request.method,
                                     self.request.path)
        if route is not None:
            callback, args, kwargs = route
            value = callback(*args, **kwargs)
        elif self._router.contains_method(self.request.method):
            value = 404 # Not found
        else:
            value = 501 # Not Implemented

        if isinstance(value, str) or isinstance(value, bytes):
            self.response.body = value

        elif isinstance(value, int) and value in Response._responses:
            self.response.status = value
            if self.response.body is None:
                self.response.body = self._get_error_page_callback()()

        elif (isinstance(value, tuple) and
              isinstance(value[0], int) and
              isinstance(value[1], str) and
              value[0] in Response._responses and
              300 <= value[0] <= 308):

            self.response.add_header('Location', value[1])
            self.response.status = value[0]
            if self.response.body is None:
                self.response.body = self._get_error_page_callback()()

        else:
            raise Error('Route callback for {} {} returned invalid '
                        'value: {}: {!r}'.format(self.request.method,
                        self.request.path, type(value).__name__, value))

        return self.response.response()

    def _get_error_page_callback(self):
        """Return an error page for the current response status."""
        if self.response.status in self._error_handlers:
            return self._error_handlers[self.response.status]
        elif None in self._error_handlers:
            return self._error_handlers[None]
        else:
            # Rudimentary error handler if no error handler was found
            self.response.media_type = 'text/plain'
            return lambda: self.response.status_line

class Router:

    """Route management and resolution."""

    def __init__(self):
        """Initialize router."""
        self._literal = collections.defaultdict(dict)
        self._wildcard = collections.defaultdict(list)
        self._regex = collections.defaultdict(list)

    def add(self, method, pattern, callback):
        """Add a route.

        Arguments:
          method (str): HTTP method, e.g. GET, POST, etc.
          pattern (str): Pattern that request paths must match.
          callback (str): Route handler that is invoked when a request
            path matches the *pattern*.
        """
        pat_type, pat = self._normalize_pattern(pattern)
        if pat_type == 'literal':
            self._literal[method][pat] = callback
        elif pat_type == 'wildcard':
            self._wildcard[method].append(WildcardRoute(pat, callback))
        else:
            self._regex[method].append(RegexRoute(pat, callback))

    def contains_method(self, method):
        """Check if there is at least one handler for *method*.

        Arguments:
          method (str): HTTP method name, e.g. GET, POST, etc.

        Returns:
          ``True`` if there is at least one route defined for *method*,
          ``False`` otherwise
        """
        return method in itertools.chain(self._literal, self._wildcard,
                                         self._regex)

    def resolve(self, method, path):
        """Resolve a request to a route handler.

        Arguments:
          method (str): HTTP method, e.g. GET, POST, etc. (type: str)
          path (str): Request path

        Returns:
          tuple or None: A tuple of three items:

            1. Route handler (callable)
            2. Positional arguments (list)
            3. Keyword arguments (dict)

          ``None`` if no route matches the request.
        """
        if method in self._literal and path in self._literal[method]:
            return self._literal[method][path], [], {}
        else:
            return self._resolve_non_literal_route(method, path)

    def _resolve_non_literal_route(self, method, path):
        """Resolve a request to a wildcard or regex route handler.

        Arguments:
          method (str): HTTP method name, e.g. GET, POST, etc.
          path (str): Request path

        Returns:
          tuple or None: A tuple of three items:

            1. Route handler (callable)
            2. Positional arguments (list)
            3. Keyword arguments (dict)

          ``None`` if no route matches the request.
        """
        for route_dict in (self._wildcard, self._regex):
            if method in route_dict:
                for route in reversed(route_dict[method]):
                    callback_data = route.match(path)
                    if callback_data is not None:
                        return callback_data
        return None

    @staticmethod
    def _normalize_pattern(pattern):
        """Return a normalized form of the pattern.

        Normalize the pattern by removing pattern type prefix if it
        exists in the pattern. Then return the pattern type and the
        pattern as a tuple of two strings.

        Arguments:
          pattern (str): Route pattern to match request paths

        Returns:
          tuple: Ruple of pattern type (str) and pattern (str)
        """
        if pattern.startswith('regex:'):
            pattern_type = 'regex'
            pattern = pattern[len('regex:'):]
        elif pattern.startswith('wildcard:'):
            pattern_type = 'wildcard'
            pattern = pattern[len('wildcard:'):]
        elif pattern.startswith('literal:'):
            pattern_type = 'literal'
            pattern = pattern[len('literal:'):]
        elif RegexRoute.like(pattern):
            pattern_type = 'regex'
        elif WildcardRoute.like(pattern):
            pattern_type = 'wildcard'
        else:
            pattern_type = 'literal'
        return pattern_type, pattern

class WildcardRoute:

    """Route containing wildcards to match request path."""

    _wildcard_re = re.compile(r'<[^<>/]*>')
    _tokenize_re = re.compile(r'<[^<>/]*>|[^<>/]+|/|<|>')

    def __init__(self, pattern, callback):
        """Initialize wildcard route.

        Arguments:
          pattern (str): Pattern associated with the route.
          callback (callable): Route handler.
        """
        self._re = []
        self._wildcards = []
        for token in WildcardRoute.tokens(pattern):
            if token and token.startswith('<') and token.endswith('>'):
                w = Wildcard(token)
                self._wildcards.append(w)
                self._re.append(w.regex())
            else:
                self._re.append(re.escape(token))
        self._re = re.compile('^' + ''.join(self._re) + '$')
        self._callback = callback

    def match(self, path):
        """Return route handler with arguments if path matches this route.

        Arguments:
          path (str): Request path

        Returns:
          tuple or None: A tuple of three items:

            1. Route handler (callable)
            2. Positional arguments (list)
            3. Keyword arguments (dict)

          ``None`` if the route does not match the path.
        """
        match = self._re.search(path)
        if match is None:
            return None
        args = []
        kwargs = {}
        for i, wildcard in enumerate(self._wildcards):
            if wildcard.name == '!':
                continue
            value = wildcard.value(match.groups()[i])
            if not wildcard.name:
                args.append(value)
            else:
                kwargs[wildcard.name] = value
        return self._callback, args, kwargs

    @staticmethod
    def like(pattern):
        """Determine if a pattern looks like a wildcard pattern.

        Arguments:
          pattern (str): Any route pattern.

        Returns:
          ``True`` if the specified pattern looks like a wildcard pattern,
          ``False`` otherwise.
        """
        return WildcardRoute._wildcard_re.search(pattern) is not None

    @staticmethod
    def tokens(pattern):
        """Return tokens in a pattern."""
        return WildcardRoute._tokenize_re.findall(pattern)

class Wildcard:

    """A single wildcard definition in a wildcard route pattern."""

    _types_re = {
        'str': r'([^/]+)',
        'path': r'(.+)',
        'int': r'(0|[1-9][0-9]*)',
        '+int': r'([1-9][0-9]*)',
        '-int': r'(0|-?[1-9][0-9]*)',
    }
    _name_re = re.compile(r'^(?:[^\d\W]\w*|!|)$') # Identifiers, '!', ''

    def __init__(self, spec):
        """Initialize wildcard definition.

        Arguments:
          spec (str): An angle-bracket delimited wildcard specification.
        """
        # Split '<foo:int>' into ['foo', 'int']
        tokens = spec[1:-1].split(':', 1)
        if len(tokens) == 1: # Split '<foo>' into ['foo', '']
            tokens.append('')
        self.name, self._type = tokens
        if not self._type:
            self._type = 'str'
        if Wildcard._name_re.search(self.name) is None:
            raise RouteError('Invalid wildcard name {!r} in {!r}'
                             .format(self.name, spec))
        if self._type not in Wildcard._types_re.keys():
            raise RouteError('Invalid wildcard type {!r} in {!r}'
                             .format(self._type, spec))

    def regex(self):
        """Convert the wildcard to a regular expression.

        Returns:
          str: A regular expression that matches strings that the
          wildcard is meant to match.
          """
        return Wildcard._types_re[self._type]

    def value(self, value):
        """Convert specified value to a value of wildcard type.

        This method does not check if the value matches the wildcard
        type. The caller of this method must ensure that the value
        passed to this method was obtained from a match by regular
        expression returned by the regex method of this class. Ensuring
        this guarantees that the value passed to the method matches the
        wildcard type.

        Arguments:
          value (str): Value to convert.

        Returns:
          str or int: Converted value.
        """
        return value if self._type in ['str', 'path'] else int(value)

class RegexRoute:

    """A regular expression pattern."""

    _group_re = re.compile(r'\(.*\)')

    def __init__(self, pattern, callback):
        """Initialize regular expression route.

        Arguments:
          pattern (str): Pattern associated with the route.
          callback (callable): Route handler.
        """
        self._re = re.compile(pattern)
        self._callback = callback

    def match(self, path):
        """Return route handler with arguments if path matches this route.

        Arguments:
          path (str): Request path

        Returns:
          tuple or None: A tuple of three items:

            1. Route handler (callable)
            2. Positional arguments (list)
            3. Keyword arguments (dict)

          ``None`` if the route does not match the path.
        """
        match = self._re.search(path)
        if match is None:
            return None
        kwargs_indexes = match.re.groupindex.values()
        args_indexes = [i for i in range(1, match.re.groups + 1)
                          if i not in kwargs_indexes]
        args = [match.group(i) for i in args_indexes]
        kwargs = {}
        for name, index in match.re.groupindex.items():
            kwargs[name] = match.group(index)
        return self._callback, args, kwargs

    @staticmethod
    def like(pattern):
        """Determine if a pattern looks like a regular expression.

        Arguments:
          pattern (str): Any route pattern.

        Returns:
          ``True`` if the specified pattern looks like a regex,
          ``False`` otherwise.
        """
        return RegexRoute._group_re.search(pattern) is not None

class Request:

    """Current request.

    Attributes:
      environ (dict): Dictionary of request environment variables.
      method (str): Request method.
      path (str): Request path.
      query (MultiDict): Key-value pairs from query string.
      form (MultiDict): Key-value pairs from form data in POST request.
      cookies (MultiDict): Key-value pairs from cookie string.
    """

    def __init__(self, environ):
        """Initialize the current request object.

        Arguments:
          environ (dict): Dictionary of environment variables.
        """
        self.environ = environ
        self.method = environ.get('REQUEST_METHOD', 'GET')
        self.path = environ.get('PATH_INFO', '/')
        if not self.path:
            self.path = '/'
        self.query = MultiDict()
        self.form = MultiDict()
        self.cookies = MultiDict()

        if 'QUERY_STRING' in environ:
            for k, v in urllib.parse.parse_qsl(environ['QUERY_STRING']):
                self.query[k] = v

        if 'wsgi.input' in environ:
            fs = cgi.FieldStorage(fp=environ['wsgi.input'],
                                  environ=environ)
            for k in fs:
                for v in fs.getlist(k):
                    self.form[k] = v

        if 'HTTP_COOKIE' in environ:
            cookies = http.cookies.SimpleCookie(environ['HTTP_COOKIE'])
            for c in cookies.values():
                self.cookies[c.key] = c.value

class Response:

    """Current response.

    Attributes:
      start (callable): Callable that starts response.
      status (int): HTTP response status code, defaults to 200.
      media_type (str): Media type of HTTP response, defaults to
        'text/html'. This together with :attr:`charset` determines the
        Content-Type response header.
      charset (str): Character set of HTTP response, defaults to
        'UTF-8'. This together with :attr:`media_type` determines the
        Content-Type response header.
      body (str or bytes): HTTP response body.
    """

    # Convert HTTP response status codes, phrases and detail in
    # http.server module into a dictionary of objects
    _Status = collections.namedtuple('_Status', ('phrase', 'detail'))

    _responses = {}

    for k, v in http.server.BaseHTTPRequestHandler.responses.items():
        _responses[k] = _Status(*v)
    del k, v

    def __init__(self, start_response_callable):
        """Initialize the current response object.

        Arguments:
          start_response_callable (callable): Callable that starts response.
        """
        self.start = start_response_callable
        self.status = 200
        self.media_type = 'text/html'
        self.charset = 'UTF-8'
        self._headers = []
        self.body = None
        self.state = {}

    def response(self):
        """Return the HTTP response body.

        Returns:
          bytes: HTTP response body as a sequence of bytes
        """
        if isinstance(self.body, bytes):
            out = self.body
        elif isinstance(self.body, str):
            out = self.body.encode(self.charset)
        else:
            out = b''
        self.add_header('Content-Type', self.content_type)
        self.add_header('Content-Length', str(len(out)))

        self.start(self.status_line, self._headers)
        return [out]

    def add_header(self, name, value):
        """Add an HTTP header to response object.

        Arguments:
          name (str): HTTP header field name
          value (str): HTTP header field value
        """
        if value is not None:
            self._headers.append((name, value))

    def set_cookie(self, name, value, attrs={}):
        """Add a Set-Cookie header to response object.

        For a description about cookie attribute values, see
        https://docs.python.org/3/library/http.cookies.html#http.cookies.Morsel.

        Arguments:
          name (str): Name of the cookie
          value (str): Value of the cookie
          attrs (dict): Dicitionary with cookie attribute keys and
                        values.
        """
        cookie = http.cookies.SimpleCookie()
        cookie[name] = value
        for key, value in attrs.items():
            cookie[name][key] = value
        self.add_header('Set-Cookie', cookie[name].OutputString())

    @property
    def status_line(self):
        """Return the HTTP response status line.

        The status line is determined from :attr:`status` code. For
        example, if the status code is 200, then '200 OK' is returned.

        Returns:
          str: Status line
        """
        return (str(self.status) + ' ' +
                Response._responses[self.status].phrase)

    @property
    def status_detail(self):
        """Return a description of the current HTTP response status.

        Returns:
          str: Response status description
        """
        return Response._responses[self.status].detail

    @property
    def content_type(self):
        """Return the value of Content-Type header field.

        The value for the Content-Type header field is determined from
        the :attr:`media_type` and :attr:`charset` data attributes.

        Returns:
          str: Value of Content-Type header field
        """
        if (self.media_type is not None and
            self.media_type.startswith('text/') and
            self.charset is not None):
            return self.media_type + '; charset=' + self.charset
        else:
            return self.media_type

class MultiDict(collections.UserDict):

    """Dictionary with multiple values for a key.

    Setting an existing key to a new value merely adds the value to the
    list of values for the key. Getting the value of an existing key
    returns the newest value set for the key.
    """

    def __setitem__(self, key, value):
        """Adds value to the list of values for the specified key.

        Arguments:
          key (object): Key
          value (object): Value
        """
        if key not in self.data:
            self.data[key] = [value]
        else:
            self.data[key].append(value)

    def __getitem__(self, key):
        """Return the newest value for the specified key.

        Arguments:
          key (object): Key

        Returns:
          object: Newest value for the specified key
        """
        return self.data[key][-1]

    def getall(self, key, default=[]):
        """Return the list of all values for the specified key.

        Arguments:
          key (object): Key
          default (list): Default value to return if the key does not
            exist, defaults to ``[]``, i.e. an empty list.

        Returns:
          list: List of all values for the specified key if the key
          exists, ``default`` otherwise.
        """
        return self.data[key] if key in self.data else default

class Error(Exception):
    """Base class for exceptions."""

class RouteError(Error):
    """Route related exception."""

class LogicError(Error):
    """Logical error that can be avoided by careful coding."""


'''rmtree, can remove directories even when shutil.rmtree fails
because of read-only files in NFS and Windows.

.. This software is released under an MIT/X11 open source license.
   Copyright 2014 Diffeo, Inc.
'''
import os
import sys
import logging
import subprocess
import shutil
import time
import traceback

logger = logging.getLogger(__name__)

def rmtree(path, use_shutil=True, followlinks=False, retries=10):
    '''remove all files and directories below path, including path
    itself; works even when shutil.rmtree fails because of read-only
    files in NFS and Windows.  Follows symlinks.

    `use_shutil` defaults to True; useful for testing

    `followlinks` defaults to False; if set to True, shutil.rmtree is
    not used.
    '''
    if use_shutil and not followlinks:
        try:
            shutil.rmtree(path)
            return
        except Exception, exc:
            logger.info('shutil.rmtree(%s) failed, so resorting to recursive delete', path)
            logger.debug('\ntrapped:\n%s', traceback.format_exc(exc))

    if not os.path.isdir(path):
        os.remove(path)
        return

    ## bottom up traversal removing files and then removing directories
    for root, dir_names, file_names in os.walk(path, topdown=False, followlinks=followlinks):
        for fname in file_names:
            fpath = os.path.join(root, fname)
            tries = 0
            while tries < retries:
                tries += 1
                try:
                    os.remove(fpath)
                    break
                except Exception, exc:
                    time.sleep(0.1)
            if os.path.exists(fpath):
                logger.critical('os.remove(%s) failed, so leaving data behind!!!', fpath)
                logger.critical('\ntrapped:\n%s', traceback.format_exc(exc))
                #logger.critical(get_open_fds())
        for dname in dir_names:
            full_path = os.path.join(root, dname)
            if os.path.islink(full_path):
                real_path = os.path.realpath(full_path)
                os.remove(full_path)
                full_path = real_path
            os.rmdir(full_path)

    if os.path.exists(path):
        os.rmdir(path)

def get_open_fds(verbose=False):
    '''return list of open files for current process

    .. warning: will only work on UNIX-like os-es.
    '''
    pid = os.getpid()
    procs = subprocess.check_output(
        [ "lsof", '-w', '-Ff', "-p", str( pid ) ] )
    if verbose:
        oprocs = subprocess.check_output(
            [ "lsof", '-w', "-p", str( pid ) ] )
        logger.info(oprocs)
    open_files = filter(
        lambda s: s and s[ 0 ] == 'f' and s[1: ].isdigit(),
        procs.split( '\n' ) )
    return open_files


#!/usr/bin/env python
'''
streamcorpus_pipeline Transform for gathering stats about file types.

This software is released under an MIT/X11 open source license.

Copyright 2012-2014 Diffeo, Inc.
'''
from __future__ import absolute_import
import re

from streamcorpus_pipeline.stages import Configured

first_letters = re.compile('''^(.|\n)*?(?P<first_letters>\S+\s?\S+)''')

## ignore up to 20 non '<' characters including line feeds looking for
## doctype XXXX, where the most common is XXX=HTML
#doctype_something_m = re.compile('''^([^<]|\r|\n){,20}\<\!doctype\s+(?P<doctype>\S+)''', re.I)
doctype_re = re.compile(
    ## skip up to 100 null bytes or whitespaces
    '''(?P<chunk>(\x00|\s|\r|\n)*?)''' + \
    ## allow an optional XML preamble
    '''(?P<xml_preamble>(\<\?xml([^>]|\n)*?\>(\s|\r|\n){,100})?)''' + \
    ## capture the doctype
    '''\<\!doctype\s+(?P<doctype>[a-zA-Z0-9]{,10})''',
    re.I)

probably_html = re.compile(
    ## skip up to twenty bytes of stuff and up to 1000 whitespaces
    ## from start of file
    #'''^(?P<chunk>([^<]|\r|\n){,20}(\s|\r|\n){,1000})''' + \
    '''[^<]*''' + \
    ## allow an optional XML preamble
    '''(?P<xml_preamble>(\<\?xml([^>]|\n)*?\>(\s|\r|\n){,1000})?)''' + \
    ## allow an optional DOCTYPE preamble without a doctype
    '''(\<\!doctype\s*\>(\s|\r|\n){,1000})?''' + \
    ## look for a first tag that is very likely to HTML
    '''(?P<first_tag>\<(html|head|frame|body|iframe|title|table|br|style|base\s+href|script|meta|link|div|\!--))''',
    re.I)

xml_ish = re.compile('''^([^<]|\r|\n){,20}(?P<intro>\<\?xml[^>]*?\>(.|\n){50})''', re.I)

pdf_start = re.compile('''^(\s|\r|\n){,20}\%pdf-(?P<version>.{5})''', re.I)

#import sys

tags = 'html|head|frame|body|iframe|title|table|br|style|base\s+href|script|meta|link|div'
def has_tags(text):
    text = text.lower()
    for tag in tags.split('|'):
        if tag in text:
            return True
    return False

def file_type_stats(config):
    '''
    returns a kba.pipeline "transform" function that generates file
    type stats from the stream_items that it sees.  Currently, these
    stats are just the first five non-whitespace characters.
    '''
    ## make a closure around config
    def _file_type_stats(stream_item, context):
        if stream_item.body and stream_item.body.raw:
            #print repr(stream_item.body.raw[:250])
            #sys.stdout.flush()
            #doctype_m = doctype_re.match(stream_item.body.raw[:250])
            #if doctype_m:
                #print 'DOCTYPE: %s' % repr(doctype_m.group('doctype').lower())
            if 'doctype html' in stream_item.body.raw[:250].lower():
                print 'DOCTYPE: html'
            else:
                #if probably_html.search(stream_item.body.raw):
                if has_tags(stream_item.body.raw[:400]):
                    print 'PROBABLY_HTML'
                else:
                    xml = xml_ish.search(stream_item.body.raw)
                    if xml:
                        print 'XML: %s' % repr(xml.group('intro'))
                    else:
                        pdf = pdf_start.search(stream_item.body.raw)
                        if pdf:
                            print 'PDF %s' % repr(pdf.group('version'))
                        else:
                            ext = stream_item.abs_url.split('.')[-1]
                            if len(ext) < 6:
                                print 'UNK ext: %s' % repr(ext)
                            else:
                                first = first_letters.match(stream_item.body.raw)
                                if first and False:
                                    print 'UNK letters: %s' % repr(first.group('first_letters'))
                                else:
                                    print 'UNK first bytes: %s' % repr(stream_item.body.raw[:50])
                    #m = first_three_letters.search(stream_item.body.raw)
                    #if m:
                    #    print repr(m.group('first_three_letters')).lower().strip()
                    #else:
                    #    print repr(stream_item.body.raw[:50]).lower().strip()
        return stream_item

    return _file_type_stats

class guess_media_type(Configured):
    '''
    returns a kba.pipeline "transform" function that populates
    body.media_type if it is empty and the content type is easily
    guessed.
    '''
    config_name = 'guess_media_type'
    def __call__(self, stream_item, context):
        if stream_item.body and stream_item.body.media_type:
            ## don't change it
            return stream_item

        if stream_item.body and stream_item.body.raw:
            if len(stream_item.body.raw) < 40:
                stream_item.body.raw = None
                stream_item.body.media_type = None

            elif 'doctype html' in stream_item.body.raw[:250].lower():
                stream_item.body.media_type = 'text/html'

            elif has_tags( stream_item.body.raw[:400] ):
                stream_item.body.media_type = 'text/html'

            elif stream_item.body.raw[:10].lower().startswith('%pdf-'):
                stream_item.body.media_type = 'application/pdf'

        if stream_item.body and not stream_item.body.media_type \
                and self.config.get('fallback_media_type'):
            stream_item.body.media_type = self.config['fallback_media_type']

        return stream_item


'''rejester task for iterating over the KBA 2014 corpus and saving
stats files

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2014 Diffeo, Inc.
'''

from __future__ import division
import gzip
import json
import logging
import os
import re
from streamcorpus import Chunk
from subprocess import Popen, PIPE
import sys
import time
import traceback
import yaml
import uuid

logging.basicConfig()

work_spec = {
    'name': 'verify_kba2014',
    'desc': 'Verify chunks from the KBA 2014 streamcorpus.',
    'min_gb': 1,
    'config': None,
    'module': 'verify_kba2014',
    'run_function': 'rejester_run',
    'terminate_function': 'rejester_terminate',
}

def rejester_run(work_unit):
     '''get a rejester.WorkUnit with KBA s3 path, fetch it, and save
     some counts about it.
     '''
     #fname = 'verify-chunks-%d-%d' % (os.getpid(), time.time())
     fname = work_unit.key.strip().split('/')[-1]

     output_dir_path = work_unit.data.get('output_dir_path', '/mnt')
     u = uuid.uuid3(uuid.UUID(int=0), work_unit.key.strip())
     path1 = u.hex[0]
     path2 = u.hex[1]
     fpath = os.path.join(output_dir_path, path1, path2, fname)
     if not os.path.exists(os.path.dirname(fpath)):
          os.makedirs(os.path.dirname(fpath))

     output = gzip.open(fpath + '-out.gz', 'wb')

     expected_si_count = int(fname.split('-')[1])

     max_tries = 20
     tries = 0
     while tries < max_tries:
          try:
               exc, si_count, serif_count, clean_visible_bytes, clean_visible_count, stream_ids = \
                   attempt_fetch(work_unit, fpath)
               if si_count != expected_si_count:
                    print 'retrying because si_count = %d != %d expected_si_count' % (si_count, expected_si_count)
                    sys.stdout.flush()
                    tries += 1
                    continue
               else:
                    print 'succeeded in reading si_count = %d' % (si_count,)
                    sys.stdout.flush()
               output.write( '%s\t%d\t%d\t%d\t%d\t%s\t%s\n' % (
                         exc, si_count, serif_count, clean_visible_bytes, clean_visible_count,
                         work_unit.key.strip(), ','.join(['%s|%s' % tup for tup in stream_ids])) )
               break
          except Exception, exc:
               print 'broken?'
               print traceback.format_exc(exc)
               sys.stdout.flush()
               tries += 1
               output.write(traceback.format_exc(exc))

     output.close()

def attempt_fetch(work_unit, fpath):
     '''attempt a fetch and iteration over a work_unit.key path in s3
     '''
     url = 'http://s3.amazonaws.com/aws-publicdatasets/' + work_unit.key.strip()

     ## cheapest way to iterate over the corpus is a few stages of
     ## streamed child processes.  Note that stderr needs to go
     ## separately to a file so that reading the stdin doesn't get
     ## blocked:
     cmd = '(wget -O - %s | gpg --no-permission-warning --trust-model always --output - --decrypt - | xz --decompress) 2> %s-err' % (url, fpath)
     print cmd
     child = Popen(cmd, stdout=PIPE, shell=True)
     print 'child launched'
     sys.stdout.flush()

     si_count = 0
     serif_count = 0
     exc = ''
     stream_ids = list()
     clean_visible_bytes = 0
     clean_visible_count = 0
     try:
         for si in Chunk(file_obj=child.stdout):
             print si.stream_id, si.abs_url
             if si.body.language:
                 lang = si.body.language.code
             else:
                 lang = ''
             stream_ids.append((lang, si.stream_id))
             if si.body.clean_visible:
                 clean_visible_count += 1
                 clean_visible_bytes += len(si.body.clean_visible)
             si_count += 1
             if 'serif' in si.body.sentences:
                 serif_count += 1
     except Exception, exc:
         exc = re.sub('\s+', ' ', str(exc)).strip()

     child.terminate()
     child.wait()
     child.stdout.close()
     return exc, si_count, serif_count, clean_visible_bytes, clean_visible_count, stream_ids

def rejester_terminate(work_unit):
     'required for rejester, but no-op here'
     pass

if __name__ == '__main__':
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument('--run', action='store_true', default=False, help='take S3 paths over stdin')
     args = parser.parse_args()

     if args.run:
          print 'generating simulated WorkUnits for processing the corpus'
          class Empty():
               pass
          for line in sys.stdin:
               print 'running', line
               sys.stdout.flush()
               work_unit = Empty()
               work_unit.data = dict(output_dir_path = '/mnt')
               work_unit.key = line.strip()
               rejester_run(work_unit)

     else:
          ws_path = 'verify_kba2014_work_spec.yaml'
          yaml.dump(work_spec, open(ws_path, 'wb'))
          print 'dumped a WorkSpec for use with rejester to %s' % ws_path

          wus_path = 'verify_kba2014_work_units.json'
          wus = open(wus_path, 'wb')
          count = 0
          for line in sys.stdin:
               count += 1
               wu = {line.strip(): dict(output_dir_path='/mnt')}
               wus.write(json.dumps(wu) + '\n')
          wus.close()

          print 'created %d records in %s' % (count, wus_path)


# coding=utf-8
"""
Print the name of a random adjective/species, more or less…
"""

import argparse
from itertools import chain
from os import path
from random import choice
import sys

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--maxlen', type=int, default=8,
                    help="Maximum length for the name, excluding any separator. (default=8)")
parser.add_argument('--sep', type=str, default='', metavar="SEPARATOR",
                    help="Separator between the adjective and species words. (default='')")
parser.add_argument('--count', type=int, default=1,
                    help="Number of adjective/species combinations to print.")
parser.add_argument('--prevent-stutter', action='store_true',
                    help="Prevent the same letter from appearing on an adjective/species boundary. "
                    "(default=True)")

def get_file_lines(file_name):
    """Return a list of non-empty lines from `file_path`."""
    file_path = path.join(path.dirname(path.abspath(__file__)), file_name)
    with open(file_path) as file_obj:
        return [line for line in file_obj.read().splitlines() if line]

def random_species():
    """Return the name of a species at random."""
    return choice(get_file_lines('species.txt'))

def get_describers():
    """
    Return a describer tuple in the form `(name, position)`,
    where position is either 'prefix' or 'suffix'.
    """
    adjectives = map(lambda x: (x, 'prefix'), get_file_lines('adjectives.txt'))
    animal_nouns = map(lambda x: (x, 'suffix'), get_file_lines('nouns.txt'))
    return list(chain(adjectives, animal_nouns))

def random_describer():
    """Return a 2-tuple, matching a describer to 'prefix' or 'suffix'."""
    return choice(get_describers())

def _random_adjspecies_pair():
    """Return an ordered 2-tuple containing a species and a describer."""
    describer, desc_position = random_describer()
    if desc_position == 'prefix':
        return (describer, random_species())
    elif desc_position == 'suffix':
        return (random_species(), describer)

def random_adjspecies_pair(maxlen=None, prevent_stutter=True):
    """
    Return an ordered 2-tuple containing a species and a describer.
    The letter-count of the pair is guarantee to not exceed `maxlen` if
    it is given. If `prevent_stutter` is True, the last letter of the
    first item of the pair will be different from the first letter of
    the second item.
    """
    while True:
        pair = _random_adjspecies_pair()
        if maxlen and len(''.join(pair)) > maxlen:
            continue
        if prevent_stutter and pair[0][-1] == pair[1][0]:
            continue
        return pair

def random_adjspecies(sep='', maxlen=8, prevent_stutter=True):
    """
    Return a random adjective/species, separated by `sep`. The keyword
    arguments `maxlen` and `prevent_stutter` are the same as for
    `random_adjspecies_pair`, but note that the maximum length argument is
    not affected by the separator.
    """
    pair = random_adjspecies_pair(maxlen, prevent_stutter)
    return pair[0] + sep + pair[1]

def main(*argv):
    args = parser.parse_args()
    for count in range(args.count):
        print(random_adjspecies(args.sep, args.maxlen))

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


# -*- coding: utf-8 -*-
"""
    Command line tools for Goo labs API
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    :author: tell-k <ffk2005@gmail.com>
    :copyright: tell-k. All Rights Reserved.
"""
from __future__ import division, print_function, absolute_import  # NOQA

import json
import locale

import click
import six

import goolabs
from goolabs import GoolabsAPI

if 0:
    from typing import Optional, IO, List, Dict, Any  # NOQA
    from click.core import Context  # NOQA

def text(s):
    # type: (unicode) -> unicode
    if isinstance(s, six.binary_type):
        return s.decode(locale.getpreferredencoding())
    return s

def clean_app_id(app_id):
    # type: (unicode) -> unicode
    if not app_id:
        raise click.UsageError('Missing option "--app-id" / "-a" '
                               'or GOOLABS_APP_ID enviroment value.')
    return app_id

def clean_sentence(sentence, sentence_file):
    # type: (unicode, Optional[IO]) -> unicode
    if not sentence and not sentence_file:
        raise click.UsageError('Missing sentence. You must set'
                               ' SENTENCE argument or --file option.')
    if not sentence and sentence_file:
        sentence = text(sentence_file.read())
    return sentence

def clean_review(review, review_file):
    # type: (unicode, Optional[IO]) -> List[unicode]
    if not review and not review_file:
        raise click.UsageError('Missing review. You must set'
                               ' REVIEW argument or --file option.')
    if not review and review_file:
        review = text(review_file.read())
    return review.split('\n')

def clean_body(body, body_file):
    # type: (unicode, Optional[IO]) -> unicode
    if not body and not body_file:
        raise click.UsageError('Missing body. You must set'
                               ' BODY argument or --file option.')
    if not body and body_file:
        body = text(body_file.read())
    return body

def clean_length(length):
    # type: (unicode) -> Optional[int]
    if length is None:
        return None
    try:
        return int(length)
    except ValueError:
        raise click.UsageError(
            '--length is not Integer. You must choice length from 60/120/180.'
        )

def format_json(json_data):
    # type: (Dict[unicode,Any]) -> unicode
    return json.dumps(json_data, indent=2, ensure_ascii=False)

@click.group()
@click.pass_context
@click.version_option(version=goolabs.__version__)
def main(ctx):
    # type: (Context) -> None
    """ Command line tools for Goo labs API(https://labs.goo.ne.jp/api/). """

@main.command()
@click.argument('sentence', required=False, type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--info-filter', '-i', 'info_filter',
              type=text, help='form,pos,read')
@click.option('--pos-filter', '-p', 'pos_filter',
              type=text, help=u'名刺,動詞活用語尾,句点..etc')
@click.option('--file', '-f', 'sentence_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def morph(ctx, app_id, sentence_file, json_flag,
          sentence, info_filter, pos_filter, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, unicode, unicode) -> None  # NOQA
    """ Morphological analysis for Japanese."""

    app_id = clean_app_id(app_id)
    sentence = clean_sentence(sentence, sentence_file)

    if info_filter:
        info_filter = info_filter.replace(',', '|')

    if pos_filter:
        pos_filter = pos_filter.replace(',', '|')

    api = GoolabsAPI(app_id)
    ret = api.morph(
        sentence=sentence,
        info_filter=info_filter,
        pos_filter=pos_filter,
        request_id=request_id,
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    for words in ret['word_list']:
        for word in words:
            click.echo(','.join(word))

@main.command()
@click.argument('query_pair', required=True, nargs=2, type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def similarity(ctx, app_id, json_flag, query_pair, request_id):
    # type: (Context, unicode, bool, List[unicode], unicode) -> None
    """ Scoring the similarity of two words. """

    app_id = clean_app_id(app_id)

    api = GoolabsAPI(app_id)
    ret = api.similarity(
        query_pair=query_pair,
        request_id=request_id
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    click.echo('{0:.16f}'.format(ret['score']))

@main.command()
@click.argument('sentence', required=False, type=text)
@click.option('--output-type', '-o', 'output_type', default='hiragana',
              type=click.Choice(['hiragana', 'katakana']))
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--file', '-f', 'sentence_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def hiragana(ctx, app_id, sentence_file,
             json_flag, sentence, output_type, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, unicode) -> None # NOQA
    """ Convert the Japanese to Hiragana or Katakana. """

    app_id = clean_app_id(app_id)
    sentence = clean_sentence(sentence, sentence_file)

    api = GoolabsAPI(app_id)
    ret = api.hiragana(
        sentence=sentence,
        output_type=output_type,
        request_id=request_id
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    click.echo(ret['converted'])

@main.command()
@click.argument('sentence', required=False, type=text)
@click.option('--class-filter', '-c', 'class_filter',
              help='ART,ORG,PSN,LOC,DAT', type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--file', '-f', 'sentence_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def entity(ctx, app_id, sentence_file,
           json_flag, sentence, class_filter, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, unicode) -> None # NOQA
    """ Extract unique representation from sentence. """

    app_id = clean_app_id(app_id)
    sentence = clean_sentence(sentence, sentence_file)

    if class_filter:
        class_filter = class_filter.replace(',', '|')

    api = GoolabsAPI(app_id)
    ret = api.entity(
        sentence=sentence,
        class_filter=class_filter,
        request_id=request_id
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    for ne in ret['ne_list']:
        click.echo(','.join(ne))

@main.command()
@click.argument('review', required=False, type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--length', '-l', type=click.Choice(['60', '120', '180']))
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--file', '-f', 'review_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def shortsum(ctx, app_id, review_file,
             json_flag, review, length, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, unicode) -> None # NOQA
    """Summarize reviews into a short summary."""

    app_id = clean_app_id(app_id)
    review_list = clean_review(review, review_file)
    length_int = clean_length(length)  # type: Optional[int]

    api = GoolabsAPI(app_id)
    ret = api.shortsum(
        review_list=review_list,
        length=length_int,
        request_id=request_id,
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    click.echo(ret['summary'])

@main.command()
@click.argument('title', required=True, type=text)
@click.argument('body', required=False, type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--max_num', '-m', type=click.INT)
@click.option('--forcus', '-fo', type=click.Choice(['ORG', 'PSN', 'LOC']))
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--file', '-f', 'body_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def keyword(ctx, app_id, body_file, json_flag,
            title, body, max_num, forcus, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, int, unicode, unicode) -> None # NOQA
    """Extract "keywords" from an input document. """

    app_id = clean_app_id(app_id)
    body = clean_body(body, body_file)

    api = GoolabsAPI(app_id)
    ret = api.keyword(
        title=title,
        body=body,
        max_num=max_num,
        forcus=forcus,
        request_id=request_id,
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    for k in ret['keywords']:
        k = dict((key.encode('utf-8'), k[key]) for key in k.keys())
        for keyword, score in six.iteritems(k):
            click.echo(u'{0},{1}'.format(text(keyword), score))

@main.command()
@click.argument('sentence', required=False, type=text)
@click.option('--app-id', '-a', 'app_id', envvar='GOOLABS_APP_ID', type=text)
@click.option('--request-id', '-r', 'request_id', type=text)
@click.option('--doc-time', '-d', 'doc_time', type=text)
@click.option('--file', '-f', 'sentence_file', type=click.File('rb'))
@click.option('--json/--no-json', '-j', 'json_flag', default=False)
@click.pass_context
def chrono(ctx, app_id, sentence_file,
           json_flag, sentence, doc_time, request_id):
    # type: (Context, unicode, Optional[IO], bool, unicode, unicode, unicode) -> None  # NOQA
    """Extract expression expressing date and time and normalize its value """

    app_id = clean_app_id(app_id)
    sentence = clean_sentence(sentence, sentence_file)

    api = GoolabsAPI(app_id)
    ret = api.chrono(
        sentence=sentence,
        doc_time=doc_time,
        request_id=request_id,
    )

    if json_flag:
        click.echo(format_json(api.response.json()))
        return

    for pair in ret['datetime_list']:
        click.echo(u'{0}: {1}'.format(text(pair[0]), pair[1]))


#!/usr/bin/env python
'''Configuration and execution of the actual pipeline.

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2015 Diffeo, Inc.

The :class:`Pipeline` itself consists of a series of
:mod:`~streamcorpus_pipeline.stages`.  These are broken into several
categories:

* Exactly one *reader* runs first, producing a sequence of
  :class:`streamcorpus.StreamItem` objects.

* The stream items are fed through *incremental transforms*, which take
  one stream item in and produce one stream item out.

* All of the stream items are written to a file, and *batch transforms*
  can operate on the entire collection of stream items at once.

* *Post-batch incremental transforms* again operate on individual
  stream items.

* Some number of *writers* send the final output somewhere.

Configuration
=============

The :command:`streamcorpus_pipeline` tool expects configuration in
a YAML file.  The configuration resulting from this can be passed
through :class:`PipelineFactory` to create :class:`Pipeline` objects.
A typical coniguration looks like:

.. code-block:: yaml

    streamcorpus_pipeline:
      # Lists of stages
      reader: from_local_chunks
      incremental_transforms: [clean_html, language]
      batch_transforms: []
      post_batch_incremental_transforms: []
      # to_local_chunks must be the last writer if it is used
      writers: [to_local_chunks]

      # Configuration for specific stages
      clean_html:
        include_language_codes: [en]

The ``streamcorpus_pipeline`` block can additionally be configured
with:

.. code-block:: yaml

    root_path: /home/user/diffeo

Any configuration variable whose name ends in "path" whose value is
not an absolute path is considered relative to this directory.  If
omitted, use the current directory.

.. code-block:: yaml

    tmp_dir_path: directory

Intermediate files are stored in a subdirectory of :file:`directory`.
The pipeline execution will make an effort to clean this up.  If
omitted, defaults to :file:`/tmp`.

.. code-block:: yaml

    third_dir_path: directory

External packages such as NLP taggers are stored in subdirectories
of :file:`directory`; these are typically individually configured
with `path_in_third` options relative to this directory.

.. code-block:: yaml

    rate_log_interval: 500

When this many items have been processed, log an INFO-level log
message giving the current progress.

.. code-block:: yaml

    input_item_limit: 500

Stop processing after reading this many stream items.  (Default:
process the entire stream)

.. code-block:: yaml

    cleanup_tmp_files: true

After execution finishes, delete the per-execution subdirectory of
``tmp_dir_path``.

.. code-block:: yaml

    assert_single_source: false

Normally a set of stream items has a consistent
:attr:`streamcorpus.StreamItem.source` value, and the pipeline
framework will stop if it sees different values here.  Set this to
``false`` to disable this check.  (Default: do assert a single source
value)

.. code-block:: yaml

    output_chunk_max_count: 500

After this many items have been written, close and re-open the output.
(Default: 500 items)

.. code-block:: yaml

    output_max_clean_visible_bytes: 1000000

After this much :attr:`streamcorpus.StreamItem.clean_visible` content
has been written, close and re-open the output.  (Default: write
entire output in one batch)

.. code-block:: yaml

    external_stages_path: stages.py

The file :file:`stages.py` is a Python module that declares a
top-level dictionary named `Stages`, a map from stage name to
implementing class.  Stages defined in this file can be used in any of
the appropriate stage lists.

.. code-block:: yaml

    external_stages_modules: [ example.stages ]

The Python module :mod:`example.stages` declares a top-level
dictionary named `Stages`, a map from stage name to implementing
class.  The named modules must be on :data:`sys.path` so that the
Python interpreter can find it.  Stages defined in these modules can
be used in any of the appropriate stage lists.

API
===

The standard execution path is to pass the
:mod:`streamcorpus_pipeline` module to
:func:`yakonfig.parse_args`, then use :class:`PipelineFactory` to
create :class:`Pipeline` objects from the resulting configuration.

.. todo:: Make the top-level configuration point
          :class:`PipelineFactory`, not the
          :mod:`streamcorpus_pipeline` module

.. autoclass:: PipelineFactory
   :members:

.. autoclass:: Pipeline
   :members:

'''

from __future__ import absolute_import
import logging
import os
import threading
import time
import uuid

try:
    import gevent
except ImportError:
    gevent = None

import streamcorpus
from streamcorpus_pipeline._exceptions import TransformGivingUp, \
    InvalidStreamItem
from streamcorpus_pipeline.util import rmtree

logger = logging.getLogger(__name__)

class PipelineFactory(object):
    '''Factory to create :class:`Pipeline` objects from configuration.

    Call this to get a :class:`Pipeline` object.  Typical programmatic
    use:

    .. code-block:: python

       parser = argparse.ArgumentParser()
       args = yakonfig.parse_args([yakonfig, streamcorpus_pipeline])
       factory = PipelineFactory(StageRegistry())
       pipeline = factory(yakonfig.get_global_config('streamcorpus_pipeline'))

    This factory class will instantiate all of the stages named in the
    `streamcorpus_pipeline` configuration.  These stages will be created
    with their corresponding configuration, except that they have two
    keys added, ``tmp_dir_path`` and ``third_dir_path``, from the
    top-level configuration.

    .. automethod:: __init__
    .. automethod:: __call__

    .. attribute:: registry

       The :class:`streamcorpus_pipeline.stages.StageRegistry` used
       to find pipeline stages.

    .. attribute:: tmp_dir_suffix

       A string value that is appended to ``tmp_dir_path`` when
       creating pipeline stages.  If :const:`None`, use the top-level
       ``tmp_dir_path`` configuration directly.

    .. attribute:: lock

       A :class:`threading.Lock` to protect against concurrent
       modification of `tmp_dir_suffix`.

    '''

    def __init__(self, registry):
        '''Create a pipeline factory.

        :param dict config: top-level "streamcorpus_pipeline" configuration
        :param registry: registry of stages
        :type registry: :class:`~streamcorpus_pipeline.stages.StageRegistry`

        '''
        super(PipelineFactory, self).__init__()
        self.registry = registry
        self.lock = threading.Lock()
        self.tmp_dir_suffix = None

    def create(self, stage, scp_config, config=None):
        '''Create a pipeline stage.

        Instantiates `stage` with `config`.  This essentially
        translates to ``stage(config)``, except that two keys from
        `scp_config` are injected into the configuration:
        ``tmp_dir_path`` is an execution-specific directory from
        combining the top-level ``tmp_dir_path`` configuration with
        :attr:`tmp_dir_suffix`; and ``third_dir_path`` is the same
        path from the top-level configuration.  `stage` may be either
        a callable returning the stage (e.g. its class), or its name
        in the configuration.

        `scp_config` is the configuration for the pipeline as a
        whole, and is required.  `config` is the configuration for
        the stage; if it is :const:`None` then it is extracted
        from `scp_config`.

        If you already have a fully formed configuration block
        and want to create a stage, you can call

        .. code-block:: python

            factory.registry[stage](stage_config)

        In most cases if you have a stage class object and want to
        instantiate it with its defaults you can call

        .. code-block:: python

            stage = stage_cls(stage_cls.default_config)

        .. note:: This mirrors
                  :meth:`yakonfig.factory.AutoFactory.create`, with
                  some thought that this factory class might migrate
                  to using that as a base in the future.

        :param stage: pipeline stage class, or its name in the registry
        :param dict scp_config: configuration block for the pipeline
        :param dict config: configuration block for the stage, or
          :const:`None` to get it from `scp_config`

        '''
        # Figure out what we have for a stage and its name
        if isinstance(stage, basestring):
            stage_name = stage
            stage_obj = self.registry[stage_name]
        else:
            stage_name = getattr(stage, 'config_name', stage.__name__)
            stage_obj = stage

        # Find the configuration; get a copy we can mutate
        if config is None:
            config = scp_config.get(stage_name, None)
        if config is None:
            config = getattr(stage_obj, 'default_config', {})
        config = dict(config)

        # Fill in more values
        if self.tmp_dir_suffix is None:
            config['tmp_dir_path'] = scp_config['tmp_dir_path']
        else:
            config['tmp_dir_path'] = os.path.join(scp_config['tmp_dir_path'],
                                                  self.tmp_dir_suffix)
        config['third_dir_path'] = scp_config['third_dir_path']

        return stage_obj(config)

    def _init_stage(self, config, name):

        '''Create a single indirect stage.

        `name` should be the name of a config item that holds the
        name of a stage, for instance, ``reader``.  This looks up
        the name of that stage, then creates and returns the
        stage named.  For instance, if the config says

        .. code-block:: yaml

            reader: from_local_chunks

        then calling ``self._init_stage(scp_config, 'reader')`` will
        return a new instance of the
        :class:`~streamcorpus_pipeline._local_storage.from_local_chunks`
        stage.

        :param dict config: `streamcorpus_pipeline` configuration block
        :param str name: name of stage name entry
        :return: new instance of the stage

        '''
        return self.create(config[name], config)

    def _init_stages(self, config, name):
        '''Create a list of indirect stages.

        `name` should be the name of a config item that holds a list
        of names of stages, for instance, ``writers``.  This looks up
        the names of those stages, then creates and returns the
        corresponding list of stage objects.  For instance, if the
        config says

        .. code-block:: yaml

            incremental_transforms: [clean_html, clean_visible]

        then calling ``self._init_stages(scp_config,
        'incremental_transforms')`` will return a list of the two
        named stage instances.

        :param dict config: `streamcorpus_pipeline` configuration block
        :param str name: name of the stage name list entry
        :return: list of new stage instances

        '''
        if name not in config:
            return []
        return [self.create(stage, config) for stage in config[name]]

    def _init_all_stages(self, config):
        '''Create stages that are used for the pipeline.

        :param dict config: `streamcorpus_pipeline` configuration
        :return: tuple of (reader, incremental transforms, batch
          transforms, post-batch incremental transforms, writers,
          temporary directory)

        '''
        reader = self._init_stage(config, 'reader')
        incremental_transforms = self._init_stages(
            config, 'incremental_transforms')
        batch_transforms = self._init_stages(config, 'batch_transforms')
        post_batch_incremental_transforms = self._init_stages(
            config, 'post_batch_incremental_transforms')
        writers = self._init_stages(config, 'writers')
        tmp_dir_path = os.path.join(config['tmp_dir_path'],
                                    self.tmp_dir_suffix)
        return (reader, incremental_transforms, batch_transforms,
                post_batch_incremental_transforms, writers, tmp_dir_path)

    def __call__(self, config):
        '''Create a :class:`Pipeline`.

        Pass in the configuration under the ``streamcorpus_pipeline``
        block, not the top-level configuration that contains it.

        If :attr:`tmp_dir_suffix` is :const:`None`, then locks the
        factory and creates stages with a temporary (UUID) value.  If
        the configuration has `cleanup_tmp_files` set to :const:`True`
        (the default) then executing the resulting pipeline will clean
        up the directory afterwards.

        :param dict config: `streamcorpus_pipeline` configuration
        :return: new pipeline instance

        '''
        tmp_dir_suffix = self.tmp_dir_suffix
        if tmp_dir_suffix is None:
            with self.lock:
                self.tmp_dir_suffix = str(uuid.uuid4())
                try:
                    (reader, incremental_transforms, batch_transforms,
                     pbi_transforms, writers,
                     tmp_dir_path) = self._init_all_stages(config)
                finally:
                    self.tmp_dir_suffix = None
        else:
            (reader, incremental_transforms, batch_transforms,
             pbi_transforms, writers,
             tmp_dir_path) = self._init_all_stages(config)

        return Pipeline(
            rate_log_interval=config['rate_log_interval'],
            input_item_limit=config.get('input_item_limit'),
            cleanup_tmp_files=config['cleanup_tmp_files'],
            tmp_dir_path=tmp_dir_path,
            assert_single_source=config['assert_single_source'],
            output_chunk_max_count=config.get('output_chunk_max_count'),
            output_max_clean_visible_bytes=config.get(
                'output_max_clean_visible_bytes'),
            reader=reader,
            incremental_transforms=incremental_transforms,
            batch_transforms=batch_transforms,
            post_batch_incremental_transforms=pbi_transforms,
            writers=writers,
        )

class Pipeline(object):
    '''Pipeline for extracting data into StreamItem instances.

    The pipeline has five sets of stages.  The *reader* stage reads
    from some input source and produces a series of StreamItem objects
    out.  *Incremental transforms* take single StreamItem objects in
    and produce single StreamItem objects out.  *Batch transforms* run
    on the entire set of StreamItem objects together.  There is a
    further set of *post-batch incremental transforms* which again run
    on individual StreamItem objects.  Finally, any number of *writers*
    send output somewhere, usually a streamcorpus.Chunk file.

    .. automethod:: __init__
    .. automethod:: run
    .. automethod:: _process_task

    '''
    def __init__(self, rate_log_interval, input_item_limit,
                 cleanup_tmp_files, tmp_dir_path, assert_single_source,
                 output_chunk_max_count, output_max_clean_visible_bytes,
                 reader, incremental_transforms, batch_transforms,
                 post_batch_incremental_transforms, writers):
        '''Create a new pipeline object.

        .. todo:: make this callable with just the lists of stages
                  and give sane (language-level) defaults for the rest

        :param int rate_log_interval: print progress every time this
          many input items have been processed
        :param int input_item_limit: stop after this many items
        :param bool cleanup_tmp_files: delete `tmp_dir_path` after
          execution if true
        :param str tmp_dir_path: path for intermediate files
        :param bool assert_single_source: require all items to have
          the same source value if true
        :param int output_chunk_max_count: restart output after
          writing this many items
        :param int output_max_clean_visible_bytes: restart output after
          writing this much content
        :param callable reader: reader stage object
        :param incremental_transforms: single-item transformation stages
        :paramtype incremental_transforms: list of callable
        :param batch_transforms: chunk-file transformation stages
        :paramtype batch_transforms: list of callable
        :param post_batch_incremental_transforms: single-item transformation
          stages
        :paramtype post_batch_incremental_transforms: list of callable
        :param writers: output stages
        :paramtype writers: list of callable

        '''
        self.rate_log_interval = rate_log_interval
        self.input_item_limit = input_item_limit
        self.cleanup_tmp_files = cleanup_tmp_files
        self.tmp_dir_path = tmp_dir_path
        self.assert_single_source = assert_single_source
        self.output_chunk_max_count = output_chunk_max_count
        self.output_max_clean_visible_bytes = output_max_clean_visible_bytes

        # stages that get passed in:
        self.reader = reader
        self.incremental_transforms = incremental_transforms
        self.batch_transforms = batch_transforms
        self.pbi_stages = post_batch_incremental_transforms
        self.writers = writers

        # current Chunk output file for incremental transforms
        self.t_chunk = None
        # context allows stages to communicate with later stages
        self.context = dict(
            i_str=None,
            data=None,
            )
        self.work_unit = None

    def _process_task(self, work_unit):
        '''Process a :class:`coordinate.WorkUnit`.

        The work unit's key is taken as the input file name.  The
        data should have ``start_count`` and ``start_chunk_time``
        values, which are passed on to :meth:`run`.

        :param work_unit: work unit to process
        :paramtype work_unit: :class:`coordinate.WorkUnit`
        :return: number of stream items processed

        '''
        self.work_unit = work_unit
        i_str = work_unit.key
        start_count = work_unit.data['start_count']
        start_chunk_time = work_unit.data['start_chunk_time']
        self.run(i_str, start_count, start_chunk_time)

    def run(self, i_str, start_count=0, start_chunk_time=None):
        '''Run the pipeline.

        This runs all of the steps described in the pipeline constructor,
        reading from some input and writing to some output.

        :param str i_str: name of the input file, or other reader-specific
          description of where to get input
        :param int start_count: index of the first stream item
        :param int start_chunk_time: timestamp for the first stream item

        '''
        try:
            if not os.path.exists(self.tmp_dir_path):
                os.makedirs(self.tmp_dir_path)

            if start_chunk_time is None:
                start_chunk_time = time.time()

            ## the reader returns generators of StreamItems
            i_chunk = self.reader(i_str)

            ## t_path points to the currently in-progress temp chunk
            t_path = None

            ## loop over all docs in the chunk processing and cutting
            ## smaller chunks if needed

            len_clean_visible = 0
            sources = set()
            next_idx = 0

            ## how many have we input and actually done processing on?
            input_item_count = 0

            for si in i_chunk:
                # TODO: break out a _process_stream_item function?
                next_idx += 1

                ## yield to the gevent hub to allow other things to run
                if gevent:
                    gevent.sleep(0)

                ## skip forward until we reach start_count
                if next_idx <= start_count:
                    continue

                if next_idx % self.rate_log_interval == 0:
                    ## indexing is zero-based, so next_idx corresponds
                    ## to length of list of SIs processed so far
                    elapsed = time.time() - start_chunk_time
                    if elapsed > 0:
                        rate = float(next_idx) / elapsed
                        logger.info('%d in %.1f --> %.1f per sec on '
                                    '(pre-partial_commit) %s',
                                    next_idx - start_count, elapsed, rate,
                                    i_str)

                if not self.t_chunk:
                    ## make a temporary chunk at a temporary path
                    # (Lazy allocation after we've read an item that might get processed out to the new chunk file)
                    # TODO: make this EVEN LAZIER by not opening the t_chunk until inside _run_incremental_transforms whe the first output si is ready
                    t_path = os.path.join(self.tmp_dir_path,
                                          't_chunk-%s' % uuid.uuid4().hex)
                    self.t_chunk = streamcorpus.Chunk(path=t_path, mode='wb')
                    assert self.t_chunk.message == streamcorpus.StreamItem_v0_3_0, self.t_chunk.message

                # TODO: a set of incremental transforms is equivalent
                # to a batch transform.  Make the pipeline explicitly
                # configurable as such:
                #
                # batch_transforms: [[incr set 1], batch op, [incr set 2], ...]
                #
                # OR: for some list of transforms (mixed incremental
                # and batch) pipeline can detect and batchify as needed

                ## incremental transforms populate t_chunk
                ## let the incremental transforms destroy the si by
                ## returning None
                si = self._run_incremental_transforms(
                    si, self.incremental_transforms)

                ## insist that every chunk has only one source string
                if si:
                    sources.add(si.source)
                    if self.assert_single_source and len(sources) != 1:
                        raise InvalidStreamItem(
                            'stream item %r had source %r, not %r '
                            '(set assert_single_source: false to suppress)' %
                            (si.stream_id, si.source, sources))

                if si and si.body and si.body.clean_visible:
                    len_clean_visible += len(si.body.clean_visible)
                    ## log binned clean_visible lengths, for quick stats estimates
                    #logger.debug('len(si.body.clean_visible)=%d' % int(10 * int(math.floor(float(len(si.body.clean_visible)) / 2**10)/10)))
                    #logger.debug('len(si.body.clean_visible)=%d' % len(si.body.clean_visible))

                if ((self.output_chunk_max_count is not None and
                     len(self.t_chunk) == self.output_chunk_max_count)):
                    logger.info('reached output_chunk_max_count (%d) at: %d',
                                len(self.t_chunk), next_idx)
                    self._process_output_chunk(
                        start_count, next_idx, sources, i_str, t_path)
                    start_count = next_idx

                elif (self.output_max_clean_visible_bytes is not None and
                      len_clean_visible >=
                      self.output_chunk_max_clean_visible_bytes):
                    logger.info(
                        'reached output_chunk_max_clean_visible_bytes '
                        '(%d) at: %d',
                        self.output_chunk_max_clean_visible_bytes,
                        len_clean_visible)
                    len_clean_visible = 0
                    self._process_output_chunk(
                        start_count, next_idx, sources, i_str, t_path)
                    start_count = next_idx

                input_item_count += 1
                if (((self.input_item_limit is not None) and
                     (input_item_count > self.input_item_limit))):
                    break

            if self.t_chunk is not None:
                self._process_output_chunk(
                    start_count, next_idx, sources, i_str, t_path)

            ## return how many stream items we processed
            return next_idx

        finally:
            if self.t_chunk is not None:
                self.t_chunk.close()
            for transform in self.batch_transforms:
                transform.shutdown()
            if self.cleanup_tmp_files:
                rmtree(self.tmp_dir_path)

    def _process_output_chunk(self, start_count, next_idx, sources, i_str,
                              t_path):
        '''
        for the current output chunk (which should be open):
          1. run batch transforms
          2. run post-batch incremental transforms
          3. run 'writers' to load-out the data to files or other storage
        return list of paths that writers wrote to
        '''
        if not self.t_chunk:
            # nothing to do
            return []
        self.t_chunk.close()

        # gather the paths as the writers run
        o_paths = None
        if len(self.t_chunk) > 0:
            # only batch transform and load if the chunk
            # isn't empty, which can happen when filtering
            # with stages like "find"

            # batch transforms act on the whole chunk in-place
            logger.info('running batch transforms on %d StreamItems',
                        len(self.t_chunk))
            self._run_batch_transforms(t_path)

            self._maybe_run_post_batch_incremental_transforms(t_path)

            # only proceed if above transforms left us with something
            if (self.t_chunk) and (len(self.t_chunk) >= 0):
                o_paths = self._run_writers(start_count, next_idx, sources,
                                            i_str, t_path)

        # we're now officially done with the chunk
        self.t_chunk = None

        # If we wrote some paths, update the data dictionary of outputs
        if self.work_unit and o_paths:
            old_o_paths = self.work_unit.data.get('output', [])
            o_paths = old_o_paths + o_paths
            self.work_unit.data['start_count'] = next_idx
            self.work_unit.data['output'] = o_paths
            self.work_unit.update()

    def _run_batch_transforms(self, chunk_path):
        '''Run all of the batch transforms over some intermediate chunk.'''
        for transform in self.batch_transforms:
            transform.process_path(chunk_path)

    def _maybe_run_post_batch_incremental_transforms(self, t_path):
        ## Run post batch incremental (pbi) transform stages.
        ## These exist because certain batch transforms have
        ## to run before certain incremental stages.
        if self.pbi_stages:
            t_path2 = os.path.join(self.tmp_dir_path, 'trec-kba-pipeline-tmp-%s' % str(uuid.uuid1()))
            # open destination for _run_incremental_transforms to write to
            self.t_chunk = streamcorpus.Chunk(path=t_path2, mode='wb')

            input_t_chunk = streamcorpus.Chunk(path=t_path, mode='rb')
            for si in input_t_chunk:
                self._run_incremental_transforms(si, self.pbi_stages)

            self.t_chunk.close()

            os.rename(t_path2, t_path)

    def _run_writers(self, start_count, next_idx, sources, i_str, t_path):
        '''Run all of the writers over some intermediate chunk.

        :param int start_count: index of the first item
        :param int next_idx: index of the next item (after the last
          item in this chunk)
        :param list sources: source strings included in this chunk
          (usually only one source)
        :param str i_str: name of input file or other input
        :param str t_path: location of intermediate chunk on disk
        :return: list of output file paths or other outputs

        '''
        # writers put the chunk somewhere, and could delete it
        name_info = dict(
            first=start_count,
            # num and md5 computed in each writers
            source=sources.pop(),
            )

        all_o_paths = []
        for writer in self.writers:
            logger.debug('running %r on %r: %r', writer, i_str, name_info)
            o_paths = writer(t_path, name_info, i_str)
            logger.debug('loaded (%d, %d) of %r into %r',
                         start_count, next_idx - 1, i_str, o_paths)
            all_o_paths += o_paths
        return all_o_paths

    def _run_incremental_transforms(self, si, transforms):
        '''
        Run transforms on stream item.
        Item may be discarded by some transform.
        Writes successful items out to current self.t_chunk
        Returns transformed item or None.
        '''
        ## operate each transform on this one StreamItem
        for transform in transforms:
            try:
                stream_id = si.stream_id
                si_new = transform(si, context=self.context)

                if si_new is None:
                    logger.warn('transform %r deleted %s abs_url=%r',
                                 transform, stream_id, si and si.abs_url)
                    return None
                si = si_new

            except TransformGivingUp:
                ## do nothing
                logger.info('transform %r giving up on %r',
                            transform, si.stream_id)

            except Exception, exc:
                logger.critical(
                    'transform %r failed on %r from i_str=%r abs_url=%r',
                    transform, si and si.stream_id, self.context.get('i_str'),
                    si and si.abs_url, exc_info=True)

        assert si is not None

        ## expect to always have a stream_time
        if not si.stream_time:
            raise InvalidStreamItem('empty stream_time: %s' % si)

        if si.stream_id is None:
            raise InvalidStreamItem('empty stream_id: %r' % si)

        ## put the StreamItem into the output
        if type(si) != streamcorpus.StreamItem_v0_3_0:
            raise InvalidStreamItem('incorrect stream item object %r' %
                                    type(si))
        self.t_chunk.add(si)
        return si


'''Generates metadata for use in naming files created by writers,
e.g. to_local_chunks and to_s3_chunks

.. This software is released under an MIT/X11 open source license.
   Copyright 2012-2014 Diffeo, Inc.
'''

import random
import datetime
import hashlib
import logging
from streamcorpus import Chunk
import re

logger = logging.getLogger(__name__)

is_hex_32 = re.compile('[a-fA-F0-9]{32}')

def get_name_info(chunk_path, assert_one_date_hour=False, i_str=None,
                  chunk_type=Chunk):
    '''
    takes a chunk blob and obtains the date_hour, md5, num

    makes fields:
    i_str
    input_fname
    input_md5 - parsed from input filename if it contains '-%(md5)s-'
    md5
    num
    epoch_ticks
    target_names
    doc_ids_8
    date_hour
    rand8
    date_now
    time_now
    date_time_now
    '''
    assert i_str is not None, 'must provide i_str as keyword arg'

    name_info = dict()
    if i_str:
        name_info['i_str'] = i_str
    else:
        name_info['i_str'] = ''

    i_fname = i_str.split('/')[-1]
    i_fname = i_fname.split('.')[0]  ## strip off .sc[.xz[.gpg]]
    name_info['input_fname'] = i_fname

    input_md5s = []
    for part in i_fname.split('-'):
        if len(part) == 32 and is_hex_32.match(part):
            input_md5s.append(part)
    name_info['input_md5'] = '-'.join(input_md5s)

    # TODO: return a dict-like object that does the expensive
    # calculation lazily, the name format might not even need that
    # value.
    ch = chunk_type(path=chunk_path, mode='rb')
    date_hours = set()
    target_names = set()
    doc_ids = set()
    epoch_ticks = None
    count = 0
    try:
        for si in ch:
            if chunk_type is Chunk:
                if epoch_ticks is None:
                    epoch_ticks = si.stream_time.epoch_ticks
                date_hours.add( si.stream_time.zulu_timestamp[:13] )
                doc_ids.add( si.doc_id )
                for annotator_id, ratings in si.ratings.items():
                    for rating in ratings:
                        target_name = rating.target.target_id.split('/')[-1]
                        target_names.add( target_name )
            count += 1
    except Exception, exc:
        logger.critical('failed to iter over chunk', exc_info=True)

    ## create the md5 property, so we can use it in the filename
    if hasattr(ch, 'md5_hexdigest'):
        name_info['md5'] = ch.md5_hexdigest
    else:
        try:
            data = open(chunk_path).read()
            name_info['md5'] = hashlib.md5(data).hexdigest()
        except Exception, exc:
            logger.critical('failed to compute md5', exc_info=True)
            name_info['md5'] = 'broken'
    name_info['num'] = count
    name_info['epoch_ticks'] = epoch_ticks

    name_info['target_names'] = '-'.join( target_names )
    name_info['doc_ids_8'] = '-'.join( [di[:8] for di in doc_ids] )

    if chunk_type is Chunk:
        if assert_one_date_hour:
            assert len(date_hours) == 1, \
                'got a chunk with other than one data_hour! ' + \
                repr(date_hours)

        if len(date_hours) > 0:
            date_hour = list(date_hours)[0]
            date_hour = date_hour.replace('T', '-')
        else:
            assert count == 0, (date_hours, count)
            date_hour = None
        name_info['date_hour'] = date_hour
    else:
        name_info['date_hour'] = 'NO-DATE-HOUR-FOR-FC'

    # TODO: in future lazy evaluation world, rand8 should return a
    # different value every time it is accessed so that a format could
    # be 'foo-{rand8}{rand8}'
    name_info['rand8'] = '%08x' % (random.randint(0, 0x7fffffff),)

    name_info['date_now'] = datetime.datetime.utcnow().strftime('%Y-%m-%d')
    name_info['time_now'] = datetime.datetime.utcnow().strftime('%H-%M-%S')
    name_info['date_time_now'] = datetime.datetime.utcnow().strftime('%Y-%m-%d-%H-%M-%S')

    return name_info


"""yakonfig declarations for streamcorpus_pipeline.

.. This software is released under an MIT/X11 open source license.
   Copyright 2014-2015 Diffeo, Inc.

If your program will use :mod:`streamcorpus_pipeline`, then include
the top-level module in your initial :mod:`yakonfig` setup, e.g.

.. code-block:: python

   with yakonfig.parse_args(parser, [streamcorpus_pipeline]):
     ...

.. todo:: Replace
          :class:`streamcorpus_pipeline._pipeline.PipelineFactory`
          with a :class:`yakonfig.factory.AutoFactory` implementation,
          which will mean that callers will need to pass a factory
          object to :mod:`yakonfig`.

.. todo:: Stop forcibly modifying stage configuration; but in turn
          figure out what we want the semantics of relative paths
          in the configuration to be.

"""
from __future__ import absolute_import
import collections
import logging
import os

import streamcorpus_pipeline
from streamcorpus_pipeline.stages import PipelineStages
from yakonfig import ConfigurationError, check_subconfig, NewSubModules

logger = logging.getLogger(__name__)

config_name = 'streamcorpus_pipeline'
default_config = {
    'tmp_dir_path': '/tmp',
    'third_dir_path': '/',
    'output_chunk_max_count': 500,
    'rate_log_interval': 100,
    'incremental_transforms': [],
    'batch_transforms': [],
    'post_batch_incremental_transforms': [],
    'cleanup_tmp_files': True,
    'assert_single_source': True,
    'reader': 'from_local_chunks',
    'writers': ['to_local_chunks'],
}
runtime_keys = {
    'tmp_dir_path': 'tmp_dir_path',
    'third_dir_path': 'third_dir_path',
}
# We can't find out our sub-modules just yet!  Since this gets evaluated
# at module import time, if an entrypoint stage tries to import
# streamcorpus_pipeline it can go wrong.
sub_modules = set()

static_stages = None

def replace_config(config, name):
    '''Replace the top-level pipeline configurable object.

    This investigates a number of sources, including
    `external_stages_path` and `external_stages_modules` configuration
    and `streamcorpus_pipeline.stages` entry points, and uses these to
    find the actual :data:`sub_modules` for
    :mod:`streamcorpus_pipeline`.

    '''
    global static_stages
    if static_stages is None:
        static_stages = PipelineStages()
        stages = static_stages
        if 'external_stages_path' in config:
            path = config['external_stages_path']
            if not os.path.isabs(path) and config.get('root_path'):
                path = os.path.join(config['root_path'], path)
            try:
                stages.load_external_stages(config['external_stages_path'])
            except IOError:
                return streamcorpus_pipeline  # let check_config re-raise this
        if 'external_stages_modules' in config:
            for mod in config['external_stages_modules']:
                try:
                    stages.load_module_stages(mod)
                except ImportError:
                    return streamcorpus_pipeline  # let check_config re-raise this
    else:
        stages = static_stages

    new_sub_modules = set(stage
                          for stage in stages.itervalues()
                          if hasattr(stage, 'config_name'))
    return NewSubModules(streamcorpus_pipeline, new_sub_modules)

def check_config(config, name):
    if 'tmp_dir_path' not in config:
        raise ConfigurationError('{0} requires tmp_dir_path setting'
                                 .format(name))

    # Checking stages:
    global static_stages
    if static_stages is None:
        static_stages = PipelineStages()
    stages = static_stages

    # (1) Push in the external stages;
    if 'external_stages_path' in config:
        try:
            stages.load_external_stages(config['external_stages_path'])
        except IOError, e:
            raise ConfigurationError(
                'invalid {0} external_stages_path {1}'
                .format(name, config['external_stages_path']), e)
    if 'external_stages_modules' in config:
            for mod in config['external_stages_modules']:
                try:
                    stages.load_module_stages(mod)
                except ImportError, e:
                    raise ConfigurationError(
                        'invalid {0} external_stages_modules value {1}'
                        .format(name, mod), e)

    # (2) Check the reader;
    if 'reader' not in config:
        raise ConfigurationError('{0} requires reader stage'
                                 .format(name))
    try:
        reader = stages[config['reader']]
    except ValueError, e:
        raise ConfigurationError(
            'invalid {0} reader {1}'
            .format(name, config['reader']))
    check_subconfig(config, name, reader)

    # (3) Check all of the intermediate and writers
    for phase in ['incremental_transforms', 'batch_transforms',
                  'post_batch_incremental_transforms', 'writers']:
        if phase not in config:
            raise ConfigurationError('{0} requires {1} stage list'
                                     .format(name, phase))
        if not isinstance(config[phase], collections.Iterable):
            raise ConfigurationError('{0} {1} must be a list of stages'
                                     .format(name, phase))
        for stagename in config[phase]:
            try:
                stage = stages[stagename]
            except KeyError, e:
                raise ConfigurationError(
                    'invalid {0} {1} {2}'
                    .format(name, phase, stagename))
            check_subconfig(config, name, stage)

def normalize_config(config):
    # Fix up all paths in our own config to be absolute
    root_path = config.get('root_path', os.getcwd())

    def fix(c, k):
        v = c[k]
        if v is None:
            return
        if isinstance(v, list):
            c[k] = [os.path.join(root_path, p) for p in v]
        if not os.path.isabs(v):
            c[k] = os.path.join(root_path, v)

    for k in config.iterkeys():
        if k.endswith('path') and k != 'root_path':
            fix(config, k)

    # Now go into all of our children and fix up their paths too.
    for c in config.itervalues():
        if isinstance(c, collections.MutableMapping):
            for k in c.iterkeys():
                if k.endswith('path'):
                    fix(c, k)