#!/usr/bin/python

# This handles the systemtap equivalent of
# $(DTRACE) $(DTRACEFLAGS) -G -s $^ -o $@
# $(DTRACE) $(DTRACEFLAGS) -h -s $^ -o $@
# which is a step that builds DTrace provider and probe definitions

# Copyright (C) 2009 Red Hat Inc.
#
# This file is part of systemtap, and is free software.  You can
# redistribute it and/or modify it under the terms of the GNU General
# Public License (GPL); either version 2, or (at your option) any
# later version.

import os,posix,sys
from subprocess import call
from tempfile import mkstemp

class provider:
    def open(self, provider, header):
        have_provider = False
        self.f = open(provider)
        self.h = open(header,mode='w')
        self.h.write("/* Generated by the Systemtap dtrace wrapper */\n")
        self.h.write("\n#include <sys/sdt.h>\n\n")
        in_comment = False
        while (True):
            line = self.f.readline()
            if (line == ""):
                break
            if (line.find("/*") >= 0):
                in_comment = True
            if (line.find("*/") >= 0):
                in_comment = False
                continue
            if (in_comment):
                continue
            if (line.find("provider") >= 0):
                tokens = line.split()
                have_provider = True
                self.provider = tokens[1]
            elif (have_provider and line.find("probe ") > 0):
                while (line.find(")") < 0):
                    line += self.f.readline()
                this_probe = line[line.find("probe ")+5:line.find("(")].strip()
                this_probe_canon = self.provider.upper() + "_" + this_probe.replace("__","_").upper()
                args = (line[line.find("(")+1:line.find(")")])
                new_args = ""
                i = 0
                c = 0
                while (i < len(args)):
                    if (args[i:i+1] == ","):
                        new_args = ('%s%s' % (new_args, args[i]))
                        c += 1
                    else:
                        new_args = new_args + args[i]
                    i += 1
                if (len(new_args) == 0):
                    c = 0
                    stap_str = "STAP_PROBE(provider,%s" % (this_probe)
                else:
                    c += 1
                    stap_str = "STAP_PROBE%d(provider,%s" % (c,this_probe)
                define_str = "#define %s(" % (this_probe_canon)
                i = 1
                while (i <= c):
                    if (i != 1):
                        define_str += ","
                    define_str = define_str + "arg%s" % (i);
                    stap_str = stap_str + ",arg%s" % (i);
                    i += 1
                self.h.write ('/* %s (%s) */\n' % (this_probe_canon,new_args))
                self.h.write ('#define %s_ENABLED() 1\n' % this_probe_canon)
                self.h.write (define_str + ") \\\n")
                self.h.write (stap_str + ")\n\n")

def usage ():
    print "Usage " + sys.argv[0] + " [-h | -G] -s File.d -o File {Files}"
    sys.exit(1)

def open_file (arg):
    if (len (sys.argv) <= arg):
        return False
    try:
        file = open(sys.argv[arg], 'r')
    except IOError:
        print (sys.argv[arg] + " not found")
        sys.exit(1)
    return file


########################################################################
# main
########################################################################

if (len (sys.argv) < 2):
    usage()

i = 1
build_header = False
build_source = False
filename = ""
while (i < len (sys.argv)):
    if (sys.argv[i] == "-o"):
        i += 1
        filename = sys.argv[i]
    elif (sys.argv[i] == "-s"):
        i += 1
        s_filename = sys.argv[i]
    elif (sys.argv[i] == "-h"):
        build_header = True
    elif (sys.argv[i] == "-G"):
        build_source = True
    i += 1
if (build_header == False and build_source == False):
    usage()
    sys.exit(1)

if (filename == ""):
    if (s_filename != ""):
	(filename,ext) = os.path.splitext(s_filename)
        filename = os.path.basename(filename)
	if (build_header):
	    filename = filename + ".h"
	elif (build_source):
	    filename = filename + ".o"
    else:
        usage
        sys.exit(1)

if (build_header):
    providers = provider()
    providers.open(s_filename, filename)
elif (build_source):
    (basename,ext) = os.path.splitext(s_filename)
    basename = os.path.basename(basename)
    (d,fn) = mkstemp(suffix=".c",prefix=basename)
    f = open(fn,mode='w')
    f.write("static __dtrace () {}\n")
    f.close()
    call(["gcc", "-fPIC", "-c", fn, "-o", filename], shell=False)
    os.remove(fn)