# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Performance regression pipeline utilities.

Generic regression detection and upload logic that can be reused across
different test scripts (perf sanity, module perf, accuracy, etc.).
"""

import json
import os
import re
import socket
from datetime import datetime
from typing import Dict, List, Optional

from defs.trt_test_alternative import print_info, print_warning

from .open_search_db_utils import add_id, get_history_data, post_new_perf_data

# Default threshold values for performance regression detection
POST_MERGE_THRESHOLD = 0.05
PRE_MERGE_THRESHOLD = 0.1

_URM_BASE = "https://urm.nvidia.com/artifactory/sw-tensorrt-generic/llm-artifacts/LLM/main"


def get_job_info():
    """Get job info from environment variables."""
    # Read environment variables
    host_node_name = os.getenv("HOST_NODE_NAME", "") or socket.gethostname()
    build_id = os.getenv("BUILD_ID", "")
    build_url = os.getenv("BUILD_URL", "")
    job_name = os.getenv("JOB_NAME", "")
    global_vars_str = os.getenv("globalVars", "{}")
    try:
        global_vars = json.loads(global_vars_str)
    except Exception:
        global_vars = {}

    # Get job_url and job_id
    job_url = ""
    job_id = ""
    parents = global_vars.get("action_info", {}).get("parents", [])
    if parents:
        last_parent = parents[-1]
        job_url = last_parent.get("url", "")
        job_id = str(last_parent.get("build_number", ""))

    # Determine job type from job_url
    is_post_merge = "PostMerge" in job_url
    is_pr_job = not is_post_merge

    raw_branch = global_vars.get("build_branch")
    branch = raw_branch if isinstance(raw_branch, str) else ""
    commit = os.getenv("gitlabCommit", "")

    # Initialize PR-specific fields
    trigger_mr_user = ""
    trigger_mr_link = ""
    trigger_mr_id = ""
    trigger_mr_commit = ""
    artifact_url = ""
    if is_pr_job:
        # Get PR info from github_pr_api_url
        github_pr_api_url = global_vars.get("github_pr_api_url", "")
        if github_pr_api_url:
            # Extract PR ID from URL like "https://api.github.com/repos/NVIDIA/TensorRT-LLM/pulls/xxxx"
            pr_match = re.search(r"/pulls/(\d+)", github_pr_api_url)
            if pr_match:
                trigger_mr_id = pr_match.group(1)
                # Convert API URL to web URL
                trigger_mr_link = github_pr_api_url.replace(
                    "api.github.com/repos/", "github.com/"
                ).replace("/pulls/", "/pull/")

        # Extract user from trigger_info
        # Pattern: target="_blank">Barry-Delaney</a><br/>Git Commit:
        trigger_info = global_vars.get("action_info", {}).get("trigger_info", "")
        # Try to extract username from patterns like 'target="_blank">username</a><br/>Git Commit:'
        user_match = re.search(r'target="_blank">([^<]+)</a><br/>Git Commit:', trigger_info)
        if user_match:
            trigger_mr_user = user_match.group(1)

        # Set trigger_mr_commit to commit
        trigger_mr_commit = commit
        artifact_url = f"{_URM_BASE}/L0_MergeRequest_PR/{job_id}" if job_id else ""
    else:
        artifact_url = f"{_URM_BASE}/L0_PostMerge/{job_id}" if job_id else ""

    return {
        "b_is_baseline": False,
        "b_is_valid": True,
        # Unique identifier
        "_id": "",
        # Job Config
        "s_host_node_name": host_node_name,
        "s_build_id": build_id,
        "s_build_url": build_url,
        "s_job_name": job_name,
        "s_job_id": job_id,
        "s_job_url": job_url,
        "s_branch": branch,
        "s_commit": commit,
        "b_is_post_merge": is_post_merge,
        "b_is_pr_job": is_pr_job,
        "s_trigger_mr_user": trigger_mr_user,
        "s_trigger_mr_link": trigger_mr_link,
        "s_trigger_mr_id": trigger_mr_id,
        "s_trigger_mr_commit": trigger_mr_commit,
        "s_artifact_url": artifact_url,
        "b_is_regression": False,
    }


def get_common_values(new_data_dict, match_keys):
    """Find keys from match_keys where all entries have identical values.

    Returns a dict with those common key-value pairs.
    Skips entries that don't have the key or have None/empty values.
    """
    if not new_data_dict or not match_keys:
        return {}

    data_list = list(new_data_dict.values())
    if not data_list:
        return {}

    common_values_dict = {}
    for key in match_keys:
        # Collect non-None, non-empty values for this key
        values = []
        for data in data_list:
            if key in data and data[key] is not None:
                values.append(data[key])

        # Skip if no valid values found
        if len(values) != len(data_list):
            continue

        # Check if all valid values are identical
        first_value = values[0]
        if all(v == first_value for v in values):
            common_values_dict[key] = first_value

    return common_values_dict


def _rolling_smooth(values, window=3):
    """Trailing rolling mean with same-length output.

    Early elements use fewer samples (i.e. the first element is itself,
    the second is the mean of the first two, etc.).
    """
    if not values:
        return []
    smoothed = []
    for i in range(len(values)):
        start = max(0, i - window + 1)
        w = values[start : i + 1]
        smoothed.append(sum(w) / len(w))
    return smoothed


def _percentile(values, p):
    """Compute the p-th percentile with linear interpolation."""
    if not values:
        return 0.0
    s = sorted(values)
    k = (p / 100.0) * (len(s) - 1)
    lo = int(k)
    hi = min(lo + 1, len(s) - 1)
    frac = k - lo
    return s[lo] + frac * (s[hi] - s[lo])


def _created_ms(data):
    """``ts_created`` of a history record in epoch milliseconds, or None.

    OpenSearch returns ``ts_created`` as epoch milliseconds; an epoch in seconds
    is accepted as well. Anything else (absent, ISO string, malformed) is None.
    """
    ts = data.get("ts_created")
    if isinstance(ts, bool) or not isinstance(ts, (int, float)):
        return None
    return ts if ts > 1e12 else ts * 1000


def _stored_baseline_predates_floor(baseline_threshold, metric, metric_history_floors_ms):
    """Whether a stored d_baseline_* for *metric* may predate the metric's history floor.

    The stored fields come from one history record (see get_history_data), whose
    ts_created says when they were set. A metric with no floor is never affected;
    under a floor, a record of unknown age counts as older than the floor.
    """
    floor_ms = (metric_history_floors_ms or {}).get(metric)
    if floor_ms is None:
        return False
    stored_ms = _created_ms(baseline_threshold or {})
    return stored_ms is None or stored_ms < floor_ms


def _daily_aggregate_values(data_list, metric, min_created_ms=None):
    """Aggregate multiple data points on the same day to a single mean value.

    When min_created_ms is given, records created before it are skipped, and so
    are records whose creation time is unknown (see _created_ms). The record
    being checked never reaches here with a ts_created, so it is unaffected by
    construction.

    Returns a list of daily-aggregated metric values sorted by date.
    """
    by_day = {}
    for data in data_list:
        if data.get("b_is_baseline"):
            continue
        val = data.get(metric)
        if val is None:
            continue
        if min_created_ms is not None:
            created_ms = _created_ms(data)
            if created_ms is None or created_ms < min_created_ms:
                continue
        ts = data.get("ts_created") or data.get("@timestamp")
        if ts is None:
            continue
        if isinstance(ts, (int, float)):
            if ts > 1e12:
                ts = ts / 1000
            day_key = datetime.fromtimestamp(ts).strftime("%Y-%m-%d")
        elif isinstance(ts, datetime):
            day_key = ts.strftime("%Y-%m-%d")
        else:
            day_key = str(ts)[:10]
        by_day.setdefault(day_key, []).append(val)
    result = []
    for day in sorted(by_day):
        vals = by_day[day]
        result.append(sum(vals) / len(vals))
    return result


def calculate_baseline_metrics(
    history_data_list,
    new_data,
    maximize_metrics,
    minimize_metrics,
    metric_history_floors_ms=None,
):
    """Calculate baseline metrics using rolling smooth + percentile algorithm.

    For each metric, aggregates data to daily values, applies a trailing
    rolling mean (window=3), then takes:
      - P95 for maximize_metrics (larger is better, e.g. throughput)
      - P5  for minimize_metrics (smaller is better, e.g. latency)

    metric_history_floors_ms optionally maps a metric to an epoch-millisecond
    ts_created floor: that metric's baseline is computed only from history
    records created at or after the floor, so values recorded under an older
    definition of the metric cannot set it. A metric with no history left
    above its floor gets no baseline, which the caller treats as "nothing to
    compare against yet".
    """
    floors = metric_history_floors_ms or {}
    all_data = []
    if history_data_list:
        all_data.extend(history_data_list)
    if isinstance(new_data, list):
        all_data.extend(new_data)
    elif new_data:
        all_data.append(new_data)

    if not all_data:
        return {}

    baseline_metrics = {}
    for metric in maximize_metrics + minimize_metrics:
        daily_vals = _daily_aggregate_values(all_data, metric, floors.get(metric))
        if not daily_vals:
            continue
        smoothed = _rolling_smooth(daily_vals, window=3)
        if metric in maximize_metrics:
            baseline_metrics[metric] = _percentile(smoothed, 95)
        else:
            baseline_metrics[metric] = _percentile(smoothed, 5)

    return baseline_metrics


def _calculate_diff(metric, new_value, baseline_value, maximize_metrics):
    """Calculate the percentage difference between new and baseline values.

    Returns a positive number if perf is better, negative if worse.
    """
    if baseline_value == 0:
        return 0.0
    if metric in maximize_metrics:
        # Larger is better: new > baseline means positive (better)
        return (new_value - baseline_value) / baseline_value * 100
    else:
        # Smaller is better: new < baseline means positive (better)
        return (baseline_value - new_value) / baseline_value * 100


def _is_regressive(value, baseline_value, threshold, maximize):
    """Return True when *value* misses *baseline_value* by more than *threshold*.

    Shared by the pre-merge verdict and by the check on main's own latest value,
    so the two can never drift apart. The comparison is strict: a value sitting
    exactly on baseline_value * (1 -/+ threshold) is not a regression.
    """
    if maximize:
        return value < baseline_value * (1 - threshold)
    return value > baseline_value * (1 + threshold)


def prepare_regressive_test_cases(
    latest_history_data_dict,
    latest_baseline_threshold_dict,
    history_data_dict,
    new_data_dict,
    maximize_metrics,
    minimize_metrics,
    regression_metrics,
    metric_history_floors_ms=None,
):
    """Update regression info for all data in new_data_dict.

    Uses baseline/threshold fields from latest_baseline_threshold_dict when
    available, otherwise falls back to calculating baseline from history data.
    For a metric listed in metric_history_floors_ms, a stored baseline is used only
    when the record it came from (its ts_created) is at or after the metric's floor;
    otherwise the floor-filtered calculation is used (see calculate_baseline_metrics).
    Its threshold is read from latest_baseline_threshold_dict either way.

    Returns the set of pre-merge cmd_idx whose regression must not fail the
    stage because the latest post-merge record for the same case is itself
    beyond the pre-merge threshold. Both pipelines compare against one shared
    baseline, so a regression landed on main would otherwise make every
    subsequent PR measure the same regressed value and fail for a change it did
    not introduce.

    The latest post-merge value is re-evaluated here against the same baseline
    and the same pre-merge threshold, rather than reusing the b_is_regression
    the post-merge run recorded at its own tighter threshold. That makes the
    exemption exactly as wide as the failure it prevents: if main sits within
    the pre-merge threshold, a PR reproducing main's value does not fail the
    gate, so there is nothing to exempt and the gate stays armed.
    """
    exempt_cmd_idxs = set()

    # If latest_history_data_dict is None (network failure), skip regression check
    if latest_history_data_dict is None:
        return exempt_cmd_idxs

    for cmd_idx in new_data_dict:
        new_data = new_data_dict[cmd_idx]
        baseline_threshold = (
            latest_baseline_threshold_dict.get(cmd_idx) if latest_baseline_threshold_dict else None
        )

        is_post_merge = new_data.get("b_is_post_merge", False)
        regressive_metrics = []
        info_lines = []

        # The latest post-merge record for this same case, used only to decide
        # whether a pre-merge regression predates the change under test.
        latest_post_merge = None if is_post_merge else latest_history_data_dict.get(cmd_idx)
        if not isinstance(latest_post_merge, dict):
            latest_post_merge = None
        post_merge_regressive_metrics = []

        # Pre-calculate fallback baseline from history if needed
        fallback_baseline = None

        # Check all metrics and build info string
        for metric in maximize_metrics + minimize_metrics:
            if metric not in new_data:
                continue

            new_value = new_data[metric]
            metric_suffix = metric[2:]  # Remove "d_" prefix

            # Get baseline value: try dedicated baseline/threshold dict first,
            # then fall back to calculating from history data
            baseline_key = f"d_baseline_{metric_suffix}"
            baseline_value = baseline_threshold.get(baseline_key) if baseline_threshold else None
            # A stored baseline recorded before a metric's floor predates the
            # metric's current definition, so the floored metric recomputes its
            # baseline from the floor-filtered history instead. One recorded at or
            # after the floor (or with no floor) is honoured as before.
            if baseline_value is not None and _stored_baseline_predates_floor(
                baseline_threshold, metric, metric_history_floors_ms
            ):
                baseline_value = None
            if baseline_value is None or baseline_value <= 0:
                if fallback_baseline is None:
                    history_list = history_data_dict.get(cmd_idx, [])
                    fallback_baseline = calculate_baseline_metrics(
                        history_list,
                        new_data,
                        maximize_metrics,
                        minimize_metrics,
                        metric_history_floors_ms=metric_history_floors_ms,
                    )
                baseline_value = fallback_baseline.get(metric)
                if baseline_value is None or baseline_value <= 0:
                    continue

            # Get threshold: try dedicated baseline/threshold dict first,
            # then fall back to default threshold
            if is_post_merge:
                threshold_key = f"d_threshold_post_merge_{metric_suffix}"
                default_threshold = POST_MERGE_THRESHOLD
            else:
                threshold_key = f"d_threshold_pre_merge_{metric_suffix}"
                default_threshold = PRE_MERGE_THRESHOLD

            threshold = baseline_threshold.get(threshold_key) if baseline_threshold else None
            if threshold is None or threshold <= 0:
                threshold = default_threshold

            diff = _calculate_diff(metric, new_value, baseline_value, maximize_metrics)
            info_lines.append(
                f"  {metric}: value={new_value:.4f} "
                f"baseline={baseline_value:.4f} "
                f"threshold={threshold * 100:.2f}% "
                f"diff={diff:+.2f}%"
            )

            # Check if this metric is regressive (only for key regression metrics)
            if metric in regression_metrics:
                maximize = metric in maximize_metrics
                if _is_regressive(new_value, baseline_value, threshold, maximize):
                    regressive_metrics.append(metric)

                # Re-evaluate main's own latest value against this same baseline
                # and this same (pre-merge) threshold. Anything unusable -- field
                # absent, null, non-numeric, non-positive -- must not exempt, so
                # a broken post-merge record leaves the gate armed.
                if latest_post_merge is not None:
                    prior_value = _safe_float(latest_post_merge.get(metric), None)
                    if (
                        prior_value is not None
                        and prior_value > 0
                        and _is_regressive(prior_value, baseline_value, threshold, maximize)
                    ):
                        post_merge_regressive_metrics.append(metric)

        test_case = new_data.get("s_test_case_name", "unknown")
        header = f"Regression in {test_case}:"
        new_data["s_regression_info"] = "\n".join([header] + info_lines)
        new_data["b_is_regression"] = len(regressive_metrics) > 0

        # A pre-merge regression is not this PR's fault when main's own latest
        # value already misses the same gate. post_merge_regressive_metrics is
        # only ever populated on a pre-merge run, so post-merge is unaffected.
        if post_merge_regressive_metrics:
            exempt_cmd_idxs.add(cmd_idx)
            if new_data["b_is_regression"]:
                new_data["s_regression_info"] += (
                    "\n  Not failing this stage: the latest post-merge run of this test "
                    "case already misses the same threshold against the same baseline "
                    f"({', '.join(post_merge_regressive_metrics)}), so this regression "
                    "predates the change under test."
                )

    return exempt_cmd_idxs


def _safe_float(value, default):
    """Convert *value* to float, returning *default* on failure.

    Handles None, non-numeric strings, and other malformed legacy values
    that may appear in OpenSearch history records.
    """
    if value is None:
        return default
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def add_baseline_fields_to_post_merge_data(
    latest_baseline_threshold_dict,
    new_data_dict,
    maximize_metrics,
    minimize_metrics,
    metric_history_floors_ms=None,
):
    """Embed baseline fields directly into each post-merge data entry.

    For each metric, only sets fields when the inherited value exists and is > 0:
      - d_baseline_{metric_suffix}
      - d_threshold_post_merge_{metric_suffix}
      - d_threshold_pre_merge_{metric_suffix}

    A d_baseline_* that predates the metric's floor in metric_history_floors_ms is
    not carried forward: the gate no longer uses it, and copying it onto a new
    record would give it a current ts_created.
    """
    if latest_baseline_threshold_dict is None:
        return

    for cmd_idx in new_data_dict:
        new_data = new_data_dict[cmd_idx]
        baseline_threshold = latest_baseline_threshold_dict.get(cmd_idx)
        if baseline_threshold is None:
            continue

        for metric in maximize_metrics + minimize_metrics:
            metric_suffix = metric[2:]  # Remove "d_" prefix
            baseline_key = f"d_baseline_{metric_suffix}"
            post_merge_key = f"d_threshold_post_merge_{metric_suffix}"
            pre_merge_key = f"d_threshold_pre_merge_{metric_suffix}"

            # Only set fields when inherited value exists and is positive.
            # Use _safe_float to handle null/non-numeric legacy values.
            post_merge_val = _safe_float(baseline_threshold.get(post_merge_key), None)
            if post_merge_val is not None and post_merge_val > 0:
                new_data[post_merge_key] = post_merge_val

            pre_merge_val = _safe_float(baseline_threshold.get(pre_merge_key), None)
            if pre_merge_val is not None and pre_merge_val > 0:
                new_data[pre_merge_key] = pre_merge_val

            baseline_val = _safe_float(baseline_threshold.get(baseline_key), None)
            if (
                baseline_val is not None
                and baseline_val > 0
                and not _stored_baseline_predates_floor(
                    baseline_threshold, metric, metric_history_floors_ms
                )
            ):
                new_data[baseline_key] = baseline_val


def check_perf_regression(new_data_dict, fail_on_regression=False, exempt_cmd_idxs=None):
    """Check performance regression by printing s_regression_info.

    Post-merge regressions log warnings. Pre-merge
    regressions raise RuntimeError when fail_on_regression is True.

    Cases in exempt_cmd_idxs are reported as warnings but never raise: their
    regression is already present on main, so failing here would block a PR for
    a change it did not introduce. See prepare_regressive_test_cases.
    """
    exempt_cmd_idxs = exempt_cmd_idxs or set()

    regressive_data_items = [
        (cmd_idx, data)
        for cmd_idx, data in new_data_dict.items()
        if data.get("b_is_regression", False)
    ]

    if not regressive_data_items:
        print_info("No regression data found.")
        return

    post_merge_regressions = [
        data for _, data in regressive_data_items if data.get("b_is_post_merge", False)
    ]
    pre_merge_regressions = [
        (cmd_idx, data)
        for cmd_idx, data in regressive_data_items
        if not data.get("b_is_post_merge", False)
    ]

    # Print post-merge regression details as warnings
    for data in post_merge_regressions:
        print_warning(data.get("s_regression_info", ""))

    # Print pre-merge regression details and raise error. An exempt case is
    # still reported -- the measurement is real -- but withheld from the error.
    error_parts = []
    for cmd_idx, data in pre_merge_regressions:
        info = data.get("s_regression_info", "")
        print_warning(info)
        if cmd_idx not in exempt_cmd_idxs:
            error_parts.append(info)

    # Guard on error_parts, not on pre_merge_regressions: when every regressive
    # case is exempt there is nothing to report and nothing to fail.
    if fail_on_regression and error_parts:
        raise RuntimeError("\n".join(error_parts))


def process_and_upload_test_results(
    new_data_dict: Dict[int, dict],
    match_keys: List[str],
    maximize_metrics: List[str],
    minimize_metrics: List[str],
    regression_metrics: List[str],
    extra_fields: Optional[dict] = None,
    upload_to_db: bool = True,
    fail_on_regression: Optional[bool] = None,
    metric_history_floors_ms: Optional[Dict[str, int]] = None,
):
    """Generic regression pipeline: enrich data, query history, detect regression, upload.

    Args:
        new_data_dict: cmd_idx -> data dict with test config + metric fields.
        match_keys: Fields that uniquely identify a test case for history matching.
        maximize_metrics: Metrics where larger is better (e.g. throughput).
        minimize_metrics: Metrics where smaller is better (e.g. latency).
        regression_metrics: Subset of metrics to check for regression.
        extra_fields: Additional fields to merge into each data entry
            (e.g. {"s_stage_name": ..., "s_test_list": ...}).
        upload_to_db: Whether to post results to the database.
        fail_on_regression: Whether to raise on pre-merge regression.
            None = auto (fail for pre-merge only).
        metric_history_floors_ms: Optional metric -> epoch-ms ts_created floor.
            History created before a metric's floor is ignored when computing
            that metric's baseline (see calculate_baseline_metrics).
    """
    if not new_data_dict:
        print_info("No data to upload to database.")
        return

    # Validate that regression_metrics is a subset of all metrics
    all_metrics = set(maximize_metrics) | set(minimize_metrics)
    invalid_metrics = set(regression_metrics) - all_metrics
    if invalid_metrics:
        raise ValueError(
            f"regression_metrics {sorted(invalid_metrics)} are not in "
            f"maximize_metrics or minimize_metrics"
        )

    # Step 1: Get job config and determine merge type
    job_config = get_job_info()
    is_post_merge = job_config["b_is_post_merge"]

    # Step 2: Enrich each entry with job config, extra fields, and ID
    for cmd_idx in new_data_dict:
        data = new_data_dict[cmd_idx]
        data.update(job_config)
        if extra_fields:
            data.update(extra_fields)
        add_id(data)

    # Step 3: For pre-merge, look history up against the baseline branch
    lookup_data_dict = new_data_dict
    if not is_post_merge and "s_branch" in match_keys:
        baseline_branch = os.environ.get("PERF_BASELINE_BRANCH", "main")
        lookup_data_dict = {
            cmd_idx: {**data, "s_branch": baseline_branch}
            for cmd_idx, data in new_data_dict.items()
        }

    # Step 4: Find common values to narrow query scope
    common_values_dict = get_common_values(lookup_data_dict, match_keys)

    # Step 5: Query history data
    latest_history_data_dict, latest_baseline_threshold_dict, history_data_dict = get_history_data(
        lookup_data_dict, match_keys, common_values_dict
    )

    # Step 6: Compute regression info. Cases whose regression is already present
    # on the latest post-merge run are reported but must not fail the stage.
    exempt_cmd_idxs = prepare_regressive_test_cases(
        latest_history_data_dict,
        latest_baseline_threshold_dict,
        history_data_dict,
        new_data_dict,
        maximize_metrics,
        minimize_metrics,
        regression_metrics,
        metric_history_floors_ms=metric_history_floors_ms,
    )

    # Step 7: For post-merge, embed baseline fields
    if is_post_merge:
        add_baseline_fields_to_post_merge_data(
            latest_baseline_threshold_dict,
            new_data_dict,
            maximize_metrics,
            minimize_metrics,
            metric_history_floors_ms=metric_history_floors_ms,
        )

    # Step 8: Upload to DB
    if upload_to_db:
        post_new_perf_data(new_data_dict)

    # Step 9: Check regression (auto-detect fail behavior if not specified)
    if fail_on_regression is None:
        fail_on_regression = not is_post_merge
    check_perf_regression(
        new_data_dict,
        fail_on_regression=fail_on_regression,
        exempt_cmd_idxs=exempt_cmd_idxs,
    )