import sys
try: import cProfile as profile
except ImportError: import profile
try: from cStringIO import StringIO
except ImportError: import StringIO
from django.conf import settings

class ProfilerMiddleware(object):
    def can(self, request):
        return settings.PROFILING and 'prof' in request.GET \
            and (not settings.INTERNAL_IPS or \
            request.META['REMOTE_ADDR'] in settings.INTERNAL_IPS)
    
    def process_view(self, request, callback, callback_args, callback_kwargs):
        if self.can(request):
            self.profiler = profile.Profile()
            args = (request,) + callback_args
            return self.profiler.runcall(callback, *args, **callback_kwargs)

    def process_response(self, request, response):
        if self.can(request):
            self.profiler.create_stats()
            out = StringIO()
            old_stdout, sys.stdout = sys.stdout, out
            self.profiler.print_stats(1)
            sys.stdout = old_stdout
            response.content = '<pre>%s</pre>' % out.getvalue()
        return response

from django.db import connection
class DatabaseProfilerMiddleware(object):
    def can(self, request):
        return settings.PROFILING and 'dbprof' in request.GET \
            and (not settings.INTERNAL_IPS or \
            request.META['REMOTE_ADDR'] in settings.INTERNAL_IPS)
    
    def process_request(self, request):
        if self.can(request):
            self.DEBUG = settings.DEBUG
            settings.DEBUG = True
    
    def process_response(self, request, response):
        if self.can(request):
            out = StringIO()
            out.write('time\tsql\n')
            total_time = 0
            for query in reversed(sorted(connection.queries, key=lambda x: x['time'])):
                total_time += float(query['time'])*1000
                out.write('%s\t%s\n' % (query['time'], query['sql']))

            response.content = '<pre style="white-space:pre-wrap">%d queries executed in %.3f seconds\n\n%s</pre>' % (len(connection.queries), total_time/1000, out.getvalue())
            settings.DEBUG = self.DEBUG
        return response