"""Distributed traces comparison tool.
Pure function executor for comparing two distributed training trace sessions
to identify performance regressions or improvements.
"""
from pathlib import Path
from wafer_core.lib.distributed_traces import (
    AnalyzerConfig,
    DistributedTracesAnalyzer,
)
from wafer_core.rollouts.dtypes import (
    Tool,
    ToolCall,
    ToolFunction,
    ToolFunctionParameter,
    ToolResult,
)
# Tool schema definition
DISTRIBUTED_TRACES_COMPARE_TOOL = Tool(
    type="function",
    function=ToolFunction(
        name="distributed_traces_compare",
        description=(
            "Compare two distributed training trace sessions to identify performance changes. "
            "Reports bandwidth, overlap, and straggler impact deltas with regression/improvement classification."
        ),
        parameters=ToolFunctionParameter(
            type="object",
            properties={
                "baseline_paths": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Paths to baseline trace files (one per rank)"
                },
                "comparison_paths": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Paths to comparison trace files (one per rank)"
                },
                "output_path": {
                    "type": "string",
                    "description": "Path to write JSON comparison report (optional)"
                },
            }
        ),
        required=["baseline_paths", "comparison_paths"]
    )
)
async def exec_distributed_traces_compare(
    tool_call: ToolCall,
    working_dir: Path,
) -> ToolResult:
    """Execute distributed traces comparison.
    Logic:
    1. Validate baseline and comparison paths exist
    2. Load both trace sessions
    3. Compare sessions
    4. Format and return comparison results
    Args:
        tool_call: Tool call with arguments
        working_dir: Working directory for relative paths
    Returns:
        ToolResult with comparison summary or error
    """
    # Validate required args
    if "baseline_paths" not in tool_call.args:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="Missing required argument: 'baseline_paths'"
        )
    if "comparison_paths" not in tool_call.args:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="Missing required argument: 'comparison_paths'"
        )
    baseline_paths = tool_call.args["baseline_paths"]
    comparison_paths = tool_call.args["comparison_paths"]
    if not isinstance(baseline_paths, list) or len(baseline_paths) == 0:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="'baseline_paths' must be a non-empty list"
        )
    if not isinstance(comparison_paths, list) or len(comparison_paths) == 0:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="'comparison_paths' must be a non-empty list"
        )
    # Resolve paths
    def resolve_paths(paths: list[str]) -> list[str]:
        resolved = []
        for path_str in paths:
            path = Path(path_str)
            if not path.is_absolute():
                path = working_dir / path
            resolved.append(str(path))
        return resolved
    baseline_resolved = resolve_paths(baseline_paths)
    comparison_resolved = resolve_paths(comparison_paths)
    output_path = tool_call.args.get("output_path")

    analyzer = DistributedTracesAnalyzer()
    baseline_session, baseline_error = analyzer.load_traces(
        baseline_resolved, session_name="baseline"
    )
    if baseline_error:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Failed to load baseline traces: {baseline_error}"
        )
    comparison_session, comparison_error = analyzer.load_traces(
        comparison_resolved, session_name="comparison"
    )
    if comparison_error:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error=f"Failed to load comparison traces: {comparison_error}"
        )
    if baseline_session is None or comparison_session is None:
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=True,
            content="",
            error="Failed to parse one or both trace sessions"
        )
    # Compare
    comparison_result = analyzer.compare_sessions(baseline_session, comparison_session)
    # Format output
    content = _format_comparison(comparison_result)
    # Write JSON if output path specified
    if output_path:
        import json
        output_path_obj = Path(output_path)
        if not output_path_obj.is_absolute():
            output_path_obj = working_dir / output_path_obj
        comparison_dict = {
            "baseline": comparison_result.baseline_summary,
            "comparison": comparison_result.comparison_summary,
            "bandwidth_delta_percent": comparison_result.bandwidth_delta,
            "overlap_delta_percent": comparison_result.overlap_delta,
            "straggler_delta_percent": comparison_result.straggler_delta,
            "regressions": comparison_result.regressions,
            "improvements": comparison_result.improvements,
            "recommendations": comparison_result.recommendations,
        }
        output_path_obj.write_text(json.dumps(comparison_dict, indent=2))
        content += f"\n\nDetailed report written to {output_path_obj}"
    return ToolResult(
        tool_call_id=tool_call.id,
        is_error=False,
        content=content
    )
def _format_comparison(result) -> str:
    """Format comparison result as readable summary."""
    lines = [
        "═══ Distributed Traces Comparison ═══",
        "",
        "─── Deltas (comparison vs baseline) ───",
        "",
    ]
    # Bandwidth
    bw_delta = result.bandwidth_delta
    bw_indicator = "↑" if bw_delta > 0 else "↓" if bw_delta < 0 else "─"
    bw_status = "improved" if bw_delta > 5 else "regressed" if bw_delta < -5 else "unchanged"
    lines.append(f"Bandwidth: {bw_indicator} {abs(bw_delta):.1f}% ({bw_status})")
    # Overlap
    overlap_delta = result.overlap_delta
    overlap_indicator = "↑" if overlap_delta > 0 else "↓" if overlap_delta < 0 else "─"
    overlap_status = "improved" if overlap_delta > 5 else "regressed" if overlap_delta < -5 else "unchanged"
    lines.append(f"Overlap:   {overlap_indicator} {abs(overlap_delta):.1f}% ({overlap_status})")
    # Straggler (lower is better)
    straggler_delta = result.straggler_delta
    straggler_indicator = "↓" if straggler_delta < 0 else "↑" if straggler_delta > 0 else "─"
    straggler_status = "improved" if straggler_delta < -10 else "regressed" if straggler_delta > 10 else "unchanged"
    lines.append(f"Straggler: {straggler_indicator} {abs(straggler_delta):.1f}% ({straggler_status})")
    # Regressions
    if result.regressions:
        lines.append("")
        lines.append("⚠️  Regressions:")
        for reg in result.regressions:
            lines.append(f"  - {reg}")
    # Improvements
    if result.improvements:
        lines.append("")
        lines.append("✓ Improvements:")
        for imp in result.improvements:
            lines.append(f"  - {imp}")
    # Recommendations
    if result.recommendations:
        lines.append("")
        lines.append("─── Recommendations ───")
        for rec in result.recommendations:
            lines.append(f"  {rec}")
    return "\n".join(lines)
