#!/usr/bin/env python

import os.path
import sys
import xpra.scripts.main
import gobject

from pycallgraph import start_trace, make_dot_graph, GlobbingFilter

#network bits:
net = ['xpra.protocol.*',
       'xpra.protocol.Protocol.start',
       'xpra.bytestreams.*',
       'xpra.rencode.*',
       'xpra.server_base.XpraServer.process_packet',
       'xpra.server_source.ServerSource.next_packet',
       'socket.*',
       '*._write_thread_loop',
       '*._read_thread_loop',
       '*._read_parse_thread_loop',
       ]
damage = ['xpra.server_source.*', 'xpra.window_source.*',
          'xpra.batch_delay_calculator.*',
          'xpra.pixbuf_to_rgb.*',
          'xpra.deque.*', 'xpra.maths.*',
          'xpra.x264.*', 'xpra.vpx.*', 'xpra.webm.*',
          'xpra.xor.*']
mouse = ['xpra.server_base.XpraServer._process_pointer_position',
         'xpra.server_base.XpraServer._process_button_action',
         'xpra.server_base.XpraServer._process_mouse_common']
keyboard = ['xpra.keys.*', 'xpra.xkbhelper.*', 'wimpiggy.keys.*',
       'xpra.server_source.KeyboardConfig.*',
       'xpra.server_source.ServerSource.make_keymask_match',
       'xpra.server_source.ServerSource._keys_changed',
       'xpra.server_base.XpraServer._keys_changed',
       'xpra.server_base.XpraServer._process_key_action',
       'xpra.server_base.XpraServer._process_key_repeat',
       'xpra.server_base.XpraServer._clear_keys_pressed',
       'xpra.server_base.XpraServer._handle_key',
       'xpra.server_base.XpraServer._key_repeat',
       'xpra.server_base.XpraServer._key_repeat_timeout',
       'xpra.server_source.ServerSource.get_keycode',
       'xpra.client.XpraClient.get_keymap_properties',
       'xpra.client.XpraClient.get_keymap_modifiers',
       'xpra.client.XpraClient.send_key_action',
       'xpra.client.nn',
       'xpra.*._key_repeat',
       'xpra.*.clear_repeat',
       'xpra.*.key_handled_as_shortcut',
       ]
cursor = [
          'xpra.*._process_cursor',
          'xpra.*.set_windows_cursor',
          ]
bell = ['wimpiggy.window.WindowModel.do_wimpiggy_xkb_event',
        'wimpiggy.wm.Wm.bell_event',
        'wimpiggy.wm.Wm.do_bell_event',
        'xpra.server.XpraServer._bell_signaled',
        'xpra.server_source.ServerSource.bell']
misc = ['xpra.dotxpra.*', 'xpra.wait_for_server.*',
        'xpra.scripts.*', 'subprocess.*',
        'wimpiggy.log*',
        'wimpiggy.gobject_compat.*',
        'wimpiggy.tray.*',
        'xpra.version_util.*', 'xpra.build_info.*']

xsettings = ['xpra.xposix.xroot_props.*', 'xpra.xposix.xsettings.*', 'wimpiggy.xsettings_prop.*']
clipboard = ['xpra.platform.clipboard_base.*',
             'xpra.platform.gdk_clipboard.*', 'wimpiggy.gdk.*',
             'xpra.nested_main.*',
             'wimpiggy.selection.*',
             'xpra.*.ClientExtras.setup_clipboard_helper']
sound = ['xpra.sound.*',
         'xpra.server_source.ServerSource.stop_sending_sound',
         'gst.*', 'pygst.*']
gl = ['xpra.gl_client_window.*', 'xpra.gl_colorspace_conversions.*', 'xpra.gl_window_backing.*']

logging = [
       'logging.*',
       'wimpiggy.log.*']

std = ['pycallgraph.*',
       'traceback.*', 'linecache.*',
       '_weakrefset.*', 'weakref.*',
       'DLFCN.*',
       'abc.*',
       're.*', 'sre_parse.*', 'sre_compile.*',
       'atexit.*', 'warnings.*',
       'posixpath.*', 'genericpath.*', 'stat.*',
       'threading.*',
       'encodings.*',
       'optparse.*', 'gettext.*', 'locale.*', 'codecs.*']

libs = ['gobject.*', 'gtk.*', 'uuid.*', 'pygtk.*', 'gio.*', 'cairo.*',
        'os.environ.*', 'os._Environ.*', 'UserDict.*', 'platform.*', 'string.split',
        'dbus.*',
        'libxml2.*', "xml.*", 'StartElement', 'EndElement',    #used by gst..
        'ctypes.*', 'hmac.*']

one_offs = ['xpra.<module>',
            'xpra.client.<module>',
            'Queue.<module>', 'Queue.Full', 'Queue.Queue', 'Queue.Empty', 'Queue.LifoQueue', 'Queue.PriorityQueue',
            'xpra.server_uuid*',
            'xpra.server.XpraServer.__init__',
            'xpra.server_base.XpraServer.__init__',
            'xpra.server_base.XpraServer.x11_init',
            'xpra.server_base.XpraServer.init_x11_atoms',
            'xpra.server_base.XpraServer.load_existing_windows',
            'xpra.server_base.XpraServer.init_packet_handlers',
            'xpra.server_base.XpraServer.reenable_keymap_changes',
            'xpra.server_base.print_ready',
            'xpra.server.DesktopManager.__init__',
            'xpra.daemon_thread.*',
            'threading.Thread.daemon', 'threading._MainThread.daemon',
            'threading._MainThread.name',
            'threading._newname',
            'threading.Thread.setDaemon', 'threading.Thread.set_daemon', 'threading.Thread._set_daemon',
            'threading.Thread.__init__',
            'threading.Condition.*', 'threading.Event.*',
            'wimpiggy.util.gtk_main_quit_forever',
            'wimpiggy.util.gtk_main_quit_really',
            'wimpiggy.util.gtk_main_quit_on_fatal_exceptions_enable',
            'wimpiggy.wm.Wm.__init__'
            'wimpiggy.wm.Wm.__setup_ewmh_window'
            'wimpiggy.wm.Wm.enableCursors',
            'wimpiggy.*.n_arg_signal',
            'wimpiggy.*.<module>',
            'wimpiggy.error._ErrorManager.__init__',
            #client bits:
            'xpra.client_base.<module>',
            'xpra.client_base.ScreenshotXpraClient',
            'xpra.client_base.XpraClientBase',
            'xpra.client_base.InfoXpraClient',
            'xpra.client_base.VersionXpraClient',
            'xpra.client_base.StopXpraClient',
            'xpra.client_base.GobjectXpraClient',
            'xpra.client.XpraClient',
            'xpra.client_base.ClientSource',
            'xpra.client_base.*.__init__',
            'xpra.*.ClientExtras.__init__',
            'xpra.xposix.gui.<module>',
            'xpra.platform.*.ClientExtrasBase',
            'xpra.xposix.ClientExtras',
            'xpra.xposix.add_client_options', 'xpra.platform.add_*_option',
            'xpra.server_base.XpraServer.add_listen_socket',
            'xpra.server_base.XpraServer._process_shutdown_server',
            'socket._socketobject.meth',
            'xpra.*.*hello',
            'xpra.xposix.get_machine_id',
            'xpra.*.init_packet_handlers',
            'xpra.*.get_screen_sizes',
            'xpra.*.get_root_size',
            'xpra.*.parse_shortcuts',
            'xpra.*.ready',
            'xpra.*.setup_pa_audio_tagging',
            'xpra.*.setup_xprop_xsettings',
            'xpra.protocol.<module>',
            'xpra.protocol.rencode.<module>',
            'xpra.protocol.Protocol.__init__',
            'xpra.*.make_uuid',
            'xpra.cursor_names.*',
            'xpra.platform.client_tray.*',
            'xpra.platform.*.<module>',
            #keyboard stuff that we only do once:
            'xpra.*._do_keys_changed',
            'xpra.*.query_xkbmap',
            'xpra.*.grok_modifier_map',
            'xpra.*.get_keyboard_repeat',
            'xpra.*.get_x11_keymap',
            'xpra.*.get_gtk_keymap',
            'xpra.*.get_keymap_modifiers',
            'xpra.*.get_keymap_spec',
            'xpra.*.get_layout_spec',
            'xpra.*.exec_get_keyboard_data',
            'xpra.*.set_modifier_mappings',
            'xpra.*.update_modmap',
            'xpra.platform.keyboard_layouts.*',
            'xpra.platform.update_modmap',
            #some network stuff only happens once too:
            'xpra.*.set_max_packet_size',
            #exit stuff:
            'xpra.*._process_connection_lost',
            'xpra.*.warn_and_quit',
            'xpra.*.quit',
            'xpra.*.cleanup',
            'xpra.*.clean_mmap',
            'xpra.*.close_about',
            'xpra.*.hide_tray',
            ]

dialogs = [
            'xpra.*.setup_tray',
            'xpra.*.setup_statusicon',
            'xpra.*.get_tray_icon_filename',
            'xpra.*.get_tray_tooltip',
            'xpra.*.get_data_dir',
            'xpra.*.get_icons_dir',
            'xpra.*._is_ubuntu*',
            'xpra.*.setup_dbusnotify',
            'xpra.*.setup_xprops',
            'xpra.*.setup_x11_bell',
            'xpra.*.supports_system_tray',
            'xpra.*.supports_clipboard',
            'xpra.*.can_notify',
            #tray menu:
            'xpra.*.supports_server',
            'xpra.*.setup_menu',
            'xpra.*.make_*menuitem',
            'xpra.*.make_*submenu',
            'xpra.*.menuitem',
            'xpra.*.checkitem',
            'xpra.*.kbitem',
            'xpra.*.handshake_menuitem',
            'xpra.*.enable*menuitem',
            'xpra.*.activate*menu',
            'xpra.*.close*menu',
            'xpra.*.show*menu',
            'xpra.*.may_enable*menu',
            'xpra.*.set_*menuitem',
            'xpra.*.set_*menu',
            'xpra.*.menu_deactivated',
            'xpra.*.CheckMenuItem',
            'xpra.*.keysort',
            'xpra.*.ClientExtras.popup_menu_workaround',
            'xpra.*.ClientExtras.setup_xprops',
            'xpra.*.ClientExtras.setup_pa_audio_tagging',
            'xpra.*.ClientExtras.clipboard_toggled',
            'xpra.*.ClientExtras.bell_toggled',
            'xpra.*.ClientExtras.cursors_toggled',
            'xpra.*.ClientExtras.keyboard_sync_toggled',
            'xpra.*.ClientExtras.set_keyboard_sync_tooltip',
            'xpra.*.ClientExtras.microphone_state',
            'xpra.*.ClientExtras.speaker_state',
            'xpra.*.ClientExtras.set_selected_layout',
            'xpra.*.ClientExtras.set_menu_title',
            'xpra.*.ClientExtras.get_image',
            'xpra.*.ClientExtras.get_pixbuf',
            'xpra.*.ClientExtras.get_icon_filename',
            'xpra.*.ClientExtras.set_window_icon',
            'xpra.platform.client_extras_base.set_tooltip_text',
            'webbrowser.*',
            #network related, but only happens rarely (user action or initial connection):
            'xpra.*.send_bell_enabled',
            'xpra.*.send_cursors_enabled',
            'xpra.*.send_deflate_level',
            'xpra.*._process_set_deflate',
            #session info:
            'xpra.*.session_info',
            'xpra.*.session_info.*',
            'xpra.platform.graph.*',
            ]

connection = ['xpra.server_base.XpraServer.send_hello',
              'xpra.server_base.XpraServer.make_hello',
              'xpra.server_base.XpraServer._get_desktop_size_capability',
              'xpra.server_source.ServerSource.hello',
              #handle connection:
              'socket._socketobject.accept',
              'socket._socketobject.__init__',
              'xpra.bytestreams.SocketConnection.__init__',
              'xpra.bytestreams.SocketConnection.__str__',
              'xpra.server_base.XpraServer._new_connection',
              'xpra.server_base.XpraServer.verify_connection_accepted',
              'xpra.server_base.XpraServer._process_hello',
              'xpra.server_base.XpraServer.sanity_checks',
              'xpra.server_base.XpraServer.get_max_screen_size',
              'xpra.server_base.XpraServer.set_best_screen_size',
              'xpra.server_base.XpraServer.set_screen_size',
              'xpra.server_base.XpraServer.send_updated_screen_size',
              'xpra.server_base.XpraServer.set_workarea',
              'xpra.server_base.XpraServer.calculate_workarea',
              'xpra.server_base.XpraServer._process_set_deflate',
              'xpra.server_base.XpraServer._screen_size_changed',
              'xpra.server_base.XpraServer.set_keymap',
              'xpra.server_source.ServerSource.set_deflate',
              'xpra.server_source.ServerSource.parse_hello',
              'xpra.server_source.ServerSource.init_mmap',
              'xpra.server_source.ServerSource.keys_changed',
              'xpra.server_source.ServerSource.set_keymap',
              'xpra.server_source.ServerSource.updated_desktop_size',
              'xpra.server_source.ServerSource.set_screen_sizes',
              'xpra.server_source.ServerSource.set_encoding',
              'xpra.server_source.ServerSource.assign_keymap_options',
              'xpra.server.XpraServer.send_windows_and_cursors',
              'xpra.protocol.Protocol.set_compression_level',
              'xpra.protocol.Protocol.enable_rencode',
              'xpra.protocol.Protocol.do_start',
              #disconnection:
              'xpra.server_base.XpraServer.send_disconnect',
              'xpra.server_base.XpraServer._process_connection_lost',
              'xpra.server_base.XpraServer.cleanup_source',
              'xpra.protocol.Protocol.close',
              'xpra.protocol.Protocol.clean',
              'xpra.protocol.Protocol.terminate_io_threads',
              'xpra.bytestreams.SocketConnection.close',
              'socket._socketobject.close']

ALL = std + keyboard + cursor + bell + misc + xsettings + clipboard + sound + gl + logging + libs + one_offs + dialogs + connection

COMMON_THREAD_NAMES = ["write", "read", "parse", "format"]
SERVER_THREAD_NAMES = ["encode", "calculate_delay"]
CLIENT_THREAD_NAMES = ["draw"]
THREAD_NAMES = COMMON_THREAD_NAMES + SERVER_THREAD_NAMES + CLIENT_THREAD_NAMES
SETS = ["ALL", "std", "cursor", "keyboard", "bell", "misc", "xsettings", "clipboard", "sound", "gl", "logging", "libs", "one_offs", "dialogs", "connection"]


def usage(msg=None):
    if msg:
        print(msg)
    cmd = os.path.basename(sys.argv[0])
    print("%s usage: -i include-set -e exclude-set -d DELAY -r RUNTIME -t thread-name -- XPRA ARGS" % sys.argv[0])
    print("The default thread is the main thread, other options are:")
    print(" - for both client and server: %s" % ", ".join(COMMON_THREAD_NAMES))
    print(" - server-only threads: %s" % ", ".join(SERVER_THREAD_NAMES))
    print(" - client-only threads: %s" % ", ".join(CLIENT_THREAD_NAMES))
    print("The include and exclude sets are defined as a coma seperated list of package groups.")
    print("The package groups available are: %s" % ", ".join(SETS))
    print("The 'ALL' set is a superset containing all the other groups")
    print("Use '*' as a wildcard group")
    print("Use the delay to start profiling after a certain amount of time and to avoid")
    print("profiling the initial setup code. This can only apply to the main thread.")
    print("Use the runtime to automatically terminate the process after the given amount of time,")
    print("the time starts counting after the start delay.")
    print("")
    print("Examples:")
    print("#profile server:")
    print("%s -i '*' -e ALL -- start :10" % cmd)
    print("#profile client:")
    print("%s -i '*' -e ALL -- attach  :10" % cmd)
    print("#profile the client's draw thread:")
    print("%s -t draw -i '*' -e ALL -- attach  :10" % cmd)
    print("#profile the client's write thread without logging for 10 seconds, xpra runs without mmap and with x264 as primary encoding:")
    print("%s -t write -i '*' -e logging -r 10 -- attach  :10  --no-mmap --encoding=x264" % cmd)
    print("#profile server data_to_packet thread, excluding standard libraries")
    print("%s -t encode -i '*' -e std -e libs -- start :10" % cmd)
    sys.exit(1)


pos = 0
for x in sys.argv:
    if x=="--":
        break
    if x=="-h" or x=="--help":
        usage()
    pos += 1
if pos==0 or pos==len(sys.argv):
    usage("invalid number of arguments")

cg_args = sys.argv[1:pos]
sys.argv = sys.argv[:1]+sys.argv[pos+1:]
if len(cg_args)%2!=0:
    usage("invalid number of arguments")

pairs = []
for i in xrange(len(cg_args)/2):
    pairs.append((cg_args[i*2], cg_args[i*2+1]))

def get_group(v):
    if v=="*":
        return ["*"]
    if v not in SETS:
        usage("invalid package group %s, options are: %s" % (v, SETS))
    return globals()[v]

exclude = []
include = []
trace_thread = None
delay = 0
runtime = 0
for a,v in pairs:
    if a not in ("-i", "-e", "-d", "-r", "-t"):
        usage("invalid argument: %s" % a)
    if a=="-i":
        include += get_group(v)
    elif a=="-e":
        exclude += get_group(v)
    elif a=="-d":
        delay  = int(v)
    elif a=="-r":
        runtime  = int(v)
    elif a=="-t":
        if trace_thread:
            usage("only one thread can be traced at a time")
        if v not in THREAD_NAMES:
            usage("invalid thread name: %s, options are: %s" % (v, THREAD_NAMES))
        trace_thread=v
    else:
        usage("impossible!")

print("")
print("include=%s" % str(include))
print("exclude=%s" % str(exclude))
print("trace_thread=%s" % trace_thread)
print("delay=%s" % delay)
print("runtime=%s" % runtime)
print("")
if delay>0 and trace_thread:
    usage("delay (-d DELAY) cannot be used with the thread parameter (-t THREAD)")

#adjust cache size:
import fnmatch
fnmatch._MAXCACHE = max(fnmatch._MAXCACHE, len(exclude), len(include)) + 10
fnmatch._purge()

filter_func = GlobbingFilter(include=include, exclude=exclude)

if trace_thread:
    from xpra import daemon_thread
    saved_make_daemon_thread = daemon_thread.make_daemon_thread
    trace_count = 0
    def make_trace_daemon_thread(target, name):
        def trace_target(*args):
            global trace_count
            if name==trace_thread and trace_count==0:
                trace_count += 1
                print("started tracing %s / %s!" % (target, name))
                start_trace(filter_func=filter_func)
            else:
                print("not tracing %s / %s!" % (target, name))
            target()
        return saved_make_daemon_thread(trace_target, name)
    daemon_thread.make_daemon_thread = make_trace_daemon_thread
else:
    def do_start_trace(*args):
        print("starting trace")
        start_trace(filter_func=filter_func)
    #trace main thread
    if delay==0:
        do_start_trace()
    else:
        gobject.timeout_add(delay*1000, do_start_trace)

if runtime>0:
    def force_exit(*args):
        print("runtime %s expired, using SIGINT to force exit" % runtime)
        import signal
        os.kill(os.getpid(), signal.SIGINT)
    gobject.timeout_add((runtime+delay)*1000, force_exit)

print("calling xpra with: %s" % str(sys.argv))
x = xpra.scripts.main.main(__file__, sys.argv)

filename = 'pycallgraph.png'
make_dot_graph(filename)

sys.exit(x)
