"""Tests for distributed traces Wevin tools.

Tests the tool definitions and executors in tools/distributed_traces_tools/.
"""

import asyncio
from pathlib import Path
import tempfile
import json

import pytest

from wafer_core.tools.distributed_traces_tools import (
    DISTRIBUTED_TRACES_ANALYZE_TOOL,
    DISTRIBUTED_TRACES_COMPARE_TOOL,
    exec_distributed_traces_analyze,
    exec_distributed_traces_compare,
)
from wafer_core.rollouts.dtypes import ToolCall


class TestToolDefinitions:
    """Test tool schema definitions."""

    def test_analyze_tool_name(self):
        """Test analyze tool has correct name."""
        assert DISTRIBUTED_TRACES_ANALYZE_TOOL.function.name == "distributed_traces_analyze"

    def test_analyze_tool_required_params(self):
        """Test analyze tool requires trace_paths."""
        assert "trace_paths" in DISTRIBUTED_TRACES_ANALYZE_TOOL.function.required

    def test_analyze_tool_has_interconnect_param(self):
        """Test analyze tool has interconnect parameter."""
        props = DISTRIBUTED_TRACES_ANALYZE_TOOL.function.parameters.properties
        assert "interconnect" in props
        assert "enum" in props["interconnect"]

    def test_analyze_tool_has_output_format_param(self):
        """Test analyze tool has output_format parameter."""
        props = DISTRIBUTED_TRACES_ANALYZE_TOOL.function.parameters.properties
        assert "output_format" in props
        assert props["output_format"]["enum"] == ["summary", "json", "csv"]

    def test_compare_tool_name(self):
        """Test compare tool has correct name."""
        assert DISTRIBUTED_TRACES_COMPARE_TOOL.function.name == "distributed_traces_compare"

    def test_compare_tool_required_params(self):
        """Test compare tool requires both baseline and comparison paths."""
        required = DISTRIBUTED_TRACES_COMPARE_TOOL.function.required
        assert "baseline_paths" in required
        assert "comparison_paths" in required


class TestAnalyzeToolExecutor:
    """Test the analyze tool executor."""

    @pytest.fixture
    def temp_traces(self, tmp_path):
        """Create temporary trace files."""
        # Create trace for rank 0
        trace0 = {
            "traceEvents": [
                {
                    "ph": "X",
                    "cat": "cpu_op",
                    "name": "nccl:all_reduce",
                    "ts": 1000,
                    "dur": 500,
                    "args": {"Input Coverage": "8192"}
                },
            ]
        }
        trace0_path = tmp_path / "rank0.json"
        trace0_path.write_text(json.dumps(trace0))

        # Create trace for rank 1
        trace1 = {
            "traceEvents": [
                {
                    "ph": "X",
                    "cat": "cpu_op",
                    "name": "nccl:all_reduce",
                    "ts": 1000,
                    "dur": 600,
                    "args": {"Input Dims": "8192"}
                },
            ]
        }
        trace1_path = tmp_path / "rank1.json"
        trace1_path.write_text(json.dumps(trace1))

        return [str(trace0_path), str(trace1_path)]

    def test_missing_trace_paths(self):
        """Test error when trace_paths is missing."""
        tool_call = ToolCall(
            id="test-1",
            name="distributed_traces_analyze",
            args={}
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, Path("/")))

        assert result.is_error is True
        assert "trace_paths" in result.error

    def test_empty_trace_paths(self):
        """Test error when trace_paths is empty."""
        tool_call = ToolCall(
            id="test-2",
            name="distributed_traces_analyze",
            args={"trace_paths": []}
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, Path("/")))

        assert result.is_error is True
        assert "non-empty" in result.error

    def test_file_not_found(self):
        """Test error when trace file doesn't exist."""
        tool_call = ToolCall(
            id="test-3",
            name="distributed_traces_analyze",
            args={"trace_paths": ["/nonexistent/path/trace.json"]}
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, Path("/")))

        assert result.is_error is True
        assert "No trace files found" in result.error

    def test_successful_analysis(self, temp_traces, tmp_path):
        """Test successful trace analysis."""
        tool_call = ToolCall(
            id="test-4",
            name="distributed_traces_analyze",
            args={"trace_paths": temp_traces}
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, tmp_path))

        assert result.is_error is False
        assert "Distributed Traces Analysis" in result.content
        assert "Ranks:" in result.content

    def test_json_output_format(self, temp_traces, tmp_path):
        """Test JSON output format."""
        tool_call = ToolCall(
            id="test-5",
            name="distributed_traces_analyze",
            args={
                "trace_paths": temp_traces,
                "output_format": "json",
            }
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, tmp_path))

        assert result.is_error is False
        # Should be valid JSON
        data = json.loads(result.content)
        assert "session" in data
        assert "analysis" in data

    def test_csv_output_format(self, temp_traces, tmp_path):
        """Test CSV output format."""
        tool_call = ToolCall(
            id="test-6",
            name="distributed_traces_analyze",
            args={
                "trace_paths": temp_traces,
                "output_format": "csv",
            }
        )

        result = asyncio.run(exec_distributed_traces_analyze(tool_call, tmp_path))

        assert result.is_error is False
        # Should have CSV header
        assert "rank" in result.content.lower()


class TestCompareToolExecutor:
    """Test the compare tool executor."""

    @pytest.fixture
    def temp_traces(self, tmp_path):
        """Create temporary trace files for comparison."""
        # Baseline traces
        baseline0 = {
            "traceEvents": [
                {"ph": "X", "cat": "cpu_op", "name": "nccl:all_reduce", "ts": 1000, "dur": 500}
            ]
        }
        baseline1 = {
            "traceEvents": [
                {"ph": "X", "cat": "cpu_op", "name": "nccl:all_reduce", "ts": 1000, "dur": 500}
            ]
        }

        # Comparison traces (slower)
        comparison0 = {
            "traceEvents": [
                {"ph": "X", "cat": "cpu_op", "name": "nccl:all_reduce", "ts": 1000, "dur": 800}
            ]
        }
        comparison1 = {
            "traceEvents": [
                {"ph": "X", "cat": "cpu_op", "name": "nccl:all_reduce", "ts": 1000, "dur": 400}
            ]
        }

        baseline0_path = tmp_path / "baseline_rank0.json"
        baseline1_path = tmp_path / "baseline_rank1.json"
        comparison0_path = tmp_path / "comparison_rank0.json"
        comparison1_path = tmp_path / "comparison_rank1.json"

        baseline0_path.write_text(json.dumps(baseline0))
        baseline1_path.write_text(json.dumps(baseline1))
        comparison0_path.write_text(json.dumps(comparison0))
        comparison1_path.write_text(json.dumps(comparison1))

        return {
            "baseline": [str(baseline0_path), str(baseline1_path)],
            "comparison": [str(comparison0_path), str(comparison1_path)],
        }

    def test_missing_baseline_paths(self):
        """Test error when baseline_paths is missing."""
        tool_call = ToolCall(
            id="test-1",
            name="distributed_traces_compare",
            args={"comparison_paths": ["/some/path"]}
        )

        result = asyncio.run(exec_distributed_traces_compare(tool_call, Path("/")))

        assert result.is_error is True
        assert "baseline_paths" in result.error

    def test_missing_comparison_paths(self):
        """Test error when comparison_paths is missing."""
        tool_call = ToolCall(
            id="test-2",
            name="distributed_traces_compare",
            args={"baseline_paths": ["/some/path"]}
        )

        result = asyncio.run(exec_distributed_traces_compare(tool_call, Path("/")))

        assert result.is_error is True
        assert "comparison_paths" in result.error

    def test_successful_comparison(self, temp_traces, tmp_path):
        """Test successful trace comparison."""
        tool_call = ToolCall(
            id="test-3",
            name="distributed_traces_compare",
            args={
                "baseline_paths": temp_traces["baseline"],
                "comparison_paths": temp_traces["comparison"],
            }
        )

        result = asyncio.run(exec_distributed_traces_compare(tool_call, tmp_path))

        assert result.is_error is False
        assert "Comparison" in result.content
        assert "Bandwidth" in result.content
        assert "Overlap" in result.content
        assert "Straggler" in result.content
