import cProfile

import ppstats

from google.appengine.ext import webapp

from google.appengine.api import memcache

import google.appengine.ext.webapp.util

from email.MIMEMultipart import MIMEMultipart

from email.Message import Message

import httplib

import logging

import os.path

import random

import re

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 new_global_profiler():

    global 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:

        gp = GAEProfiler()

        gp.profile_obj = load_pstats_from_memcache(key)

        gp.profile_key = key

        return gp

    else:

        gp = get_global_profiler()

        if not gp.has_profiler():

            return None

        return gp

def mime_upload_data_as_file(field_name, filename, body):

    part = Message()

    part['Content-Disposition'] = 'form-data; name="%s"; filename="%s"' % (field_name, filename)

    part['Content-Transfer-Encoding'] = 'binary'

    part['Content-Type'] = 'application/octet-stream'

    part['Content-Length'] = str(len(body))

    part.set_payload(body)

    return part

def mime_form_value(name, value):

    part = Message()

    part['Content-Disposition'] = 'form-data; name="%s"' % name

    part.set_payload(value)

    return part

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.profile_obj.set_output(self.response.out)

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

        ps.profile_obj.sort_stats(sort)

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

        ps.profile_obj.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.profile_obj.dump_stats_pickle()

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

        self.response.out.write(output)

class send_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            

        dest = self.request.get('dest', '')

        if not dest:

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

        upload_form = MIMEMultipart('form-data')

        upload_filename =  'profile.%s.pstats' % ps.profile_key

        upload_field_name = 'profile_file'

        upload_form.attach(mime_upload_data_as_file('profile_file', upload_field_name, zlib.compress(ps.profile_obj.dump_stats_pickle())))

        upload_form.attach(mime_form_value('key_only', '1'))

        http_conn = httplib.HTTPConnection(dest)

        http_conn.connect()

        http_conn.request('POST', '/upload_profile', upload_form.as_string(), 

                          {'Content-Type': 'multipart/form-data; boundary=%s' % upload_form.get_boundary()})

        http_resp = http_conn.getresponse()

        remote_data = http_resp.read()

        if http_resp.status == 200:

            remote_url = "http://%s/view_profile?key=%s" % (dest, remote_data)

            self.response.out.write("<html><body>Success! <a href='%s'>%s</a></body></html>" % (remote_url, remote_url))

        else:

            self.response.out.write("Failure!\n%s: %s\n%s" % (http_resp.status, http_resp.reason, remote_data))

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 = new_global_profiler()

        gp.start_profiling()

        self.response.headers['Content-Type'] = "text/plain"

        self.response.out.write("Started profiling (key: %s).\n" % gp.profile_key)

        self.response.out.write("Retrieve saved results at <a href='/profiler/show?key=%(key)s'>/profiler/show?key=%(key)s).\n" % {'key':gp.profile_key})

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),

     ('/profiler/send', send_profile_data),

     ],

    debug=True)

def main():

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

if __name__ == '__main__':

    main()