"""Improved kernel matching that preserves fusion information.

Key improvements over v1:
1. Uses existing classifier.py instead of reimplementing
2. Marks fusion differences instead of filtering them out
3. Provides detailed fusion analysis
4. Handles sequence alignment when platforms have different kernel counts
"""

from dataclasses import dataclass, field
from typing import Any

import orjson

from .classifier import classify


@dataclass
class FusionDifference:
    """A fusion difference between platforms."""

    platform_with_kernel: str  # "AMD" or "NVIDIA"
    kernel_name: str
    operation_type: str
    position: int
    likely_fused_into: str | None = None  # Best guess of where this work went


@dataclass
class KernelMatch:
    """A matched pair of kernels, or a fusion difference."""

    graph_index: int
    amd_position: int | None  # None if this is NVIDIA-only
    nvidia_position: int | None  # None if this is AMD-only

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

    operation_type: str
    confidence: float  # 1.0 = perfect, 0.5 = fusion difference

    # If this is a fusion difference
    is_fusion_difference: bool = False
    fusion_info: FusionDifference | None = None

    amd_correlation: int | None = None
    nvidia_correlation: int | None = None


@dataclass
class GraphPair:
    """Matched CUDA graph pair with fusion analysis."""

    graph_index: int
    amd_correlation: int
    nvidia_correlation: int

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

    is_layer: bool
    fusion_differences: list[FusionDifference] = field(default_factory=list)


def load_graph_execution_order(trace_path: str) -> list[tuple[int, int]]:
    """Load CUDA graph 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))

    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 correlation, keeping ALL kernels including fusion differences."""
    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:
                kernels.append({
                    "name": event.get("name", ""),
                    "ts": event.get("ts"),
                    "dur": event.get("dur", 0),
                    "correlation": corr_id,
                    "args": event.get("args", {}),
                })

    kernels.sort(key=lambda k: k["ts"])
    return kernels


def align_sequences_with_fusion(
    amd_kernels: list[dict[str, Any]],
    nvidia_kernels: list[dict[str, Any]],
    platform_amd: str = "AMD",
    platform_nvidia: str = "NVIDIA",
) -> list[tuple[int | None, int | None, str]]:
    """Align two kernel sequences, identifying fusion differences.

    Returns:
        List of (amd_index, nvidia_index, alignment_type) where:
        - alignment_type is "match", "amd_only", or "nvidia_only"
    """
    # Classify all kernels
    amd_ops = [classify(k["name"], platform_amd)[0].value for k in amd_kernels]
    nvidia_ops = [classify(k["name"], platform_nvidia)[0].value for k in nvidia_kernels]

    # Use simple sequence alignment
    # For now, handle the common case: AMD has extra "reduce_segments" operations
    alignments = []
    amd_i = 0
    nv_i = 0

    while amd_i < len(amd_ops) or nv_i < len(nvidia_ops):
        if amd_i >= len(amd_ops):
            # Remaining NVIDIA ops
            alignments.append((None, nv_i, "nvidia_only"))
            nv_i += 1
        elif nv_i >= len(nvidia_ops):
            # Remaining AMD ops
            alignments.append((amd_i, None, "amd_only"))
            amd_i += 1
        elif amd_ops[amd_i] == nvidia_ops[nv_i]:
            # Match
            alignments.append((amd_i, nv_i, "match"))
            amd_i += 1
            nv_i += 1
        else:
            # Mismatch - check if AMD has an extra operation
            amd_kernel_name = amd_kernels[amd_i]["name"].lower()

            # Known fusion differences
            if "reduce_segments" in amd_kernel_name:
                # AMD has reduce_segments, NVIDIA fuses it
                alignments.append((amd_i, None, "amd_only"))
                amd_i += 1
            else:
                # Unknown mismatch - try to match anyway
                alignments.append((amd_i, nv_i, "match"))
                amd_i += 1
                nv_i += 1

    return alignments


def match_traces(
    amd_trace_path: str,
    nvidia_trace_path: str,
) -> tuple[list[GraphPair], list[KernelMatch]]:
    """Match traces with fusion difference detection."""
    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={len(amd_graphs)}, NVIDIA={len(nvidia_graphs)}"
        )

    # Detect platform from first kernel
    with open(amd_trace_path, "rb") as f:
        amd_trace = orjson.loads(f.read())
    props = amd_trace.get("deviceProperties", [{}])[0]
    platform_amd = "AMD" if amd_trace.get("roctracer_version") or props.get("warpSize") == 64 else "NVIDIA"

    with open(nvidia_trace_path, "rb") as f:
        nvidia_trace = orjson.loads(f.read())
    props = nvidia_trace.get("deviceProperties", [{}])[0]
    platform_nvidia = "AMD" if nvidia_trace.get("roctracer_version") or props.get("warpSize") == 64 else "NVIDIA"

    graph_pairs = []
    kernel_matches = []

    for graph_idx, ((amd_ts, amd_corr), (nv_ts, nv_corr)) in enumerate(
        zip(amd_graphs, nvidia_graphs)
    ):
        amd_kernels = load_kernels_for_correlation(amd_trace_path, amd_corr, platform_amd)
        nvidia_kernels = load_kernels_for_correlation(
            nvidia_trace_path, nv_corr, platform_nvidia
        )

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

        # Align sequences
        alignments = align_sequences_with_fusion(
            amd_kernels, nvidia_kernels, platform_amd, platform_nvidia
        )

        # Create matches and track fusion differences
        fusion_diffs = []

        for amd_i, nv_i, align_type in alignments:
            if align_type == "match":
                # Perfect match
                amd_k = amd_kernels[amd_i]
                nv_k = nvidia_kernels[nv_i]

                amd_op, _ = classify(amd_k["name"], platform_amd)
                nv_op, _ = classify(nv_k["name"], platform_nvidia)

                confidence = 1.0 if amd_op == nv_op else 0.8

                kernel_matches.append(
                    KernelMatch(
                        graph_index=graph_idx,
                        amd_position=amd_i,
                        nvidia_position=nv_i,
                        amd_kernel=amd_k,
                        nvidia_kernel=nv_k,
                        operation_type=amd_op.value,
                        confidence=confidence,
                        amd_correlation=amd_corr,
                        nvidia_correlation=nv_corr,
                    )
                )
            elif align_type == "amd_only":
                # Fusion difference: AMD has this, NVIDIA fused it
                amd_k = amd_kernels[amd_i]
                amd_op, _ = classify(amd_k["name"], platform_amd)

                fusion_diff = FusionDifference(
                    platform_with_kernel="AMD",
                    kernel_name=amd_k["name"],
                    operation_type=amd_op.value,
                    position=amd_i,
                    likely_fused_into="adjacent operation (NVIDIA fuses)",
                )
                fusion_diffs.append(fusion_diff)

                kernel_matches.append(
                    KernelMatch(
                        graph_index=graph_idx,
                        amd_position=amd_i,
                        nvidia_position=None,
                        amd_kernel=amd_k,
                        nvidia_kernel=None,
                        operation_type=amd_op.value,
                        confidence=0.5,
                        is_fusion_difference=True,
                        fusion_info=fusion_diff,
                        amd_correlation=amd_corr,
                        nvidia_correlation=nv_corr,
                    )
                )
            elif align_type == "nvidia_only":
                # Fusion difference: NVIDIA has this, AMD fused it
                nv_k = nvidia_kernels[nv_i]
                nv_op, _ = classify(nv_k["name"], platform_nvidia)

                fusion_diff = FusionDifference(
                    platform_with_kernel="NVIDIA",
                    kernel_name=nv_k["name"],
                    operation_type=nv_op.value,
                    position=nv_i,
                    likely_fused_into="adjacent operation (AMD fuses)",
                )
                fusion_diffs.append(fusion_diff)

                kernel_matches.append(
                    KernelMatch(
                        graph_index=graph_idx,
                        amd_position=None,
                        nvidia_position=nv_i,
                        amd_kernel=None,
                        nvidia_kernel=nv_k,
                        operation_type=nv_op.value,
                        confidence=0.5,
                        is_fusion_difference=True,
                        fusion_info=fusion_diff,
                        amd_correlation=amd_corr,
                        nvidia_correlation=nv_corr,
                    )
                )

        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,
                fusion_differences=fusion_diffs,
            )
        )

    return graph_pairs, kernel_matches


def get_matching_statistics(kernel_matches: list[KernelMatch]) -> dict[str, Any]:
    """Calculate statistics including fusion analysis."""
    total = len(kernel_matches)
    perfect = sum(1 for m in kernel_matches if m.confidence == 1.0)
    fusion_diffs = sum(1 for m in kernel_matches if m.is_fusion_difference)

    # Breakdown by operation type
    from collections import defaultdict

    by_operation = defaultdict(lambda: {"total": 0, "perfect": 0, "fusion": 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
        if match.is_fusion_difference:
            by_operation[op]["fusion"] += 1

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