import cProfile

import ppstats

from google.appengine.ext import webapp

from google.appengine.api import memcache

import google.appengine.ext.webapp.util

import os.path

import logging

import re

import random

import string

import zlib

mc_client = memcache.Client()

alphanumeric = string.letters + string.digits

global_profiler = None

class GAEProfiler(object):

    _save_every = 10

    def __init__(self):

        self.is_profiling = False

        self._profiler = None

        self.num_requests = 0

        self.requests_profiled = 0

        self.request_regex = None

        self.profile_key = ''.join([random.choice(alphanumeric) for x in range(4)])

    def start_profiling(self, request_regex=None, num_requests=0):

        "start profiling with this object, setting # of requests and filter"

        if self.is_profiling:

            return

        self.is_profiling = True

        if self._profiler is None:

            self._profiler = cProfile.Profile()

        self.num_requests = num_requests

        if request_regex:

            self.request_regex = re.compile(request_regex)

    def stop_profiling(self):

        self.is_profiling = False

    def resume_profiling(self):

        self.is_profiling = True

    def has_profiler(self):

        return self._profiler is not None

    def get_pstats(self):

        "return a ppstats object from current profile data"

        gae_base_dir = '/'.join(webapp.__file__.split('/')[:-5])

        sys_base_dir = '/'.join(logging.__file__.split('/')[:-2])

        stats = ppstats.Stats(self._profiler)

        stats.hide_directory(gae_base_dir, 'GAEHome')

        stats.hide_directory(sys_base_dir, 'SysHome')

        stats.strip_dirs()

        return stats

    def runcall(self, func, *args, **kwargs):

        "profile one call, incrementing requests_profiled and maybe saving stats"

        self.requests_profiled += 1

        if self._profiler:

            ret = self._profiler.runcall(func, *args, **kwargs)

        else:

            ret = func(*args, **kwargs)

#        if (self.requests_profiled % self._save_every) == 0 or \

#                self.requests_profiled == self.num_requests:

#            self.save_pstats_to_memcache()

        self.save_pstats_to_memcache()

        return ret

    def should_profile_request(self):

        "check for # of requests profiled and that SCRIPT_NAME matches regex"

        env = dict(os.environ)

        script_name = env.get('SCRIPT_NAME', '')

        logging.info(script_name)

        if self.num_requests and self.requests_profiled >= self.num_requests:

            return False

        if self.request_regex and not self.request_regex.search(script_name):

            return False

        return True

    def save_pstats_to_memcache(self):

        "save stats from profiler object to memcache"

        ps = self.get_pstats()

        output = ps.dump_stats_pickle()

        compressed_data = zlib.compress(output, 3)

        cache_key = cache_key_for_profile(self.profile_key)

        mc_client.set(cache_key, compressed_data)

        logging.info("Saved pstats to memcache with key %s" % cache_key)

def get_global_profiler():

    global global_profiler

    if not global_profiler:

        global_profiler = GAEProfiler()

    return global_profiler

def cache_key_for_profile(profile_key):

    "generate a memcache key"

    return "ProfileData.%s" % profile_key

def load_pstats_from_memcache(profile_key):

    "retrieve ppstats object"

    mc_data = mc_client.get(cache_key_for_profile(profile_key))

    if not mc_data:

        return None

    return ppstats.from_gz(mc_data)

def get_stats_from_global_or_request(request_obj):

    "get pstats for a key, or the global pstats"

    key = request_obj.get('key', '')

    if key:

        return load_pstats_from_memcache(key)

    else:

        gp = get_global_profiler()

        if not gp.has_profiler():

            return None

        return gp.get_pstats()

class show_profile(webapp.RequestHandler):

    def get(self):

        ps = get_stats_from_global_or_request(self.request)

        if not ps:

            self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

            return

        ps.set_output(self.response.out)

        sort = self.request.get('sort', 'time')

        ps.sort_stats(sort)

        self.response.out.write("<body><html><pre>\n")

        ps.print_stats(30)

        self.response.out.write("</pre></html></body>")

class download_profile_data(webapp.RequestHandler):

    def get(self):

        ps = get_stats_from_global_or_request(self.request)

        if not ps:

            self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

            return            

        output = ps.dump_stats_pickle()

        self.response.headers['Content-Type'] = 'application/octet-stream'

        self.response.out.write(output)

class show_profiler_status(webapp.RequestHandler):

    def get(self):

        gp = get_global_profiler()

        if not gp.has_profiler:

            self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

            return

        self.response.out.write("<html><body>")

        self.response.out.write("<b>Currently profiling:</b> %s<br>" % gp.is_profiling)

        self.response.out.write("<b>Profile Key</b>: %s<br>" % gp.profile_key)

        self.response.out.write("<b>Requests profiled so far:</b> %s<br>" % gp.requests_profiled)

        self.response.out.write("<b>Requests to profile:</b> %s<br>" % gp.num_requests)

        self.response.out.write("<b>Request regex:</b> %s<br>" % gp.request_regex)

        self.response.out.write("</body></html>")

class start_profiler(webapp.RequestHandler):

    def get(self):

        gp = get_global_profiler()

        gp.start_profiling()

class stop_profiler(webapp.RequestHandler):

    def get(self):

        gp = get_global_profiler()

        gp.stop_profiling()

        self.request.out.write("Content-Type: text/plain\n\n")

        self.request.out.write("done.")

class save_profile_data(webapp.RequestHandler):

    def get(self):

        gp = get_global_profiler()

def _add_our_endpoints(application):

    "insert our URLs into the application map"

    url_mapping = [(regex.pattern, handler) for (regex, handler) in application._url_mapping]

    return webapp.WSGIApplication(url_mapping, debug=True)

#

#  wrapper to for webapp applications

#

def run_wsgi_app(application):

    "proxy webapp.util's call to profile when needed"

    gp = get_global_profiler()

    if gp.is_profiling and gp.should_profile_request():

        return gp.runcall(google.appengine.ext.webapp.util.run_wsgi_app, *(application,))

    else:

        return google.appengine.ext.webapp.util.run_wsgi_app(application)

#

# middleware for django applications

#

class ProfileMiddleware(object):

    def __init__(self):

        self.profiler = None

    def process_request(self, request):

        self.profiler = get_global_profiler()

    def process_view(self, request, callback, callback_args, callback_kwargs):

        if self.profiler.is_profiling:

            return self.profiler.runcall(callback, request, *callback_args, **callback_kwargs)

application = webapp.WSGIApplication(

    [('/profiler/start', start_profiler),

     ('/profiler/stop', stop_profiler),

     ('/profiler/show', show_profile),

     ('/profiler/download', download_profile_data),

     ('/profiler/status', show_profiler_status),

     ],

    debug=True)

def main():

    google.appengine.ext.webapp.util.run_wsgi_app(application)

if __name__ == '__main__':

    main()