"""Distributed traces analysis tool.
Pure function executor for analyzing NCCL/RCCL collective operations
in distributed training traces.
"""
from pathlib import Path
from wafer_core.lib.distributed_traces import (
    AnalyzerConfig,
    DistributedTracesAnalyzer,
)
from wafer_core.lib.distributed_traces.analysis.bandwidth import InterconnectType
from wafer_core.rollouts.dtypes import (
    Tool,
    ToolCall,
    ToolFunction,
    ToolFunctionParameter,
    ToolResult,
)
# Tool schema definition
DISTRIBUTED_TRACES_ANALYZE_TOOL = Tool(
    type="function",
    function=ToolFunction(
        name="distributed_traces_analyze",
        description=(
            "Analyze distributed training traces to identify communication bottlenecks. "
            "Supports PyTorch Profiler, NCCL Inspector, rocprofv3, Nsight Systems, and ROCm Systems Profiler formats. "
            "Provides bandwidth efficiency, compute/communication overlap, straggler detection, and hang risk analysis."
        ),
        parameters=ToolFunctionParameter(
            type="object",
            properties={
                "trace_paths": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Paths to trace files (one per rank, or a directory containing trace files)"
                },
                "interconnect": {
                    "type": "string",
                    "enum": ["auto", "nvlink3", "nvlink4", "nvlink5", "pcie4", "pcie5", "xgmi", "ib-hdr", "ib-ndr"],
                    "description": "Interconnect type for bandwidth calculation (default: auto-detect)"
                },
                "timeout_sec": {
                    "type": "number",
                    "description": "NCCL timeout threshold in seconds for hang risk analysis (default: 600)"
                },
                "output_format": {
                    "type": "string",
                    "enum": ["summary", "json", "csv"],
                    "description": "Output format (default: summary)"
                },
                "output_path": {
                    "type": "string",
                    "description": "Path to write output file (optional)"
                },
            }
        ),
        required=["trace_paths"]
    )
)
# Interconnect string to enum mapping
_INTERCONNECT_MAP = {
    "auto": InterconnectType.UNKNOWN,
    "nvlink3": InterconnectType.NVLINK_3,
    "nvlink4": InterconnectType.NVLINK_4,
    "nvlink5": InterconnectType.NVLINK_5,
    "pcie4": InterconnectType.PCIE_4,
    "pcie5": InterconnectType.PCIE_5,
    "xgmi": InterconnectType.XGMI,
    "ib-hdr": InterconnectType.IB_HDR,
    "ib-ndr": InterconnectType.IB_NDR,
}
async def exec_distributed_traces_analyze(
    tool_call: ToolCall,
    working_dir: Path,
) -> ToolResult:
    """Execute distributed traces analysis.
    Logic:
    1. Validate trace_paths argument exists
    2. Resolve relative paths against working_dir
    3. Expand directories to trace files
    4. Parse and analyze traces
    5. Format and return results
    Args:
        tool_call: Tool call with arguments
        working_dir: Working directory for relative paths
    Returns:
        ToolResult with analysis summary or error
    """
    # Validate required args
    if "trace_paths" not in tool_call.args:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="Missing required argument: 'trace_paths'"
        )
    trace_paths_arg = tool_call.args["trace_paths"]
    if not isinstance(trace_paths_arg, list) or len(trace_paths_arg) == 0:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="'trace_paths' must be a non-empty list of file paths"
        )
    # Resolve and expand paths
    resolved_paths = []
    trace_extensions = {".json", ".jsonl", ".csv", ".nsys-rep", ".sqlite", ".rocpd", ".db", ".gz"}
    for path_str in trace_paths_arg:
        path = Path(path_str)
        if not path.is_absolute():
            path = working_dir / path
        if path.is_dir():
            # Expand directory to trace files
            for ext in trace_extensions:
                resolved_paths.extend(str(f) for f in path.glob(f"*{ext}"))
        elif path.exists():
            resolved_paths.append(str(path))
        else:
            # Try glob pattern
            import glob
            matches = glob.glob(str(path))
            resolved_paths.extend(matches)
    if not resolved_paths:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="No trace files found in the specified paths"
        )

    interconnect_str = tool_call.args.get("interconnect", "auto")
    timeout_sec = tool_call.args.get("timeout_sec", 600.0)
    output_format = tool_call.args.get("output_format", "summary")
    output_path = tool_call.args.get("output_path")

    interconnect = _INTERCONNECT_MAP.get(interconnect_str, InterconnectType.UNKNOWN)
    config = AnalyzerConfig(
        interconnect=interconnect,
        timeout_sec=float(timeout_sec),
    )
    # Run analysis
    analyzer = DistributedTracesAnalyzer(config)
    session, load_error = analyzer.load_traces(resolved_paths)
    if load_error:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Failed to load traces: {load_error}"
        )
    if session is None:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="No traces could be parsed"
        )
    result = analyzer.analyze(session)
    # Format output
    if output_format == "json":
        output_str = analyzer.export_json(result, session, output_path)
        if output_path:
            content = f"✓ Analysis complete. JSON written to {output_path}"
        else:
            content = output_str
    elif output_format == "csv":
        output_str = analyzer.export_csv(session, output_path)
        if output_path:
            content = f"✓ Analysis complete. CSV written to {output_path}"
        else:
            content = output_str
    else:
        # Summary format
        content = _format_summary(result, session)
        if output_path:
            output_str = analyzer.export_json(result, session, output_path)
            content += f"\n\nDetailed JSON written to {output_path}"
    return ToolResult(
        tool_call_id=tool_call.id,
        is_error=False,
        content=content
    )
def _format_summary(result, session) -> str:
    """Format analysis result as readable summary."""
    summary = result.summary
    lines = [
        "═══ Distributed Traces Analysis ═══",
        "",
        f"Session: {summary.get('session_name', 'unnamed')}",
        f"Format:  {summary.get('trace_format', 'unknown')}",
        f"Ranks:   {summary.get('num_ranks', 0)} / {summary.get('world_size', 0)} world size",
        f"Duration: {summary.get('duration_ms', 0):.2f} ms",
        f"Collectives: {summary.get('total_collectives', 0)} ({summary.get('num_matched_collectives', 0)} matched)",
        "",
        "─── Performance Metrics ───",
        "",
        f"Bandwidth Efficiency: {summary.get('bandwidth_efficiency_percent', 0):.1f}%",
        f"  Achieved: {summary.get('achieved_bandwidth_gbps', 0):.2f} GB/s",
        "",
        f"Overlap Ratio: {summary.get('overlap_percent', 0):.1f}%",
        f"  (Communication hidden by compute)",
        "",
        f"Straggler Impact: {summary.get('total_straggler_impact_ms', 0):.2f} ms",
    ]
    flagged = summary.get('flagged_straggler_ranks', [])
    if flagged:
        lines.append(f"  Flagged ranks: {flagged}")
    hang_dist = summary.get('hang_risk_distribution', {})
    critical = hang_dist.get('critical', 0)
    warning = hang_dist.get('warning', 0)
    if critical > 0:
        lines.append(f"\n⚠️  CRITICAL hang risk: {critical} collectives at risk")
    elif warning > 0:
        lines.append(f"\n⚡ Warning: {warning} collectives with elevated duration")
    else:
        lines.append(f"\n✓ No hang risk detected")
    # Collectives by type
    by_type = summary.get('collectives_by_type', {})
    if by_type:
        lines.append("")
        lines.append("─── Collectives by Type ───")
        for coll_type, count in sorted(by_type.items(), key=lambda x: -x[1]):
            lines.append(f"  {coll_type}: {count}")
    return "\n".join(lines)
