# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import TypedDict

# --- Type Aliases ---
# "{data_type[::phi::dtype::bfloat16]; data_layout[STRIDED]; place[Place(gpu:0)]; library_type[PLAIN]}"
KernelConfig = dict[str, str]
# "fused_transpose": [
#     "{data_type[uint8_t]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
#     "{data_type[int8_t]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
#     "{data_type[::phi::dtype::bfloat16]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
#     "{data_type[float]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}"
# ]
KernelManifest = dict[str, list[KernelConfig]]


class GpuKernelChangeSummary(TypedDict):
    # New kernel names that contain "gpu".
    new_kernels_with_gpu: list[str]
    # Existing kernels that have new GPU configurations added.
    kernels_with_new_gpu_support: list[str]
    # Existing GPU kernels that ignore data_type changes.
    gpu_kernels_with_new_data_types: dict[str, list[str]]


class KernelManifestComparer:
    def __init__(self, baseline_manifest_data: dict[str, list[str]]) -> None:
        self.baseline_manifest: KernelManifest = self._process_raw_manifest(
            baseline_manifest_data
        )
        self._preprocessed_baseline = {
            kernel_name: {
                "configs_set": {tuple(sorted(cfg.items())) for cfg in configs},
                "gpu_signatures": {
                    self._create_config_signature_without_dtype(cfg)
                    for cfg in configs
                    if "gpu" in cfg.get("place", "")
                },
            }
            for kernel_name, configs in self.baseline_manifest.items()
        }

    def _parse_config_string(self, config_str: str) -> KernelConfig:
        pattern = r"([\w:]+)\[([^\]]+)\]"
        matches = re.findall(pattern, config_str)
        return dict(matches)

    def _create_config_signature_without_dtype(
        self, config: KernelConfig
    ) -> tuple[tuple[str, str], ...]:
        signature_items = {
            key: value for key, value in config.items() if key != "data_type"
        }
        return tuple(sorted(signature_items.items()))

    def _process_raw_manifest(
        self, raw_manifest: dict[str, list[str]]
    ) -> KernelManifest:
        processed_manifest: KernelManifest = {}
        for kernel_name, config_strings in raw_manifest.items():
            processed_manifest[kernel_name] = [
                self._parse_config_string(cs) for cs in config_strings
            ]
        return processed_manifest

    def compare(
        self, target_manifest_data: dict[str, list[str]]
    ) -> GpuKernelChangeSummary:
        target_manifest = self._process_raw_manifest(target_manifest_data)

        new_kernels_with_gpu: set[str] = set()
        kernels_with_new_gpu_support: set[str] = set()
        gpu_kernels_with_new_data_types: dict[str, set[str]] = {}

        baseline_kernel_names = set(self._preprocessed_baseline.keys())
        target_kernel_names = set(target_manifest.keys())

        # 1. Find newly added kernels that are GPU-related.
        added_kernel_names = target_kernel_names - baseline_kernel_names
        for kernel_name in added_kernel_names:
            if any(
                "gpu" in cfg.get("place", "")
                for cfg in target_manifest[kernel_name]
            ):
                new_kernels_with_gpu.add(kernel_name)

        # 2. Find changes within existing kernels.
        common_kernel_names = baseline_kernel_names.intersection(
            target_kernel_names
        )
        for kernel_name in common_kernel_names:
            baseline_data = self._preprocessed_baseline[kernel_name]
            target_configs = target_manifest[kernel_name]

            target_configs_set = {
                tuple(sorted(cfg.items())) for cfg in target_configs
            }

            if baseline_data["configs_set"] == target_configs_set:
                continue

            added_config_tuples = (
                target_configs_set - baseline_data["configs_set"]
            )
            added_configs = [dict(t) for t in added_config_tuples]

            baseline_gpu_signatures = baseline_data["gpu_signatures"]
            has_baseline_gpu_support = bool(baseline_gpu_signatures)

            for added_cfg in added_configs:
                if "gpu" in added_cfg.get("place", ""):
                    kernels_with_new_gpu_support.add(kernel_name)
                    added_signature = (
                        self._create_config_signature_without_dtype(added_cfg)
                    )
                    if (
                        has_baseline_gpu_support
                        and added_signature in baseline_gpu_signatures
                    ):
                        new_data_type = added_cfg.get("data_type", "N/A")
                        gpu_kernels_with_new_data_types.setdefault(
                            kernel_name, set()
                        ).add(new_data_type)

        return {
            "new_kernels_with_gpu": sorted(new_kernels_with_gpu),
            "kernels_with_new_gpu_support": sorted(
                kernels_with_new_gpu_support
            ),
            "gpu_kernels_with_new_data_types": {
                k: sorted(v) for k, v in gpu_kernels_with_new_data_types.items()
            },
        }


def cli():
    parser = argparse.ArgumentParser(
        description="Compare GPU kernel configurations between two manifests."
    )
    parser.add_argument(
        "baseline_manifest",
        type=Path,
        help="Path to the baseline manifest JSON file.",
    )
    parser.add_argument(
        "target_manifest",
        type=Path,
        help="Path to the target manifest JSON file.",
    )
    parser.add_argument(
        "--ignore-data-type-changes",
        action="store_true",
        help="If set, ignore data_type changes in GPU kernel comparisons.",
    )
    args = parser.parse_args()
    return args


def main():
    args = cli()
    if not args.baseline_manifest.exists():
        raise ValueError(
            f"Baseline manifest file not found: {args.baseline_manifest.resolve()}"
        )
    if not args.target_manifest.exists():
        raise ValueError(
            f"Target manifest file not found: {args.target_manifest.resolve()}"
        )

    with args.baseline_manifest.open("r", encoding="utf-8") as f:
        baseline_data = json.load(f)

    with args.target_manifest.open("r", encoding="utf-8") as f:
        target_data = json.load(f)

    comparer = KernelManifestComparer(baseline_data)
    summary = comparer.compare(target_data)

    has_reportable_changes = False
    if args.ignore_data_type_changes:
        all_gpu_related_changes = set(summary["new_kernels_with_gpu"]) | set(
            summary["kernels_with_new_gpu_support"]
        )
        kernels_with_only_dtype_changes = set(
            summary["gpu_kernels_with_new_data_types"].keys()
        )
        unignored_changes = (
            all_gpu_related_changes - kernels_with_only_dtype_changes
        )
        if unignored_changes:
            has_reportable_changes = True
    else:
        if (
            summary["new_kernels_with_gpu"]
            or summary["kernels_with_new_gpu_support"]
        ):
            has_reportable_changes = True

    if summary["new_kernels_with_gpu"]:
        print("New GPU Kernels Added:")
        for kernel in summary["new_kernels_with_gpu"]:
            print(f"  - {kernel}")

    if summary["kernels_with_new_gpu_support"]:
        print("\nKernels with New GPU Support:")
        for kernel in summary["kernels_with_new_gpu_support"]:
            print(f"  - {kernel}")

    if summary["gpu_kernels_with_new_data_types"]:
        print("\nGPU Kernels with new Data Type:")
        for kernel, data_types in summary[
            "gpu_kernels_with_new_data_types"
        ].items():
            data_types_str = ", ".join(data_types)
            print(f"  - {kernel}: New data types - {data_types_str}")

    if has_reportable_changes:
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()