"""Timeline export for visualization.
Exports trace data to formats suitable for visualization:
- Custom JSON format for wevin-extension timeline view
- Perfetto format for external viewing
"""
import json
from pathlib import Path
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import CollectiveType
from wafer_core.lib.distributed_traces.models.trace_session import (
    CollectiveMatch,
    TraceSession,
)
def export_timeline_json(
    session: TraceSession,
    matches: Optional[list[CollectiveMatch]] = None,
    output_path: Optional[str | Path] = None,
) -> str:
    """Export session to JSON format optimized for timeline visualization.
    The output format is designed for the wevin-extension multi-rank
    timeline view, with data structured for efficient rendering.
    Args:
        session: TraceSession to export
        matches: Optional collective matches for alignment lines
        output_path: Optional path to write to
    Returns:
        JSON string
    """

    min_time = session.start_time_ns
    max_time = session.end_time_ns
    duration = max_time - min_time
    data = {
        "metadata": {
            "name": session.name,
            "world_size": session.world_size,
            "num_ranks": session.num_ranks,
            "start_time_ns": min_time,
            "end_time_ns": max_time,
            "duration_ns": duration,
            "duration_ms": duration / 1_000_000,
        },
        "ranks": [],
        "alignment_lines": [],
    }
    # Export per-rank data
    for rank in sorted(session.timelines.keys()):
        timeline = session.timelines[rank]
        rank_data = {
            "rank": rank,
            "device_id": timeline.device_id,
            "hostname": timeline.hostname,
            "events": [],
        }

        for c in timeline.collectives:
            rank_data["events"].append({
                "type": "collective",
                "name": c.collective_type.value,
                "start_ns": c.start_time_ns - min_time,  # Relative time
                "end_ns": c.end_time_ns - min_time,
                "duration_ns": c.duration_ns,
                "message_size_bytes": c.message_size_bytes,
                "collective_type": c.collective_type.value,
                "sequence": c.sequence_number,
                "algorithm": c.algorithm,
            })

        for e in timeline.compute_events[:1000]:
            rank_data["events"].append({
                "type": "compute",
                "name": e.name,
                "start_ns": e.start_time_ns - min_time,
                "end_ns": e.end_time_ns - min_time,
                "duration_ns": e.duration_ns,
                "stream_id": e.stream_id,
            })

        rank_data["events"].sort(key=lambda e: e["start_ns"])
        data["ranks"].append(rank_data)

    if matches:
        for match in matches:
            if match.num_ranks < 2:
                continue

            points = []
            for rank, coll in sorted(match.collectives.items()):
                points.append({
                    "rank": rank,
                    "start_ns": coll.start_time_ns - min_time,
                    "end_ns": coll.end_time_ns - min_time,
                    "duration_ns": coll.duration_ns,
                })
            data["alignment_lines"].append({
                "sequence": match.sequence_number,
                "collective_type": match.collective_type.value,
                "straggler_rank": match.straggler_rank,
                "points": points,
            })
    json_str = json.dumps(data, indent=2)
    if output_path:
        Path(output_path).write_text(json_str)
    return json_str
def export_perfetto_format(
    session: TraceSession,
    output_path: Optional[str | Path] = None,
) -> str:
    """Export session to Perfetto-compatible Chrome trace format.
    This allows viewing the trace in Perfetto UI or Chrome's trace viewer.
    Each rank becomes a separate process in the trace.
    Args:
        session: TraceSession to export
        output_path: Optional path to write to
    Returns:
        JSON string in Chrome trace format
    """
    trace_events = []
    # Process metadata events
    for rank, timeline in session.timelines.items():
        # Process name
        trace_events.append({
            "ph": "M",
            "cat": "__metadata",
            "name": "process_name",
            "pid": rank,
            "tid": 0,
            "args": {"name": f"Rank {rank}"},
        })
        # Thread names
        trace_events.append({
            "ph": "M",
            "cat": "__metadata",
            "name": "thread_name",
            "pid": rank,
            "tid": 0,
            "args": {"name": "Collectives"},
        })
        trace_events.append({
            "ph": "M",
            "cat": "__metadata",
            "name": "thread_name",
            "pid": rank,
            "tid": 1,
            "args": {"name": "Compute"},
        })

    for rank, timeline in session.timelines.items():
        for c in timeline.collectives:
            trace_events.append({
                "ph": "X",
                "cat": "collective",
                "name": c.collective_type.value,
                "pid": rank,
                "tid": 0,
                "ts": c.start_time_ns / 1000,  # Convert to microseconds
                "dur": c.duration_ns / 1000,
                "args": {
                    "message_size_bytes": c.message_size_bytes,
                    "sequence": c.sequence_number,
                    "algorithm": c.algorithm,
                    "process_group": c.process_group.comm_id if c.process_group else None,
                },
            })

    for rank, timeline in session.timelines.items():
        for e in timeline.compute_events[:5000]:
            trace_events.append({
                "ph": "X",
                "cat": "compute",
                "name": e.name,
                "pid": rank,
                "tid": 1,  # Separate thread for compute
                "ts": e.start_time_ns / 1000,
                "dur": e.duration_ns / 1000,
                "args": {
                    "stream_id": e.stream_id,
                    "device_id": e.device_id,
                },
            })
    data = {
        "traceEvents": trace_events,
        "metadata": {
            "source": "wafer-distributed-traces",
            "session_name": session.name,
            "world_size": session.world_size,
        },
    }
    json_str = json.dumps(data, indent=2)
    if output_path:
        Path(output_path).write_text(json_str)
    return json_str
def export_heatmap_data(
    matches: list[CollectiveMatch],
    output_path: Optional[str | Path] = None,
) -> str:
    """Export data for heatmap visualization.
    Creates a matrix of rank x collective with duration as values.
    Args:
        matches: Collective matches across ranks
        output_path: Optional path to write to
    Returns:
        JSON string with heatmap data
    """
    if not matches:
        return json.dumps({"ranks": [], "collectives": [], "data": []})

    all_ranks = set()
    for match in matches:
        all_ranks.update(match.collectives.keys())
    ranks = sorted(all_ranks)

    data = {
        "ranks": ranks,
        "collectives": [],
        "data": [],  # 2D array: [rank_idx][collective_idx] = duration_ms
    }

    for _ in ranks:
        data["data"].append([])
    for match in matches:
        data["collectives"].append({
            "sequence": match.sequence_number,
            "type": match.collective_type.value,
            "median_duration_ms": match.median_duration_ns / 1_000_000,
        })
        for rank_idx, rank in enumerate(ranks):
            if rank in match.collectives:
                duration_ms = match.collectives[rank].duration_ms
            else:
                duration_ms = None
            data["data"][rank_idx].append(duration_ms)
    json_str = json.dumps(data, indent=2)
    if output_path:
        Path(output_path).write_text(json_str)
    return json_str
