"""Per-rank event timeline representation.
Defines data structures for representing a timeline of events (both
communication and compute) for a single rank.
"""
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective
class WindowType(Enum):
    """Type of timeline window."""
    COMPUTE_ONLY = "compute_only"
    COMMUNICATION_ONLY = "communication_only"
    OVERLAPPED = "overlapped"
    IDLE = "idle"
@dataclass(frozen=True)
class ComputeEvent:
    """Represents a compute operation (GPU kernel not related to NCCL/RCCL).
    Attributes:
        name: Kernel or operation name
        start_time_ns: Start timestamp in nanoseconds
        end_time_ns: End timestamp in nanoseconds
        stream_id: CUDA/HIP stream ID
        device_id: GPU device ID
        grid_size: Grid dimensions (x, y, z)
        block_size: Block dimensions (x, y, z)
        shared_memory_bytes: Shared memory used
        registers_per_thread: Registers per thread
        correlation_id: Correlation ID for linking to API calls
        pytorch_op_name: Name of PyTorch operator
        extra: Additional metadata
    """
    name: str
    start_time_ns: int
    end_time_ns: int
    stream_id: int = 0
    device_id: int = 0
    grid_size: Optional[tuple[int, int, int]] = None
    block_size: Optional[tuple[int, int, int]] = None
    shared_memory_bytes: int = 0
    registers_per_thread: int = 0
    correlation_id: Optional[int] = None
    pytorch_op_name: Optional[str] = None
    extra: dict = field(default_factory=dict)
    @property
    def duration_ns(self) -> int:
        """Duration in nanoseconds."""
        return self.end_time_ns - self.start_time_ns
    @property
    def duration_ms(self) -> float:
        """Duration in milliseconds."""
        return self.duration_ns / 1_000_000
@dataclass(frozen=True)
class TimelineEvent:
    """Union type for timeline events.
    A timeline event can be either a Collective or a ComputeEvent.
    Attributes:
        event: The actual event (Collective or ComputeEvent)
        is_collective: True if this is a collective operation
    """
    event: Collective | ComputeEvent
    is_collective: bool
    @property
    def start_time_ns(self) -> int:
        """Start timestamp in nanoseconds."""
        return self.event.start_time_ns
    @property
    def end_time_ns(self) -> int:
        """End timestamp in nanoseconds."""
        return self.event.end_time_ns
    @property
    def duration_ns(self) -> int:
        """Duration in nanoseconds."""
        return self.event.duration_ns
    @classmethod
    def from_collective(cls, collective: Collective) -> "TimelineEvent":
        """Create a TimelineEvent from a Collective."""
        return cls(event=collective, is_collective=True)
    @classmethod
    def from_compute(cls, compute: ComputeEvent) -> "TimelineEvent":
        """Create a TimelineEvent from a ComputeEvent."""
        return cls(event=compute, is_collective=False)
@dataclass(frozen=True)
class TimelineWindow:
    """Represents a time window with a specific type.
    Used for overlap analysis to segment the timeline into compute-only,
    communication-only, overlapped, and idle windows.
    Attributes:
        window_type: Type of this window
        start_time_ns: Start timestamp in nanoseconds
        end_time_ns: End timestamp in nanoseconds
        events: Events active during this window
    """
    window_type: WindowType
    start_time_ns: int
    end_time_ns: int
    events: tuple[TimelineEvent, ...] = field(default_factory=tuple)
    @property
    def duration_ns(self) -> int:
        """Duration in nanoseconds."""
        return self.end_time_ns - self.start_time_ns
    @property
    def duration_ms(self) -> float:
        """Duration in milliseconds."""
        return self.duration_ns / 1_000_000
@dataclass
class RankTimeline:
    """Timeline of events for a single rank.
    Contains all collectives and compute events for one rank,
    with methods for querying and analyzing the timeline.
    Note: This class is mutable to allow incremental building of timelines.
    Attributes:
        rank: Global rank number
        device_id: Local GPU device ID
        hostname: Hostname where this rank ran
        pid: Process ID
        collectives: List of collective operations
        compute_events: List of compute operations
        clock_offset_ns: Clock offset for alignment with other ranks
    """
    rank: int
    device_id: int = 0
    hostname: Optional[str] = None
    pid: Optional[int] = None
    collectives: list[Collective] = field(default_factory=list)
    compute_events: list[ComputeEvent] = field(default_factory=list)
    clock_offset_ns: int = 0
    @property
    def start_time_ns(self) -> int:
        """Earliest timestamp in this timeline."""
        times = []
        if self.collectives:
            times.append(min(c.start_time_ns for c in self.collectives))
        if self.compute_events:
            times.append(min(e.start_time_ns for e in self.compute_events))
        return min(times) if times else 0
    @property
    def end_time_ns(self) -> int:
        """Latest timestamp in this timeline."""
        times = []
        if self.collectives:
            times.append(max(c.end_time_ns for c in self.collectives))
        if self.compute_events:
            times.append(max(e.end_time_ns for e in self.compute_events))
        return max(times) if times else 0
    @property
    def duration_ns(self) -> int:
        """Total duration of this timeline."""
        return self.end_time_ns - self.start_time_ns
    @property
    def total_collective_time_ns(self) -> int:
        """Total time spent in collectives."""
        return sum(c.duration_ns for c in self.collectives)
    @property
    def total_compute_time_ns(self) -> int:
        """Total time spent in compute."""
        return sum(e.duration_ns for e in self.compute_events)
    def add_collective(self, collective: Collective) -> None:
        """Add a collective to the timeline."""
        self.collectives.append(collective)
    def add_compute_event(self, event: ComputeEvent) -> None:
        """Add a compute event to the timeline."""
        self.compute_events.append(event)
    def get_events_sorted(self) -> list[TimelineEvent]:
        """Get all events sorted by start time."""
        events = []
        for c in self.collectives:
            events.append(TimelineEvent.from_collective(c))
        for e in self.compute_events:
            events.append(TimelineEvent.from_compute(e))
        return sorted(events, key=lambda e: e.start_time_ns)
    def get_collectives_by_type(
        self,
        collective_type: "CollectiveType",  # noqa: F821
    ) -> list[Collective]:
        """Get collectives of a specific type."""
        from wafer_core.lib.distributed_traces.models.collective import CollectiveType
        return [c for c in self.collectives if c.collective_type == collective_type]
    def apply_clock_offset(self) -> None:
        """Apply clock offset to all events.
        This modifies the timeline in place to adjust for clock skew.
        Should only be called once after alignment.
        """
        if self.clock_offset_ns == 0:
            return

        # Since Collective is frozen, we need to create new instances
        from dataclasses import replace
        self.collectives = [
            replace(
                c,
                start_time_ns=c.start_time_ns + self.clock_offset_ns,
                end_time_ns=c.end_time_ns + self.clock_offset_ns,
            )
            for c in self.collectives
        ]
        self.compute_events = [
            replace(
                e,
                start_time_ns=e.start_time_ns + self.clock_offset_ns,
                end_time_ns=e.end_time_ns + self.clock_offset_ns,
            )
            for e in self.compute_events
        ]
        # Reset offset after applying
        self.clock_offset_ns = 0
