"""
$URL: svn+ssh://svn.mems-exchange.org/repos/trunk/qp/lib/profiler.py $
$Id: profiler.py 28413 2006-05-25 19:36:51Z dbinger $
"""
from cStringIO import StringIO
from profile import runctx
from pstats import Stats
from qp.pub.common import get_publisher
from tempfile import NamedTemporaryFile
import sys


class Profile (object):

    def __init__(self, PATH_INFO='/', user=None, content='', **env):
        publisher = get_publisher()
        if isinstance(user, basestring):
            user = publisher.get_users().get(user)
        if user:
            session_id = 'bogus'
            site = publisher.get_site().get_name()
            env.setdefault('HTTP_COOKIE', '%s=%s' % (site, session_id))
            def abort():
                publisher.__class__.abort(publisher)
                session = publisher.create_session()
                session.owner = user
                session.effective_user = user
                get_publisher().get_sessions()[session_id] = session
            publisher.abort = abort
        env['PATH_INFO'] = PATH_INFO
        env.setdefault('SCRIPT_NAME', '')
        env.setdefault('HTTPS', 'on')
        env.setdefault('SERVER_NAME', 'example.org')
        if content:
            env.setdefault('REQUEST_METHOD', 'POST')
        else:
            env.setdefault('REQUEST_METHOD', 'GET')
        def commit():
            def ignore(*args, **kwargs):
                pass
            publisher.connection.storage.end = ignore
            publisher.__class__.commit(publisher)
            publisher.__class__.abort(publisher)
        publisher.commit = commit
        input = StringIO(content)
        output = StringIO()
        def process():
            hit = publisher.process(input, env)
            hit.get_response().write(output)
            return hit, output.getvalue()
        tmpfile = NamedTemporaryFile(suffix=".profile")
        context = dict()
        runctx("hit, page = process()", locals(), context, tmpfile.name)
        self.stats = Stats(tmpfile.name)
        self.stats.sort_stats('time', 'calls')
        tmpfile.close()
        self.hit = context.get('hit')
        self.response_string = context.get('page')
        if 'abort' in vars(publisher):
            del publisher.abort
        del publisher.commit

    def get_stats(self):
        return self.stats

    def get_response_string(self):
        return self.response_string

    def get_hit(self):
        return self.hit

    def get_status_code(self):
        return self.get_hit().get_response().get_status_code()

    def sort(self, *columns):
        self.stats.sort_stats(*columns)

    def format(self, *restrictions):
        stdout = sys.stdout
        output = StringIO()
        try:
            sys.stdout = output
            self.stats.print_stats(*restrictions)
        finally:
            sys.stdout = stdout
        return output.getvalue()

    def calls(self, function_name):
        return [value[0]
            for key, value in self.stats.stats.items()
            if key[-1] == function_name]

    def loads(self):
        return sum(self.calls('_p_load_state'))

    def __str__(self):
        self.sort('time', 'calls')
        return self.format(10)

    def callers(self, *args):
        self.stats.print_callers(*args)

    def callees(self, *args):
        self.stats.print_callees(*args)