# Forked from paste.cgitb_catcher, which is available under the MIT licens.
# http://www.opensource.org/licenses/mit-license.php
from __future__ import absolute_import

import cgitb
from cStringIO import StringIO
import sys

class CgitbMiddleware(object):
    """Forked from paste.cgitb_catcher.

    If bazjunk.catch_errors is False in environ, display a HTML
    error page instead, generated by a passed-in callable, unless
    the response has already started.
    
    Fix handling of GeneratorExit/.close()"""

    def __init__(self, app,
                 gen_html,
                 display=True,
                 logdir=None,
                 context=5,
                 format="html"):
        self.app = app
        self.gen_html = gen_html
        self.display = display
        self.logdir = logdir
        self.context = int(context)
        self.format = format
        
    def __call__(self, environ, start_response):
        try:
            app_iter = self.app(environ, start_response)
            return self.catching_iter(app_iter, environ)
        except:
            if 'bazjunk.throw_errors' in environ:
                raise
            exc_info = sys.exc_info()
            start_response('500 Internal Server Error',
                           [('content-type', 'text/html')],
                           exc_info)
            if environ.get('bazjunk.catch_errors',
                           lambda e: False)(environ):
                response = self.exception_handler(exc_info, environ)
            else:
                response = self.gen_html()
            return [response]

    def catching_iter(self, app_iter, environ):
        if not app_iter:
            raise StopIteration
        error_on_close = False
        try:
            for v in app_iter:
                yield v
            if hasattr(app_iter, 'close'):
                error_on_close = True
                app_iter.close()
        except GeneratorExit:
            if not error_on_close and hasattr(app_iter, 'close'):
                app_iter.close()
        except:
            if 'bazjunk.throw_errors' in environ:
                raise
            response = self.exception_handler(sys.exc_info(), environ)
            if not error_on_close and hasattr(app_iter, 'close'):
                try:
                    app_iter.close()
                except:
                    close_response = self.exception_handler(
                        sys.exc_info(), environ)
                    response += (
                        '<hr noshade>Error in .close():<br>%s'
                        % close_response)
            yield response

    def exception_handler(self, exc_info, environ):
        if not environ.get('bazjunk.catch_errors', lambda e: False)(environ):
            raise
        dummy_file = StringIO()
        hook = cgitb.Hook(file=dummy_file,
                          display=self.display,
                          logdir=self.logdir,
                          context=self.context,
                          format=self.format)
        hook(*exc_info)
        return dummy_file.getvalue()