#!/usr/bin/env python

import sys
import nsysstats

from collections import defaultdict
from functools import cached_property


class CumulativeDuration:
    def __init__(self):
        self._duration = 0
        self._last_end_time = 0

    def add(self, start: int, end: int):
        if (self._last_end_time > end):
            return None

        self._duration = self._duration + end - max(start, self._last_end_time)
        self._last_end_time = end

    def get(self):
        return self._duration


class WddmData:
    def __init__(self):
        self._hw_queues_to_cumulative = defaultdict(CumulativeDuration)
        self._start_time = sys.maxsize
        self._end_time = -sys.maxsize

    def add(self, start_time: int, end_time: int, gpu: int, context: int, engine_type: int, node_ordinal: int,
            type: int):
        cumulative = self._hw_queues_to_cumulative[(gpu, context, engine_type, node_ordinal, type)]
        cumulative.add(start_time, end_time)

        self._start_time = min(self._start_time, start_time)
        self._end_time = max(self._end_time, end_time)

    def get_percent_of_utilization_time(self, keys: tuple, margin=1.0):
        cumulative = self._hw_queues_to_cumulative[keys]
        utilization_duration = cumulative.get()

        duration = self._end_time - self._start_time
        normalized_duration = duration * margin
        percentages = utilization_duration * 100.0 / normalized_duration

        return percentages

    def get_data(self):
        return self._hw_queues_to_cumulative


class WddmQueuesSummary(nsysstats.Report):
    def __init__(self, dbfile, args=[]):
        dburi_query = {
            'mode': 'rw',
            'nolock': '1',
        }

        super().__init__(dbfile, args, dburi_query)

        self._processes_to_wddm_data = defaultdict(WddmData)
        self._parent_hw_queue_descriptors = {}
        self._nodeordinal_to_friendlyname = {}

        self.init()

    def add_row(self, row: tuple, type: int):
        start_time = row[0]
        end_time = row[1]
        pid = row[2]
        gpu = row[4]
        context = row[5]
        engine_type = row[6]
        node_ordinal = row[7]

        wddm_data = self._processes_to_wddm_data[pid]
        wddm_data.add(start_time, end_time, gpu, context, engine_type, node_ordinal, type)

    @cached_property
    def get_wddm_contexts(self):
        table = 'TARGET_INFO_WDDM_CONTEXTS'
        get_wddm_contexts_query = f"""
            SELECT
                *
            FROM
                TARGET_INFO_WDDM_CONTEXTS
            """

        rows = None
        if (self.table_exists(table)):
            cursor = cur = self._dbcon.cursor()
            qcur = cursor.execute(get_wddm_contexts_query)
            rows = qcur.fetchall()

        return rows

    def get_wddm_context_info(self, context: int):
        rows = self.get_wddm_contexts
        if rows is not None:
            for row in rows:
                if (row[0] == context):
                    return row

        return None

    def insert_rows_to_wddm_dma_packets_table(self):
        query = f"""
            SELECT
                START.start as start,
                STOP.start as end,
                START.globalTid / 0x1000000 % 0x1000000 as pid,
                START.globalTid % 0x1000000 AS tid,
                START.gpu as gpu,
                START.context as context,
                START.engineType as engineType
            FROM
                WDDM_DMA_PACKET_START_EVENTS as START,
                WDDM_DMA_PACKET_STOP_EVENTS as STOP
            WHERE
                START.context == STOP.context AND START.uliSubmissionId == STOP.uliCompletionId
        """

        insert_query = 'INSERT INTO WDDM_DMA_PACKETS VALUES'
        cursor = self._dbcon.cursor()
        qcur = cursor.execute(query)
        rows = qcur.fetchall()
        for row in rows:
            context = row[5]
            wddm_context_info = self.get_wddm_context_info(context)
            node_ordinal = 0 if wddm_context_info is None else wddm_context_info[2]
            insert_query += '({}, {}, {}, {}, {}, {} ,{}, {}),'.format(*row, node_ordinal)

        cursor.execute(insert_query[:-1])

    def create_wddm_dma_packets_table(self):
        table = 'WDDM_DMA_PACKET_START_EVENTS'
        drop_dma_packet_table = f"""
        DROP TABLE IF EXISTS WDDM_DMA_PACKETS
        """

        create_dma_packet_table = f"""
CREATE TABLE WDDM_DMA_PACKETS (
    -- Processed ETW DMA packets events.

    start                       INTEGER,                               -- Event start timestamp (ns).
    end                         INTEGER,                               -- Event end timestamp (ns).
    pid                         INTEGER,                               -- The process ID.
    tid                         INTEGER,                               -- The thread ID.
    gpu                         INTEGER,                               -- GPU index.
    context                     INTEGER,                               -- The context ID of WDDM queue.
    engineType                  INTEGER,                               -- The engine type.
    nodeOrdinal                 INTEGER                                -- The node ordinal ID.
)
        """

        if (self.table_exists(table)):
            cursor = self._dbcon.cursor()
            cursor.execute(drop_dma_packet_table)
            cursor.execute(create_dma_packet_table)
            self.insert_rows_to_wddm_dma_packets_table()
            self._dbcon.commit()

    def insert_rows_to_wddm_queue_packets_table(self):
        query = f"""
            SELECT
                START.start as start,
                STOP.start as end,
                START.globalTid / 0x1000000 % 0x1000000 as pid,
                START.globalTid % 0x1000000 AS tid,
                START.gpu as gpu,
                START.context as context,
                START.engineType as engineType
            FROM
                WDDM_QUEUE_PACKET_START_EVENTS as START,
                WDDM_QUEUE_PACKET_STOP_EVENTS as STOP
            WHERE
                START.context == STOP.context AND START.submitSequence == STOP.submitSequence
            """

        insert_query = 'INSERT INTO WDDM_QUEUE_PACKETS VALUES'
        cursor = self._dbcon.cursor()
        qcur = cursor.execute(query)
        rows = qcur.fetchall()
        for row in rows:
            context = row[5]
            wddm_context_info = self.get_wddm_context_info(context)
            node_ordinal = 0 if wddm_context_info is None else wddm_context_info[2]
            insert_query += '({}, {}, {}, {}, {}, {} ,{}, {}),'.format(*row, node_ordinal)

        cursor.execute(insert_query[:-1])

    def create_wddm_queue_packets_table(self):
        table = 'WDDM_QUEUE_PACKET_START_EVENTS'

        drop_queue_packet_table = f"""
        DROP TABLE IF EXISTS WDDM_QUEUE_PACKETS
        """

        create_queue_packet_table = f"""
CREATE TABLE WDDM_QUEUE_PACKETS (
    -- Processed ETW queue packets events.

    start                       INTEGER,                               -- Event start timestamp (ns).
    end                         INTEGER,                               -- Event end timestamp (ns).
    pid                         INTEGER,                               -- The process ID.
    tid                         INTEGER,                               -- The thread ID.
    gpu                         INTEGER,                               -- GPU index.
    context                     INTEGER,                               -- The context ID of WDDM queue.
    engineType                  INTEGER,                               -- The engine type.
    nodeOrdinal                 INTEGER                                -- The node ordinal ID.
)
        """

        if (self.table_exists(table)):
            cursor = self._dbcon.cursor()
            cursor.execute(drop_queue_packet_table)
            cursor.execute(create_queue_packet_table)
            self.insert_rows_to_wddm_queue_packets_table()
            self._dbcon.commit()

    def create_wddm_paging_queue_packets_table(self):
        table = 'WDDM_PAGING_QUEUE_PACKET_START_EVENTS'

        drop_paging_queue_packet_table = f"""
        DROP TABLE IF EXISTS WDDM_PAGING_QUEUE_PACKETS
        """

        create_paging_queue_packet_table = f"""
CREATE TABLE WDDM_PAGING_QUEUE_PACKETS (
    -- Processed ETW paging queue packets events.

    start                       INTEGER,                               -- Event start timestamp (ns).
    end                         INTEGER,                               -- Event end timestamp (ns).
    pid                         INTEGER,                               -- The process ID.
    tid                         INTEGER,                               -- The thread ID.
    gpu                         INTEGER                                -- GPU index.
)
        """

        insert_into_paging_queue_packet_table = f"""
        INSERT INTO WDDM_PAGING_QUEUE_PACKETS
            SELECT
                START.start as start,
                STOP.start as end,
                START.globalTid / 0x1000000 % 0x1000000 as pid,
                START.globalTid % 0x1000000 AS tid,
                START.gpu as gpu
            FROM
                WDDM_PAGING_QUEUE_PACKET_START_EVENTS as START,
                WDDM_PAGING_QUEUE_PACKET_STOP_EVENTS as STOP
            WHERE
                START.pagingQueue == STOP.pagingQueue AND START.sequenceId == STOP.sequenceId
        """

        if (self.table_exists(table)):
            cursor = self._dbcon.cursor()
            cursor.execute(drop_paging_queue_packet_table)
            cursor.execute(create_paging_queue_packet_table)
            cursor.execute(insert_into_paging_queue_packet_table)
            self._dbcon.commit()

    def accumulate_dma_packets(self):
        table = 'WDDM_DMA_PACKET_START_EVENTS'
        queue_type = 1
        query = f"""
        SELECT
            *
        FROM
            WDDM_DMA_PACKETS
        """

        self.accumulate_packets(table, query, queue_type)

    def accumulate_queue_packets(self):
        table = 'WDDM_QUEUE_PACKET_START_EVENTS'
        queue_type = 0
        query = f"""
            SELECT
                *
            FROM
                WDDM_QUEUE_PACKETS
            """

        self.accumulate_packets(table, query, queue_type)

    def accumulate_packets(self, table: str, query: str, queue_type: int):
        if (self.table_exists(table)):
            cursor = self._dbcon.cursor()
            qcur = cursor.execute(query)
            rows = qcur.fetchall()
            for row in rows:
                self.add_row(row, queue_type)

    def map_parent_hw_queue_descriptors(self):
        table = 'WDDM_HW_QUEUE_EVENTS'
        query = f"""
            SELECT
                *
            FROM
                WDDM_HW_QUEUE_EVENTS
            """

        if (self.table_exists(table)):
            cursor = self._dbcon.cursor()
            qcur = cursor.execute(query)
            rows = qcur.fetchall()
            for row in rows:
                parent_hw_queue = row[6]
                hw_context = row[4]
                if (self._parent_hw_queue_descriptors.get(parent_hw_queue) is None):
                    self._parent_hw_queue_descriptors[parent_hw_queue] = hw_context

    def create_wddm_queues_utilization_table(self):
        drop_wddm_queues_utilization_table = f"""
        DROP TABLE IF EXISTS WDDM_QUEUES_UTILIZATION
        """

        create_wddm_queues_utilization_table = f"""
CREATE TABLE WDDM_QUEUES_UTILIZATION (
    -- WDDM queue utilization details.

    friendlyName                TEXT,                                  -- The queue-friendly name.
    pid                         INTEGER,                               -- The process ID.
    gpu                         INTEGER,                               -- GPU index.
    context                     INTEGER,                               -- The context ID of WDDM queue.
    engineType                  INTEGER,                               -- The engine type.
    nodeOrdinal                 INTEGER,                               -- The node ordinal ID.
    queueType                   INTEGER,                               -- The type of the packet.
    utilization                 REAL                                   -- The percentage of queue utilization.
)
        """

        cursor = self._dbcon.cursor()
        cursor.execute(drop_wddm_queues_utilization_table)
        cursor.execute(create_wddm_queues_utilization_table)
        self._dbcon.commit()

    def get_real_context(self, context: int):
        if (self._parent_hw_queue_descriptors.get(context) is not None):
            return self._parent_hw_queue_descriptors[context]

        return context

    def does_context_node_use_scg(self, hw_context: int, is_3d_engine_type: bool):
        queues_count = 0
        min_queues_for_scg = 2

        for hw_queue, hw_contex in self._parent_hw_queue_descriptors.items():
            if hw_contex == hw_context:
                queues_count += 1

        return queues_count > min_queues_for_scg and is_3d_engine_type

    def get_friendly_name(self, context: int, engine_type: int, nodeordinal: int):
        dxgk_engine_type_strings = [
            "Other",
            "3D",
            "Video Decode",
            "Video Encode",
            "Video Processing",
            "Scene Assembly",
            "Copy",
            "Overlay",
            "Crypto"
        ]

        real_context = self.get_real_context(context)
        if (real_context != context):
            is_3d_engine_type = True if engine_type == 1 else False
            if (self.does_context_node_use_scg(real_context, is_3d_engine_type)):
                return "3D/Comp"

        if (engine_type != 0):
            return dxgk_engine_type_strings[engine_type]

        if (self._nodeordinal_to_friendlyname.get(nodeordinal) is not None):
            return self._nodeordinal_to_friendlyname[nodeordinal]

        wddm_context_info = self.get_wddm_context_info(context)
        if (wddm_context_info is not None):
            node_ordinal = wddm_context_info[2]
            friendly_name = wddm_context_info[3]
            if friendly_name:
                self._nodeordinal_to_friendlyname[nodeordinal] = friendly_name
                return friendly_name

        return "Other"

    def insert_row_to_wddm_queues_utilization_table(self, key: tuple, pid: int, percentage: int):
        if (key):
            context = key[1]
            engine_type = key[2]
            node_ordinal = key[3]
            friendly_name = self.get_friendly_name(context, engine_type, node_ordinal)
            query = 'INSERT INTO WDDM_QUEUES_UTILIZATION VALUES (\'{}\', {}, {}, {}, {}, {}, {}, {:.2f})'.format(
                friendly_name, pid, *key, percentage)
            cursor = self._dbcon.cursor()
            cursor.execute(query)

    usage = f"""{{SCRIPT}} -- WDDM Queues Utilization Summary

    No arguments.

    Output: All time values default to nanoseconds
        FriendlyName : The queue-friendly name
        ProcessId : The process ID
        GPU : GPU index
        Context : The context ID of WDDM queue
        EngineType : The engine ID
        NodeOrdinal : The node ordinal ID
        QueueType : The queue type
        Utilization : The percentage of queue utilization
"""

    query = f"""
        SELECT
            *
        FROM
            WDDM_QUEUES_UTILIZATION as UTILIZATION
        ORDER BY utilization DESC
        """

    def init(self):
        self.create_wddm_dma_packets_table()
        self.create_wddm_queue_packets_table()
        self.create_wddm_paging_queue_packets_table()
        self.map_parent_hw_queue_descriptors()

        self.accumulate_dma_packets()
        self.accumulate_queue_packets()
        self.create_wddm_queues_utilization_table()
        for pid, wddm_data in self._processes_to_wddm_data.items():
            for key, data in wddm_data.get_data().items():
                self.insert_row_to_wddm_queues_utilization_table(
                    key, pid, wddm_data.get_percent_of_utilization_time(key))

        self._dbcon.commit()


if __name__ == "__main__":
    WddmQueuesSummary.Main()
