# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

from asyncore import file_dispatcher
import importlib
import numpy as np
import os
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import subprocess
import sys

from pathlib import Path

import nsysstats

from nsys_recipe.log import logger
from nsys_recipe.lib import exceptions, nsys_path
from nsys_recipe import nsys_constants

EVENT_TYPE_NVTX_DOMAIN_CREATE = 75
EVENT_TYPE_NVTX_PUSHPOP_RANGE = 59
EVENT_TYPE_NVTX_STARTEND_RANGE = 60


class _Loader:
    output_suffix = ""

    def __init__(self, report_path):
        self._report_path = report_path

    def _validate_export_time(self, path):
        if not Path(path).exists():
            return False

        return os.path.getctime(self._report_path) < os.path.getctime(path)

    def validate_table(self, path, table, columns):
        raise NotImplementedError()

    def read_table(self, handle, table, columns):
        raise NotImplementedError()

    def list_tables(self, path):
        raise NotImplementedError


class _ParquetLoader(_Loader):
    output_suffix = "_pqtdir"
    file_extension = ".parquet"

    def validate_table(self, path, table, columns):
        file_path = Path(path) / f"{table}{self.file_extension}"

        if not self._validate_export_time(file_path):
            return None

        parquet_file = pq.ParquetFile(file_path)
        parquet_schema = parquet_file.schema
        parquet_columns = parquet_schema.names

        if columns and any(column not in parquet_columns for column in columns):
            return None

        return parquet_file

    def read_table(self, handle, table, columns):
        if not columns:
            columns = None

        return handle.read(columns).to_pandas()

    def list_tables(self, path):
        path = Path(path)

        if not path.exists():
            return []

        files = list(path.glob(f"*{self.file_extension}"))
        return [file.stem for file in files]


class _ArrowLoader(_Loader):
    output_suffix = "_arwdir"
    file_extension = ".arrow"

    def validate_table(self, path, table, columns):
        file_path = Path(path) / f"{table}{self.file_extension}"

        if not self._validate_export_time(file_path):
            return None

        reader = pa.RecordBatchStreamReader(file_path)
        arrow_schema = reader.schema
        arrow_columns = arrow_schema.names

        if columns and any(column not in arrow_columns for column in columns):
            return None

        return reader

    def read_table(self, handle, table, columns):
        df = handle.read_pandas()
        return df[columns] if columns else df

    def list_tables(self, path):
        path = Path(path)

        if not path.exists():
            return []

        files = list(path.glob(f"*{self.file_extension}"))
        return [file.stem for file in files]


class _SqliteLoader(_Loader):
    output_suffix = ".sqlite"
    file_extension = ".sqlite"

    def validate_tables(self, path, tables):
        if not self._validate_export_time(path):
            return None

        try:
            sql_report = nsysstats.Report(path)
        except nsysstats.Report.Error_InvalidDatabaseFile:
            return None

        if sql_report is None:
            return None

        if tables and any(not sql_report.table_exists(table) for table in tables):
            return None

        return sql_report

    def validate_table(self, path, table, columns):
        sql_report = self.validate_tables(path, [table])

        if sql_report is None:
            return None

        if columns and any(
            not sql_report.table_col_exists(table, column) for column in columns
        ):
            return None

        return sql_report

    def read_sql_query(self, handle, query):
        df = pd.read_sql(query, handle.dbcon)
        return df

    def read_table(self, handle, table, columns):
        column_query = ",".join(columns) if columns else "*"
        query = f"SELECT {column_query} FROM {table}"
        return self.read_sql_query(handle, query)

    def list_tables(self, path):
        if Path(path) is None:
            return []

        try:
            return nsysstats.Report(path).tables
        except Exception:
            return []


class StatsService:
    def __init__(self, report_path):
        self._data_service = DataService(report_path)
        self._sqlite_file = Path(report_path).with_suffix(".sqlite")

    def _get_stats_class(self, module, class_name):
        try:
            stats_module_path = f"{nsys_constants.NSYS_REPORT_DIR}.{module}"
            module = importlib.import_module(stats_module_path)
        except ModuleNotFoundError as e:
            try:
                rule_module_path = f"{nsys_constants.NSYS_RULE_DIR}.{module}"
                module = importlib.import_module(rule_module_path)
            except ModuleNotFoundError:
                raise exceptions.StatsModuleNotFoundError(e) from e

        return getattr(module, class_name)

    def _setup(self, stats_class, parsed_args):
        report, exitval, errmsg = stats_class.Setup(self._sqlite_file, parsed_args)

        if report is not None:
            return report

        if exitval == stats_class.EXIT_NODATA:
            return None

        raise exceptions.StatsInternalError(errmsg)

    def read_sql_report(self, module, class_name, parsed_args):
        stats_class = self._get_stats_class(module, class_name)
        report = None

        if self._data_service._get_service("sqlite")._validate_export_time(
            self._sqlite_file
        ):
            report = self._setup(stats_class, parsed_args)

        if report is None:
            hints = {"format": "sqlite"}
            if not self._data_service.export(None, hints):
                return None

            report = self._setup(stats_class, parsed_args)

        if report is None:
            return None

        return pd.read_sql(report.get_query(), report.dbcon)


class DataService:
    def __init__(self, report_path):
        if not Path(report_path).exists:
            raise FileNotFoundError(f"{report_path} does not exist.")

        self._report_path = Path(report_path)
        self._service_instances = {}

    def _create_service(self, format):
        loader_map = {
            "parquetdir": _ParquetLoader,
            "arrowdir": _ArrowLoader,
            "sqlite": _SqliteLoader,
        }

        if format not in loader_map:
            raise NotImplementedError("Invalid format type.")

        return loader_map[format](self._report_path)

    def _get_service(self, format):
        if format not in self._service_instances:
            self._service_instances[format] = self._create_service(format)

        return self._service_instances[format]

    def _get_output_path(self, hints):
        service = self._get_service(hints.get("format", "parquetdir"))

        default_output_path = (
            str(self._report_path.with_suffix("")) + service.output_suffix
        )
        return hints.get("path", default_output_path)

    def _get_export_args(self, tables, hints):
        format_type = hints.get("format", "parquetdir")
        mode = hints.get("mode", "partial")
        output_path = self._get_output_path(hints)

        export_args = [
            "--type",
            format_type,
            "--output",
            output_path,
            "--force-overwrite",
            "true",
            "--lazy",
            "false",
        ]

        if mode == "partial" and tables:
            export_args += ["--tables", ",".join(tables)]

        return export_args

    def export(self, tables, hints):
        format_type = hints.get("format", "parquetdir")
        output_path = self._get_output_path(hints)

        # TODO(DTSP-15985): Support directory of files for SQLite.
        # We do not currently offer support for partial exports in SQLite.
        # Regardless of the 'table' argument, all tables will always be imported.
        if format_type == "sqlite":
            tables = None

        if isinstance(tables, str):
            tables = [tables]

        export_args = self._get_export_args(tables, hints)

        nsys_exe = nsys_path.find_installed_file(nsys_constants.NSYS_EXE_NAME)
        cmd = [nsys_exe, "export", *export_args, self._report_path]
        if tables:
            logger.info(
                f"Exporting {tables} from {self._report_path} to {output_path}..."
            )
        else:
            logger.info(f"Exporting {self._report_path} to {output_path}...")

        try:
            p = subprocess.run(cmd, text=True, capture_output=True)
        except Exception as e:
            logger.error(f"Failed to export {self._report_path}: {e}")
            return False

        if p.returncode:
            logger.error(f"Failed to export {self._report_path}: {p.stderr}")
            return False

        return True

    def read_tables(self, table_column_dict, hints=None):
        """Read tables into dataframes.

        Parameters
        ----------
        table_column_dict : dict
            Dictionary of tables and columns.
        hints : dict
            Additional configurations.
        """
        if hints is None:
            hints = {}

        result_dict = {}
        tables_to_export = []

        output_path = self._get_output_path(hints)
        overwrite = hints.get("overwrite", False)
        service = self._get_service(hints.get("format", "parquetdir"))

        if not overwrite:
            for table, columns in table_column_dict.items():
                handle = service.validate_table(output_path, table, columns)
                if handle is not None:
                    result_dict[table] = service.read_table(handle, table, columns)
                else:
                    tables_to_export.append(table)
        else:
            tables_to_export = list(table_column_dict.keys())

        if tables_to_export and not self.export(tables_to_export, hints):
            return None

        for table in tables_to_export:
            columns = table_column_dict[table]
            handle = service.validate_table(output_path, table, columns)
            if handle is not None:
                result_dict[table] = service.read_table(handle, table, columns)
            else:
                logger.error(
                    f"Could not validate '{table}'."
                    " Please ensure the table and column names are correct or"
                    " re-try with a recent version of Nsight Systems."
                )
                return None

        return result_dict

    def read_sql_query(self, query, tables=None, hints=None):
        """Read the SQL query into a dataframe.

        Parameters
        ----------
        query : str
            SQL query to execute.
        tables : list of str or str
            If specified, the function will export the tables before executing
            the query and check whether the table names are valid. If no
            tables are provided, all tables will be exported, and no checks
            will be made before executing the query.
        hints : dict
            Additional configurations.
        """
        if hints is None:
            hints = {}

        format_type = hints.get("format", "sqlite")
        if format_type != "sqlite":
            raise NotImplementedError("Invalid format type")

        output_path = self._get_output_path(hints)
        overwrite = hints.get("overwrite", False)
        service = self._get_service(format_type)

        if isinstance(tables, str):
            tables = [tables]

        if not overwrite and tables:
            handle = service.validate_tables(output_path, tables)
            if handle is not None:
                return service.read_sql_query(handle, query)

        if not self.export(tables, hints):
            return None

        handle = service.validate_tables(output_path, tables)
        if handle is not None:
            return service.read_sql_query(handle, query)

        logger.error(
            f"Could not validate {tables}."
            " Please ensure the table names are correct or"
            " re-try with a recent version of Nsight Systems."
        )

    def list_tables(self, hints=None):
        """List the available tables in the report file."""
        if hints is None:
            hints = {}

        service = self._get_service(hints.get("format", "parquetdir"))
        return service.list_tables(self._get_output_path(hints))


def _get_time_cols(df):
    if "start" in df.columns:
        if "end" in df.columns:
            # Time range.
            return ("start", "end")
        else:
            # Point in time.
            return ("start", "start")

    if "timestamp" in df.columns:
        # Point in time.
        return ("timestamp", "timestamp")
    elif "rawTimestamp" in df.columns:
        # Point in time.
        return ("rawTimestamp", "rawTimestamp")
    else:
        raise NotImplementedError


def filter_by_time_range(dfs, start_time=0, end_time=sys.maxsize):
    """Filter the dataframe(s) to retain only events that start
    or end within the given range.

    Parameters
    ----------
    dfs : list of dataframes or dataframe
        Dataframes to filter.
    start_time : int
        Start time of the desired range.
    end_time : int
        End time of the desired range.
    """
    if not isinstance(dfs, list):
        dfs = [dfs]

    for df in dfs:
        if df.empty:
            continue

        start_col, end_col = _get_time_cols(df)

        mask = pd.Series(True, index=df.index)
        if end_time is not None:
            mask &= df[start_col] <= end_time
        if start_time is not None:
            mask &= df[end_col] >= start_time

        df.drop(df[~mask].index, inplace=True)


def apply_time_offset(dfs, session_offset):
    """Synchronize session start times.

    Parameters
    ----------
    dfs : list of dataframes or dataframe
        Dataframes to filter.
    session_offset : int
        Offset of the session time
    """
    if not isinstance(dfs, list):
        dfs = [dfs]

    for df in dfs:
        if df.empty:
            continue

        start_col, end_col = _get_time_cols(df)

        df.loc[:, start_col] += session_offset

        if start_col != end_col:
            df.loc[:, end_col] += session_offset


def compute_session_duration(analysis_df, target_info_df, min_session):
    profile_duration = analysis_df.at[0, "duration"]
    session_time = target_info_df.at[0, "utcEpochNs"]

    session_offset = session_time - min_session
    profile_duration += session_offset

    return session_offset, profile_duration


def replace_id_with_value(main_df, str_df, id_column):
    """Replace the values in 'id_column' of 'main_df' with the corresponding
    string value stored in 'str_df'.

    Parameters
    ----------
    main_df : dataframe
        Dataframe containing 'id_column'.
    str_df : dataframe
        Dataframe 'StringId' that maps IDs to string values.
    id_column : str
        Name of the column that should be replaced with the corresponding
        string values.
    """
    renamed_str_df = str_df.rename(columns={"id": id_column})
    merged_df = main_df.merge(renamed_str_df, on=id_column, how="left")

    # Drop the original 'id_column' column and rename 'value' column to 'id_column'.
    merged_df = merged_df.drop(columns=[id_column])

    merged_df = merged_df.rename(columns={"value": id_column})

    return merged_df


def get_domain_range_df(nvtx_df, domain_name):
    """Get push/pop and start/end ranges of the specified domain."""
    domain_df = nvtx_df[
        (nvtx_df["eventType"] == EVENT_TYPE_NVTX_DOMAIN_CREATE)
        & (nvtx_df["text"] == domain_name)
    ]

    if domain_df.empty:
        return domain_df

    domain_id = domain_df["domainId"].iloc[0]

    return nvtx_df[
        (nvtx_df["domainId"] == domain_id)
        & (
            nvtx_df["eventType"].isin(
                [EVENT_TYPE_NVTX_PUSHPOP_RANGE, EVENT_TYPE_NVTX_STARTEND_RANGE]
            )
        )
    ]


def combine_text_fields(nvtx_df, str_df):
    """Combine the 'text' and 'textId' fields of the NVTX dataframe.

    This function simplifies the lookup process for accessing the event
    message. The 'text' field corresponds to the NVTX event message passed
    through 'nvtxDomainRegisterString', while the 'textId' field corresponds
    to the other case. By merging these fields, we streamline the process of
    retrieving the message.
    """
    if not nvtx_df["textId"].notnull().any():
        return nvtx_df.copy()

    nvtx_textId_df = replace_id_with_value(nvtx_df, str_df, "textId")
    mask = ~nvtx_textId_df["textId"].isna()
    nvtx_textId_df.loc[mask, "text"] = nvtx_textId_df.loc[mask, "textId"]
    return nvtx_textId_df.drop(columns=["textId"])


def extract_pid(global_id_series):
    """Extract the PID from the global ID.

    Parameters
    ----------
    global_id_series: series
        Either the 'globalPid' or 'globalTid' column of the dataframe.
    """
    pid = (global_id_series >> np.array(24)) & 0x00FFFFFF
    return pid


def combine_api_gpu_dfs(runtime_df, gpu_df):
    """Combine the runtime dataframe and the gpu dataframes.

    This function merges all the dataframes based on the 'correlationId'. The
    'start' and 'end' columns of the GPU dataframes are renamed to 'gpu_start'
    and 'gpu_end'.
    """
    gpu_df["pid"] = extract_pid(gpu_df["globalPid"])
    gpu_df = gpu_df.rename(columns={"start": "gpu_start", "end": "gpu_end"})

    api_df = runtime_df.copy()
    api_df["pid"] = extract_pid(api_df["globalTid"])
    api_df = api_df.merge(gpu_df, on=["correlationId", "pid"], how="inner")

    # Both PID and globalPid can be retrieved from globalTid.
    return api_df.drop(columns=["pid", "globalPid"])