"""Utilities for optimizing graph matching results for extension use.
The extension needs to store graph matches in state for UI display, but the full
result (~200MB+) is too large. These utilities extract only what's needed.
"""
from typing import Any
def extract_minimal_matches_for_fusion(graph_result: dict[str, Any]) -> dict[str, Any]:
    """Extract minimal data needed for fusion analysis from graph matching result.
    This reduces the size from ~200MB to ~70KB by keeping only correlation ID pairs.
    The all_matches field is removed as it's not used by fusion analysis or the UI.
    Args:
        graph_result: Full result from match_traces()
    Returns:
        Optimized dictionary containing:
            - amd_platform, nv_platform: Platform names
            - summary: Match statistics
            - graph_pairs: Minimal correlation ID pairs (just IDs, no kernel data)
    """
    minimal_graph_pairs = [
        {
            'amd_corr_id': pair.get('amd_corr_id', pair.get('amd_correlation')),
            'nv_corr_id': pair.get('nv_corr_id', pair.get('nv_correlation')),
        }
        for pair in graph_result['graph_pairs']
    ]
    return {
        'amd_platform': graph_result['amd_platform'],
        'nv_platform': graph_result['nv_platform'],
        'summary': graph_result['summary'],
        'graph_pairs': minimal_graph_pairs,
    }
def extract_matches_for_ui(graph_result: dict[str, Any], max_pairs: int | None = 500, use_name_dict: bool = True) -> dict[str, Any]:
    """Extract data needed for kernel matching UI from graph matching result.
    This keeps essential kernel data and match information while still optimizing size.
    For large traces, limits to first N graph pairs to avoid slow postMessage transfer.
    Args:
        graph_result: Full result from match_traces()
        max_pairs: Maximum number of graph pairs to include (default 100, None for all)
        use_name_dict: If True, store kernel names in a dictionary and use indices (default True)
    Returns:
        Optimized dictionary containing:
            - amd_platform, nv_platform: Platform names
            - amd_graphs, nv_graphs: Empty arrays (not needed for UI)
            - summary: Match statistics (includes total count before truncation)
            - graph_pairs: First N graph pairs with essential kernel data and matches
            - kernel_name_dict: Optional dictionary of kernel names (if use_name_dict is True)
            - truncated: Boolean indicating if results were truncated
            - uses_name_dict: Boolean indicating if kernel names are in dictionary
    """
    name_to_idx: dict[str, int] = {}
    name_list: list[str] = []
    type_to_idx: dict[str, int] = {}
    type_list: list[str] = []
    status_to_idx: dict[str, int] = {}
    status_list: list[str] = []
    def get_name_idx(name: str) -> int:
        """Get or create index for a kernel name."""

        if len(name) > 80:
            # Try to keep the meaningful part before template parameters
            if '<' in name:
                base = name.split('<')[0]
                name = f"{base}<...>" if len(base) < 70 else f"{base[:70]}..."
            else:
                name = name[:77] + "..."

        if name not in name_to_idx:
            idx = len(name_list)
            name_list.append(name)
            name_to_idx[name] = idx
        return name_to_idx[name]
    def get_type_idx(type_str: str) -> int:
        """Get or create index for an operation type."""
        if type_str not in type_to_idx:
            idx = len(type_list)
            type_list.append(type_str)
            type_to_idx[type_str] = idx
        return type_to_idx[type_str]
    def get_status_idx(status: str) -> int:
        """Get or create index for a match status."""
        if status not in status_to_idx:
            idx = len(status_list)
            status_list.append(status)
            status_to_idx[status] = idx
        return status_to_idx[status]
    # Keep essential fields from kernel events
    def minimize_kernel(k: dict[str, Any]) -> dict[str, Any]:
        """Keep only essential kernel fields for UI display.
        UI displays: name, dur, phase, ts (for timeline view)
        """
        if use_name_dict:

            name_idx = get_name_idx(k.get('name', ''))
            return {
                'n': name_idx,  # Shorter key saves space
                'd': k.get('dur', 0),
                'p': k.get('phase', 'decode'),  # Phase from vLLM annotations
                't': k.get('ts', 0),  # Timestamp for timeline view
            }
        else:
            # Legacy: store full name
            name = k.get('name', '')
            if len(name) > 80:
                if '<' in name:
                    base = name.split('<')[0]
                    name = f"{base}<...>" if len(base) < 70 else f"{base[:70]}..."
                else:
                    name = name[:77] + "..."
            return {
                'name': name,
                'dur': k.get('dur', 0),
                'phase': k.get('phase', 'decode'),
                'ts': k.get('ts', 0),  # Timestamp for timeline view
            }
    def build_kernel_index_map(kernels: list[dict[str, Any]]) -> dict[str, int]:
        """Build a map from (name, ts) to kernel index.
        Note: We use name:ts as the key because kernels at the same position
        in the sequence may have the same name but different timestamps.
        """
        return {f"{k.get('name', '')}:{k.get('ts', 0)}": i for i, k in enumerate(kernels)}
    # Minimize matches - use indices for everything
    def minimize_match(match: dict[str, Any], amd_idx_map: dict[str, int], nv_idx_map: dict[str, int]) -> dict[str, Any]:
        """Keep match data with indices instead of strings.
        Uses dictionaries for:
        - Kernel references (already implemented)
        - Operation types (amd_type, nv_type)
        - Match status (MATCH, AMD_ONLY, NV_ONLY, MISMATCH)
        This saves ~11 MB across all matches by deduplicating repeated type/status strings.
        """
        if use_name_dict:
            result: dict[str, Any] = {
                'amd_type_idx': get_type_idx(match.get('amd_type', '-')),
                'nv_type_idx': get_type_idx(match.get('nv_type', '-')),
                'status_idx': get_status_idx(match.get('status', 'MISMATCH')),
            }
        else:
            # Legacy: use full strings
            result: dict[str, Any] = {
                'amd_type': match.get('amd_type', '-'),
                'nv_type': match.get('nv_type', '-'),
                'status': match.get('status', 'MISMATCH'),
            }
        # Reference kernels by index instead of embedding
        amd_kernel = match.get('amd_kernel')
        nv_kernel = match.get('nv_kernel')
        if amd_kernel:
            key = f"{amd_kernel.get('name', '')}:{amd_kernel.get('ts', 0)}"
            idx = amd_idx_map.get(key, -1)
            if idx == -1:
                # Debug: log why lookup failed
                import sys
                print(f"[WARN] AMD kernel lookup failed for key: {key[:80]}", file=sys.stderr)
                print(f"  Available keys sample: {list(amd_idx_map.keys())[:3]}", file=sys.stderr)
            result['amd_kernel_idx'] = idx
        else:
            result['amd_kernel_idx'] = None
        if nv_kernel:
            key = f"{nv_kernel.get('name', '')}:{nv_kernel.get('ts', 0)}"
            idx = nv_idx_map.get(key, -1)
            if idx == -1:
                # Debug: log why lookup failed
                import sys
                print(f"[WARN] NV kernel lookup failed for key: {key[:80]}", file=sys.stderr)
                print(f"  Available keys sample: {list(nv_idx_map.keys())[:3]}", file=sys.stderr)
            result['nv_kernel_idx'] = idx
        else:
            result['nv_kernel_idx'] = None
        return result
    all_pairs = graph_result.get('graph_pairs', [])
    pairs_to_process = all_pairs[:max_pairs] if max_pairs else all_pairs
    truncated = max_pairs is not None and len(all_pairs) > max_pairs
    optimized_graph_pairs = []
    for pair in pairs_to_process:
        amd_kernels = pair.get('amd_kernels', [])
        nv_kernels = pair.get('nv_kernels', [])
        amd_idx_map = build_kernel_index_map(amd_kernels)
        nv_idx_map = build_kernel_index_map(nv_kernels)
        # Minimize matches first (needs original kernels for lookup)
        minimized_matches = [minimize_match(m, amd_idx_map, nv_idx_map) for m in pair.get('matches', [])]
        # Then minimize kernels (keeps ts for timeline view)
        minimized_amd_kernels = [minimize_kernel(k) for k in amd_kernels]
        minimized_nv_kernels = [minimize_kernel(k) for k in nv_kernels]
        optimized_graph_pairs.append({
            'amd_correlation': pair.get('amd_correlation', 0),
            'nv_correlation': pair.get('nv_correlation', 0),
            'amd_corr_id': pair.get('amd_corr_id', pair.get('amd_correlation', 0)),
            'nv_corr_id': pair.get('nv_corr_id', pair.get('nv_correlation', 0)),
            'amd_kernels': minimized_amd_kernels,
            'nv_kernels': minimized_nv_kernels,
            'matches': minimized_matches,
            'graph_type': pair.get('graph_type', 'small'),
        })

    summary = dict(graph_result['summary'])
    if truncated:
        summary['total_graph_pairs'] = len(all_pairs)
        summary['displayed_graph_pairs'] = len(optimized_graph_pairs)
    result = {
        'amd_platform': graph_result['amd_platform'],
        'nv_platform': graph_result['nv_platform'],
        'amd_graphs': [],  # Not needed for UI
        'nv_graphs': [],   # Not needed for UI
        'summary': summary,
        'graph_pairs': optimized_graph_pairs,
        'truncated': truncated,
        'uses_kernel_indices': True,  # Indicate that matches use indices
        'uses_name_dict': use_name_dict,  # Indicate if names are in dictionary
    }
    if use_name_dict:
        result['kernel_name_dict'] = name_list
        result['op_type_dict'] = type_list
        result['status_dict'] = status_list
    return result
def can_use_for_fusion(graph_matches: dict[str, Any]) -> bool:
    """Check if graph_matches contains enough data for fusion analysis.
    Args:
        graph_matches: Graph matching result (full or optimized)
    Returns:
        True if fusion analysis can use this data, False if need to re-run
    """
    # Need graph_pairs for fusion detection
    if 'graph_pairs' not in graph_matches:
        return False
    if not graph_matches['graph_pairs']:
        return True  # Empty is fine
    first_pair = graph_matches['graph_pairs'][0]
    required_keys = {'amd_corr_id', 'nv_corr_id'}
    return required_keys.issubset(first_pair.keys())
