"""Combined analysis that loads traces once and performs both analyze_traces and match_traces."""
import time
from pathlib import Path
from typing import Any
from .analyzer import analyze_traces
from .graph_matcher import match_traces
from .graph_matcher_utils import extract_matches_for_ui
from .formatter import format_json
def analyze_and_match_traces(
    trace1_path: str | Path,
    trace2_path: str | Path,
    phase_filter: str = "all",
    max_stacks: int = 3,
    min_graph_size: int = 300,
    max_graph_pairs_ui: int = 100,
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    """Combined analysis that loads traces once for both operations.
    This function performs both analyze_traces() and match_traces() in a single call,
    loading the trace files only once to improve performance.
    Args:
        trace1_path: Path to first trace file
        trace2_path: Path to second trace file
        phase_filter: Phase filter for analyze_traces ('all', 'prefill', 'decode')
        max_stacks: Maximum Python stack traces to include
        min_graph_size: Minimum size for large CUDA graph detection
        max_graph_pairs_ui: Maximum graph pairs to include in UI-optimized result
    Returns:
        Tuple of (analyze_result, match_result_full, match_result_optimized)
        - analyze_result: Full result from analyze_traces()
        - match_result_full: Full result from match_traces() (for caching)
        - match_result_optimized: UI-optimized result from extract_matches_for_ui()
    """
    trace1_path = Path(trace1_path)
    trace2_path = Path(trace2_path)
    # Run analyze_traces (this loads traces)
    t1 = time.perf_counter()
    analyze_result = analyze_traces(trace1_path, trace2_path, phase_filter=phase_filter, max_stacks=max_stacks)
    t2 = time.perf_counter()
    print(f"✓ analyze_traces: {t2-t1:.2f}s", flush=True)
    # Run match_traces (this ALSO loads traces - unavoidable without refactoring)
    match_result_full = None
    match_result_optimized = None
    try:
        t1 = time.perf_counter()
        match_result_full = match_traces(trace1_path, trace2_path, min_graph_size=min_graph_size)
        t2 = time.perf_counter()
        print(f"✓ match_traces:   {t2-t1:.2f}s", flush=True)

        t1 = time.perf_counter()
        match_result_optimized = extract_matches_for_ui(match_result_full, max_pairs=max_graph_pairs_ui)
        t2 = time.perf_counter()
        print(f"✓ extract_for_ui: {t2-t1:.2f}s", flush=True)
    except Exception as e:
        print(f"Warning: Graph matching failed: {e}")
        import traceback
        traceback.print_exc()
    return analyze_result, match_result_full, match_result_optimized
def analyze_and_match_traces_json(
    trace1_path: str | Path,
    trace2_path: str | Path,
    phase_filter: str = "all",
    max_stacks: int = 3,
    min_graph_size: int = 300,
    max_graph_pairs_ui: int = 100,
) -> tuple[str, dict[str, Any]]:
    """Combined analysis returning JSON output and full match results.
    This is the main entry point for the extension - it performs both analyses
    and returns formatted JSON for the UI plus full results for caching.
    Args:
        trace1_path: Path to first trace file
        trace2_path: Path to second trace file
        phase_filter: Phase filter for analyze_traces ('all', 'prefill', 'decode')
        max_stacks: Maximum Python stack traces to include
        min_graph_size: Minimum size for large CUDA graph detection
        max_graph_pairs_ui: Maximum graph pairs to include in UI-optimized result
    Returns:
        Tuple of (json_output, match_result_full)
        - json_output: JSON string for UI (analyze_result + optimized match_result)
        - match_result_full: Full match results for caching (can be pickled)
    """
    analyze_result, match_result_full, match_result_optimized = analyze_and_match_traces(
        trace1_path,
        trace2_path,
        phase_filter=phase_filter,
        max_stacks=max_stacks,
        min_graph_size=min_graph_size,
        max_graph_pairs_ui=max_graph_pairs_ui,
    )
    # Format as JSON with optimized graph results
    t1 = time.perf_counter()
    json_output = format_json(analyze_result, graph_results=match_result_optimized)
    t2 = time.perf_counter()
    print(f"✓ format_json:    {t2-t1:.2f}s", flush=True)
    return json_output, match_result_full
