"""Deterministic kernel matching using CUDA graph execution order.

This module provides 98-99% deterministic matching by leveraging the fact that
both AMD and NVIDIA traces execute CUDA graphs in identical order, and kernels
within each graph execute in deterministic timestamp order.
"""

from dataclasses import dataclass
from typing import Any

import orjson


@dataclass
class KernelMatch:
    """A matched pair of kernels from AMD and NVIDIA traces."""

    graph_index: int  # Which graph execution this belongs to (0-184)
    position_in_graph: int  # Position within the graph (0-based)

    amd_kernel: dict[str, Any]
    nvidia_kernel: dict[str, Any]

    operation_type: str  # GEMM, ATTN, RMS, etc.
    confidence: float  # 1.0 = perfect match, <1.0 = potential fusion difference

    # For debugging/validation
    amd_correlation: int
    nvidia_correlation: int


@dataclass
class GraphPair:
    """A pair of matched CUDA graphs from AMD and NVIDIA traces."""

    graph_index: int
    amd_correlation: int
    nvidia_correlation: int

    amd_kernels: list[dict[str, Any]]
    nvidia_kernels: list[dict[str, Any]]

    is_layer: bool  # True if this is a transformer layer (>100 kernels)


def classify_kernel(name: str) -> str:
    """Classify kernel by operation type.

    This is a coarse classification for matching purposes.
    """
    nl = name.lower()

    if 'cijk_' in nl or 'nvjet' in nl:
        return 'GEMM'
    elif 'attention' in nl or 'fmha' in nl:
        return 'ATTN'
    elif 'reshape_and_cache' in nl:
        return 'KV'
    elif 'triton_per' in nl and 'rsqrt' in nl:
        return 'RMS'
    elif 'triton_poi' in nl and 'silu' in nl:
        return 'SILU'
    elif 'triton_poi' in nl:
        return 'POI'
    elif 'triton_red' in nl:
        return 'RED'
    elif 'reduce_segments' in nl:
        return 'RSEG'
    else:
        return 'OTH'


def is_platform_specific_kernel(name: str, platform: str) -> bool:
    """Check if kernel is platform-specific and should be excluded from matching.

    AMD runs reduce_segments after attention operations, but NVIDIA fuses this
    into adjacent kernels. We need to filter these out for accurate matching.
    """
    if platform == "AMD":
        return 'reduce_segments' in name.lower()

    # Add NVIDIA-specific exclusions here if discovered
    return False


def load_graph_execution_order(trace_path: str) -> list[tuple[int, int]]:
    """Load CUDA graph execution order from trace.

    Returns:
        List of (timestamp, correlation_id) tuples in execution order
    """
    with open(trace_path, "rb") as f:
        trace = orjson.loads(f.read())

    graph_launches = []
    for event in trace.get("traceEvents", []):
        if event.get("cat") == "cuda_runtime":
            name = event.get("name", "")
            if "GraphLaunch" in name or "graphLaunch" in name.lower():
                ts = event.get("ts")
                corr_id = event.get("args", {}).get("correlation")
                if ts is not None and corr_id is not None:
                    graph_launches.append((ts, corr_id))

    # Sort by timestamp to get execution order
    graph_launches.sort()
    return graph_launches


def load_kernels_for_correlation(trace_path: str, correlation_id: int, platform: str) -> list[dict[str, Any]]:
    """Load all kernels for a given correlation ID in timestamp order.

    Args:
        trace_path: Path to trace JSON
        correlation_id: Correlation ID to filter by
        platform: "AMD" or "NVIDIA" for platform-specific filtering

    Returns:
        List of kernel events sorted by timestamp, with platform-specific kernels removed
    """
    with open(trace_path, "rb") as f:
        trace = orjson.loads(f.read())

    kernels = []
    for event in trace.get("traceEvents", []):
        if event.get("cat") == "kernel":
            corr_id = event.get("args", {}).get("correlation")
            if corr_id == correlation_id:
                name = event.get("name", "")

                # Skip platform-specific kernels
                if is_platform_specific_kernel(name, platform):
                    continue

                kernels.append({
                    "name": name,
                    "ts": event.get("ts"),
                    "dur": event.get("dur", 0),
                    "correlation": corr_id,
                    "args": event.get("args", {}),
                })

    # Sort by timestamp for deterministic ordering
    kernels.sort(key=lambda k: k["ts"])
    return kernels


def match_traces(
    amd_trace_path: str,
    nvidia_trace_path: str
) -> tuple[list[GraphPair], list[KernelMatch]]:
    """Match kernels between AMD and NVIDIA traces using graph execution order.

    This provides 98-99% deterministic matching by:
    1. Matching graphs by execution order (100% deterministic)
    2. Matching kernels by position within graphs (98-99% deterministic)
    3. Filtering platform-specific operations (e.g., AMD's reduce_segments)

    Args:
        amd_trace_path: Path to AMD trace JSON
        nvidia_trace_path: Path to NVIDIA trace JSON

    Returns:
        Tuple of (graph_pairs, kernel_matches)
        - graph_pairs: List of matched CUDA graph pairs
        - kernel_matches: List of all kernel matches across all graphs
    """
    # Step 1: Get graph execution order from both traces
    amd_graphs = load_graph_execution_order(amd_trace_path)
    nvidia_graphs = load_graph_execution_order(nvidia_trace_path)

    if len(amd_graphs) != len(nvidia_graphs):
        raise ValueError(
            f"Graph count mismatch: AMD has {len(amd_graphs)} graphs, "
            f"NVIDIA has {len(nvidia_graphs)} graphs. "
            "Traces may be from different workloads."
        )

    graph_pairs = []
    kernel_matches = []

    # Step 2: Match graphs by execution order
    for graph_idx, ((amd_ts, amd_corr), (nv_ts, nv_corr)) in enumerate(zip(amd_graphs, nvidia_graphs)):
        # Load kernels for this correlation
        amd_kernels = load_kernels_for_correlation(amd_trace_path, amd_corr, "AMD")
        nvidia_kernels = load_kernels_for_correlation(nvidia_trace_path, nv_corr, "NVIDIA")

        is_layer = len(amd_kernels) > 100 or len(nvidia_kernels) > 100

        graph_pairs.append(GraphPair(
            graph_index=graph_idx,
            amd_correlation=amd_corr,
            nvidia_correlation=nv_corr,
            amd_kernels=amd_kernels,
            nvidia_kernels=nvidia_kernels,
            is_layer=is_layer,
        ))

        # Step 3: Match kernels within this graph by position
        matches = match_kernels_in_graph(
            graph_idx=graph_idx,
            amd_corr=amd_corr,
            nvidia_corr=nv_corr,
            amd_kernels=amd_kernels,
            nvidia_kernels=nvidia_kernels,
        )
        kernel_matches.extend(matches)

    return graph_pairs, kernel_matches


def match_kernels_in_graph(
    graph_idx: int,
    amd_corr: int,
    nvidia_corr: int,
    amd_kernels: list[dict[str, Any]],
    nvidia_kernels: list[dict[str, Any]],
) -> list[KernelMatch]:
    """Match kernels within a single CUDA graph by position.

    Args:
        graph_idx: Index of this graph in execution order
        amd_corr: AMD correlation ID
        nvidia_corr: NVIDIA correlation ID
        amd_kernels: AMD kernels (already sorted by timestamp, filtered for platform-specific ops)
        nvidia_kernels: NVIDIA kernels (already sorted by timestamp, filtered for platform-specific ops)

    Returns:
        List of kernel matches with confidence scores
    """
    matches = []

    # If kernel counts don't match after filtering, something unexpected happened
    if len(amd_kernels) != len(nvidia_kernels):
        # Handle gracefully: match what we can
        min_len = min(len(amd_kernels), len(nvidia_kernels))

        for i in range(min_len):
            amd_k = amd_kernels[i]
            nv_k = nvidia_kernels[i]

            amd_type = classify_kernel(amd_k["name"])
            nv_type = classify_kernel(nv_k["name"])

            # Lower confidence if operation types don't match
            confidence = 1.0 if amd_type == nv_type else 0.5

            matches.append(KernelMatch(
                graph_index=graph_idx,
                position_in_graph=i,
                amd_kernel=amd_k,
                nvidia_kernel=nv_k,
                operation_type=amd_type,
                confidence=confidence,
                amd_correlation=amd_corr,
                nvidia_correlation=nvidia_corr,
            ))

        # Note: Unmatched kernels are implicitly dropped
        # Could add logging here for debugging
    else:
        # Perfect length match - match by position
        for i, (amd_k, nv_k) in enumerate(zip(amd_kernels, nvidia_kernels)):
            amd_type = classify_kernel(amd_k["name"])
            nv_type = classify_kernel(nv_k["name"])

            # Confidence = 1.0 if operation types match, else 0.8
            # (0.8 because position-based matching is still very reliable)
            confidence = 1.0 if amd_type == nv_type else 0.8

            matches.append(KernelMatch(
                graph_index=graph_idx,
                position_in_graph=i,
                amd_kernel=amd_k,
                nvidia_kernel=nv_k,
                operation_type=amd_type if amd_type == nv_type else f"{amd_type}→{nv_type}",
                confidence=confidence,
                amd_correlation=amd_corr,
                nvidia_correlation=nvidia_corr,
            ))

    return matches


def get_matching_statistics(kernel_matches: list[KernelMatch]) -> dict[str, Any]:
    """Calculate statistics about matching quality.

    Returns:
        Dict with:
        - total_matches: Total kernel pairs matched
        - perfect_matches: Matches with confidence=1.0
        - fuzzy_matches: Matches with confidence<1.0
        - match_rate: Percentage of perfect matches
        - by_operation: Breakdown by operation type
    """
    total = len(kernel_matches)
    perfect = sum(1 for m in kernel_matches if m.confidence == 1.0)

    # Breakdown by operation type
    from collections import defaultdict
    by_operation = defaultdict(lambda: {"total": 0, "perfect": 0})

    for match in kernel_matches:
        op = match.operation_type
        by_operation[op]["total"] += 1
        if match.confidence == 1.0:
            by_operation[op]["perfect"] += 1

    return {
        "total_matches": total,
        "perfect_matches": perfect,
        "fuzzy_matches": total - perfect,
        "match_rate": perfect / total if total > 0 else 0.0,
        "by_operation": dict(by_operation),
    }
