from __future__ import print_function
import ompdModule
from ompd_handles import ompd_thread, ompd_task, ompd_parallel
import gdb
import sys
import traceback
from enum import Enum


class ompd_scope(Enum):
    ompd_scope_global = 1
    ompd_scope_address_space = 2
    ompd_scope_thread = 3
    ompd_scope_parallel = 4
    ompd_scope_implicit_task = 5
    ompd_scope_task = 6


class ompd_address_space(object):
    def __init__(self):
        """Initializes an ompd_address_space object by calling ompd_initialize
        in ompdModule.c
        """
        self.addr_space = ompdModule.call_ompd_initialize()
        # maps thread_num (thread id given by gdb) to ompd_thread object with thread handle
        self.threads = {}
        self.states = None
        self.icv_map = None
        self.ompd_tool_test_bp = None
        self.scope_map = {
            1: "global",
            2: "address_space",
            3: "thread",
            4: "parallel",
            5: "implicit_task",
            6: "task",
        }
        self.sched_map = {1: "static", 2: "dynamic", 3: "guided", 4: "auto"}
        gdb.events.stop.connect(self.handle_stop_event)
        self.new_thread_breakpoint = gdb.Breakpoint(
            "ompd_bp_thread_begin", internal=True
        )
        tool_break_symbol = gdb.lookup_global_symbol("ompd_tool_break")
        if tool_break_symbol is not None:
            self.ompd_tool_test_bp = gdb.Breakpoint("ompd_tool_break", internal=True)

    def handle_stop_event(self, event):
        """Sets a breakpoint at different events, e.g. when a new OpenMP
        thread is created.
        """
        if isinstance(event, gdb.BreakpointEvent):
            # check if breakpoint has already been hit
            if self.new_thread_breakpoint in event.breakpoints:
                self.add_thread()
                gdb.execute("continue")
                return
            elif (
                self.ompd_tool_test_bp is not None
                and self.ompd_tool_test_bp in event.breakpoints
            ):
                try:
                    self.compare_ompt_data()
                    gdb.execute("continue")
                except ():
                    traceback.print_exc()
        elif isinstance(event, gdb.SignalEvent):
            # TODO: what do we need to do on SIGNALS?
            pass
        else:
            # TODO: probably not possible?
            pass

    def get_icv_map(self):
        """Fills ICV map."""
        self.icv_map = {}
        current = 0
        more = 1
        while more > 0:
            tup = ompdModule.call_ompd_enumerate_icvs(self.addr_space, current)
            (current, next_icv, next_scope, more) = tup
            self.icv_map[next_icv] = (current, next_scope, self.scope_map[next_scope])
        print("Initialized ICV map successfully for checking OMP API values.")

    def compare_ompt_data(self):
        """Compares OMPT tool data about parallel region to data returned by OMPD functions."""
        # make sure all threads and states are set
        self.list_threads(False)

        thread_id = gdb.selected_thread().ptid[1]
        curr_thread = self.get_curr_thread()

        # check if current thread is LWP thread; return if "ompd_rc_unavailable"
        thread_handle = ompdModule.get_thread_handle(thread_id, self.addr_space)
        if thread_handle == -1:
            print("Skipping OMPT-OMPD checks for non-LWP thread.")
            return

        print("Comparing OMPT data to OMPD data...")
        field_names = [i.name for i in gdb.parse_and_eval("thread_data").type.fields()]
        thread_data = gdb.parse_and_eval("thread_data")

        if self.icv_map is None:
            self.get_icv_map()

        # compare state values
        if "ompt_state" in field_names:
            if self.states is None:
                self.enumerate_states()
            ompt_state = str(thread_data["ompt_state"])
            ompd_state = str(self.states[curr_thread.get_state()[0]])
            if ompt_state != ompd_state:
                print(
                    "OMPT-OMPD mismatch: ompt_state (%s) does not match OMPD state (%s)!"
                    % (ompt_state, ompd_state)
                )

        # compare wait_id values
        if "ompt_wait_id" in field_names:
            ompt_wait_id = thread_data["ompt_wait_id"]
            ompd_wait_id = curr_thread.get_state()[1]
            if ompt_wait_id != ompd_wait_id:
                print(
                    "OMPT-OMPD mismatch: ompt_wait_id (%d) does not match OMPD wait id (%d)!"
                    % (ompt_wait_id, ompd_wait_id)
                )

        # compare thread id
        if "omp_thread_num" in field_names and "thread-num-var" in self.icv_map:
            ompt_thread_num = thread_data["omp_thread_num"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.thread_handle,
                self.icv_map["thread-num-var"][1],
                self.icv_map["thread-num-var"][0],
            )
            if ompt_thread_num != icv_value:
                print(
                    "OMPT-OMPD mismatch: omp_thread_num (%d) does not match OMPD thread num according to ICVs (%d)!"
                    % (ompt_thread_num, icv_value)
                )

        # compare thread data
        if "ompt_thread_data" in field_names:
            ompt_thread_data = thread_data["ompt_thread_data"].dereference()["value"]
            ompd_value = ompdModule.call_ompd_get_tool_data(
                3, curr_thread.thread_handle
            )[0]
            if ompt_thread_data != ompd_value:
                print(
                    "OMPT-OMPD mismatch: value of ompt_thread_data (%d) does not match that of OMPD data union (%d)!"
                    % (ompt_thread_data, ompd_value)
                )

        # compare number of threads
        if "omp_num_threads" in field_names and "team-size-var" in self.icv_map:
            ompt_num_threads = thread_data["omp_num_threads"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_parallel_handle(),
                self.icv_map["team-size-var"][1],
                self.icv_map["team-size-var"][0],
            )
            if ompt_num_threads != icv_value:
                print(
                    "OMPT-OMPD mismatch: omp_num_threads (%d) does not match OMPD num threads according to ICVs (%d)!"
                    % (ompt_num_threads, icv_value)
                )

        # compare omp level
        if "omp_level" in field_names and "levels-var" in self.icv_map:
            ompt_levels = thread_data["omp_level"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_parallel_handle(),
                self.icv_map["levels-var"][1],
                self.icv_map["levels-var"][0],
            )
            if ompt_levels != icv_value:
                print(
                    "OMPT-OMPD mismatch: omp_level (%d) does not match OMPD levels according to ICVs (%d)!"
                    % (ompt_levels, icv_value)
                )

        # compare active level
        if "omp_active_level" in field_names and "active-levels-var" in self.icv_map:
            ompt_active_levels = thread_data["omp_active_level"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_parallel_handle(),
                self.icv_map["active-levels-var"][1],
                self.icv_map["active-levels-var"][0],
            )
            if ompt_active_levels != icv_value:
                print(
                    "OMPT-OMPD mismatch: active levels (%d) do not match active levels according to ICVs (%d)!"
                    % (ompt_active_levels, icv_value)
                )

        # compare parallel data
        if "ompt_parallel_data" in field_names:
            ompt_parallel_data = thread_data["ompt_parallel_data"].dereference()[
                "value"
            ]
            current_parallel_handle = curr_thread.get_current_parallel_handle()
            ompd_value = ompdModule.call_ompd_get_tool_data(4, current_parallel_handle)[
                0
            ]
            if ompt_parallel_data != ompd_value:
                print(
                    "OMPT-OMPD mismatch: value of ompt_parallel_data (%d) does not match that of OMPD data union (%d)!"
                    % (ompt_parallel_data, ompd_value)
                )

        # compare max threads
        if "omp_max_threads" in field_names and "nthreads-var" in self.icv_map:
            ompt_max_threads = thread_data["omp_max_threads"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.thread_handle,
                self.icv_map["nthreads-var"][1],
                self.icv_map["nthreads-var"][0],
            )
            if icv_value is None:
                icv_string = ompdModule.call_ompd_get_icv_string_from_scope(
                    curr_thread.thread_handle,
                    self.icv_map["nthreads-var"][1],
                    self.icv_map["nthreads-var"][0],
                )
                if icv_string is None:
                    print(
                        "OMPT-OMPD mismatch: omp_max_threads (%d) does not match OMPD thread limit according to ICVs (None Object)"
                        % (ompt_max_threads)
                    )
                else:
                    if ompt_max_threads != int(icv_string.split(",")[0]):
                        print(
                            "OMPT-OMPD mismatch: omp_max_threads (%d) does not match OMPD thread limit according to ICVs (%d)!"
                            % (ompt_max_threads, int(icv_string.split(",")[0]))
                        )
            else:
                if ompt_max_threads != icv_value:
                    print(
                        "OMPT-OMPD mismatch: omp_max_threads (%d) does not match OMPD thread limit according to ICVs (%d)!"
                        % (ompt_max_threads, icv_value)
                    )

        # compare omp_parallel
        # NOTE: omp_parallel = true if active-levels-var > 0
        if "omp_parallel" in field_names:
            ompt_parallel = thread_data["omp_parallel"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_parallel_handle(),
                self.icv_map["active-levels-var"][1],
                self.icv_map["active-levels-var"][0],
            )
            if (
                ompt_parallel == 1
                and icv_value <= 0
                or ompt_parallel == 0
                and icv_value > 0
            ):
                print(
                    "OMPT-OMPD mismatch: ompt_parallel (%d) does not match OMPD parallel according to ICVs (%d)!"
                    % (ompt_parallel, icv_value)
                )

        # compare omp_final
        if "omp_final" in field_names and "final-task-var" in self.icv_map:
            ompt_final = thread_data["omp_final"]
            current_task_handle = curr_thread.get_current_task_handle()
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                current_task_handle,
                self.icv_map["final-task-var"][1],
                self.icv_map["final-task-var"][0],
            )
            if icv_value != ompt_final:
                print(
                    "OMPT-OMPD mismatch: omp_final (%d) does not match OMPD final according to ICVs (%d)!"
                    % (ompt_final, icv_value)
                )

        # compare omp_dynamic
        if "omp_dynamic" in field_names and "dyn-var" in self.icv_map:
            ompt_dynamic = thread_data["omp_dynamic"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.thread_handle,
                self.icv_map["dyn-var"][1],
                self.icv_map["dyn-var"][0],
            )
            if icv_value != ompt_dynamic:
                print(
                    "OMPT-OMPD mismatch: omp_dynamic (%d) does not match OMPD dynamic according to ICVs (%d)!"
                    % (ompt_dynamic, icv_value)
                )

        # compare omp_max_active_levels
        if (
            "omp_max_active_levels" in field_names
            and "max-active-levels-var" in self.icv_map
        ):
            ompt_max_active_levels = thread_data["omp_max_active_levels"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_task_handle(),
                self.icv_map["max-active-levels-var"][1],
                self.icv_map["max-active-levels-var"][0],
            )
            if ompt_max_active_levels != icv_value:
                print(
                    "OMPT-OMPD mismatch: omp_max_active_levels (%d) does not match OMPD max active levels (%d)!"
                    % (ompt_max_active_levels, icv_value)
                )

                # compare omp_kind: TODO: Add the test for monotonic/nonmonotonic modifier
        if "omp_kind" in field_names and "run-sched-var" in self.icv_map:
            ompt_sched_kind = thread_data["omp_kind"]
            icv_value = ompdModule.call_ompd_get_icv_string_from_scope(
                curr_thread.get_current_task_handle(),
                self.icv_map["run-sched-var"][1],
                self.icv_map["run-sched-var"][0],
            )
            ompd_sched_kind = icv_value.split(",")[0]
            if self.sched_map.get(int(ompt_sched_kind)) != ompd_sched_kind:
                print(
                    "OMPT-OMPD mismatch: omp_kind kind (%s) does not match OMPD schedule kind according to ICVs (%s)!"
                    % (self.sched_map.get(int(ompt_sched_kind)), ompd_sched_kind)
                )

        # compare omp_modifier
        if "omp_modifier" in field_names and "run-sched-var" in self.icv_map:
            ompt_sched_mod = thread_data["omp_modifier"]
            icv_value = ompdModule.call_ompd_get_icv_string_from_scope(
                curr_thread.get_current_task_handle(),
                self.icv_map["run-sched-var"][1],
                self.icv_map["run-sched-var"][0],
            )
            token = icv_value.split(",")[1]
            if token is not None:
                ompd_sched_mod = int(token)
            else:
                ompd_sched_mod = 0
            if ompt_sched_mod != ompd_sched_mod:
                print(
                    "OMPT-OMPD mismatch: omp_kind modifier does not match OMPD schedule modifier according to ICVs!"
                )

        # compare omp_proc_bind
        if "omp_proc_bind" in field_names and "bind-var" in self.icv_map:
            ompt_proc_bind = thread_data["omp_proc_bind"]
            icv_value = ompdModule.call_ompd_get_icv_from_scope(
                curr_thread.get_current_task_handle(),
                self.icv_map["bind-var"][1],
                self.icv_map["bind-var"][0],
            )
            if icv_value is None:
                icv_string = ompdModule.call_ompd_get_icv_string_from_scope(
                    curr_thread.get_current_task_handle(),
                    self.icv_map["bind-var"][1],
                    self.icv_map["bind-var"][0],
                )
                if icv_string is None:
                    print(
                        "OMPT-OMPD mismatch: omp_proc_bind (%d) does not match OMPD proc bind according to ICVs (None Object)"
                        % (ompt_proc_bind)
                    )
                else:
                    if ompt_proc_bind != int(icv_string.split(",")[0]):
                        print(
                            "OMPT-OMPD mismatch: omp_proc_bind (%d) does not match OMPD proc bind according to ICVs (%d)!"
                            % (ompt_proc_bind, int(icv_string.split(",")[0]))
                        )
            else:
                if ompt_proc_bind != icv_value:
                    print(
                        "OMPT-OMPD mismatch: omp_proc_bind (%d) does not match OMPD proc bind according to ICVs (%d)!"
                        % (ompt_proc_bind, icv_value)
                    )

        # compare enter and exit frames
        if "ompt_frame_list" in field_names:
            ompt_task_frame_dict = thread_data["ompt_frame_list"].dereference()
            ompt_task_frames = (
                int(ompt_task_frame_dict["enter_frame"].cast(gdb.lookup_type("long"))),
                int(ompt_task_frame_dict["exit_frame"].cast(gdb.lookup_type("long"))),
            )
            current_task = curr_thread.get_current_task()
            ompd_task_frames = current_task.get_task_frame()
            if ompt_task_frames != ompd_task_frames:
                print(
                    "OMPT-OMPD mismatch: ompt_task_frames (%s) do not match OMPD task frames (%s)!"
                    % (ompt_task_frames, ompd_task_frames)
                )

        # compare task data
        if "ompt_task_data" in field_names:
            ompt_task_data = thread_data["ompt_task_data"].dereference()["value"]
            current_task_handle = curr_thread.get_current_task_handle()
            ompd_value = ompdModule.call_ompd_get_tool_data(6, current_task_handle)[0]
            if ompt_task_data != ompd_value:
                print(
                    "OMPT-OMPD mismatch: value of ompt_task_data (%d) does not match that of OMPD data union (%d)!"
                    % (ompt_task_data, ompd_value)
                )

    def save_thread_object(self, thread_num, thread_id, addr_space):
        """Saves thread object for thread_num inside threads dictionary."""
        thread_handle = ompdModule.get_thread_handle(thread_id, addr_space)
        self.threads[int(thread_num)] = ompd_thread(thread_handle)

    def get_thread(self, thread_num):
        """Get thread object from map."""
        return self.threads[int(thread_num)]

    def get_curr_thread(self):
        """Get current thread object from map or add new one to map, if missing."""
        thread_num = int(gdb.selected_thread().num)
        if thread_num not in self.threads:
            self.add_thread()
        return self.threads[thread_num]

    def add_thread(self):
        """Add currently selected (*) thread to dictionary threads."""
        inf_thread = gdb.selected_thread()
        try:
            self.save_thread_object(inf_thread.num, inf_thread.ptid[1], self.addr_space)
        except:
            traceback.print_exc()

    def list_threads(self, verbose):
        """Prints OpenMP threads only that are being tracking inside the "threads" dictionary.
        See handle_stop_event and add_thread.
        """
        list_tids = []
        curr_inferior = gdb.selected_inferior()

        for inf_thread in curr_inferior.threads():
            list_tids.append((inf_thread.num, inf_thread.ptid))
        if verbose:
            if self.states is None:
                self.enumerate_states()
            for (thread_num, thread_ptid) in sorted(list_tids):
                if thread_num in self.threads:
                    try:
                        print(
                            "Thread %i (%i) is an OpenMP thread; state: %s"
                            % (
                                thread_num,
                                thread_ptid[1],
                                self.states[self.threads[thread_num].get_state()[0]],
                            )
                        )
                    except:
                        traceback.print_exc()
                else:
                    print(
                        "Thread %i (%i) is no OpenMP thread"
                        % (thread_num, thread_ptid[1])
                    )

    def enumerate_states(self):
        """Helper function for list_threads: initializes map of OMPD states for output of
        'ompd threads'.
        """
        if self.states is None:
            self.states = {}
            current = int("0x102", 0)
            count = 0
            more = 1

            while more > 0:
                tup = ompdModule.call_ompd_enumerate_states(self.addr_space, current)
                (next_state, next_state_name, more) = tup

                self.states[next_state] = next_state_name
                current = next_state