"""Match collectives across ranks by type, size, ordering.

This module provides functionality to identify the same logical collective
operation across multiple ranks. This is essential for:
- Straggler detection (finding which rank was slowest)
- Cross-rank timing analysis
- Validating collective correctness

Matching Algorithm:
1. Group collectives by process group
2. Within each group, match by (collective_type, sequence_number, message_size)
3. Handle missing sequence numbers by falling back to ordering
"""

from dataclasses import dataclass
from typing import Optional

from wafer_core.lib.distributed_traces.models.collective import Collective, CollectiveType, ProcessGroup
from wafer_core.lib.distributed_traces.models.trace_session import CollectiveMatch, TraceSession


@dataclass
class MatchingConfig:
    """Configuration for collective matching.

    Attributes:
        size_tolerance: Tolerance for message size matching (0.0-1.0)
        require_all_ranks: If True, only match when all ranks participate
        max_time_diff_ns: Maximum time difference between matched collectives
    """

    size_tolerance: float = 0.1  # 10% tolerance
    require_all_ranks: bool = False
    max_time_diff_ns: Optional[int] = None  # No limit by default


class CollectiveMatcher:
    """Matches collective operations across ranks.

    Identifies the same logical collective operation on different ranks
    to enable cross-rank analysis and straggler detection.
    """

    def __init__(
        self,
        session: TraceSession,
        config: Optional[MatchingConfig] = None,
    ) -> None:
        """Initialize the matcher.

        Args:
            session: TraceSession with aligned timelines
            config: Optional matching configuration
        """
        self.session = session
        self.config = config or MatchingConfig()
        self._matches: list[CollectiveMatch] = []

    def match_all(self) -> list[CollectiveMatch]:
        """Match all collectives across ranks.

        Returns:
            List of CollectiveMatch objects, each representing one
            logical collective operation across multiple ranks.
        """
        self._matches = []

        # Get all unique process groups
        process_groups = self._get_process_groups()

        # For each process group, match collectives
        for pg_id in process_groups:
            pg_matches = self._match_within_process_group(pg_id)
            self._matches.extend(pg_matches)

        return self._matches

    def _get_process_groups(self) -> set[str]:
        """Get all unique process group IDs in the session."""
        pg_ids: set[str] = set()

        for timeline in self.session.timelines.values():
            for coll in timeline.collectives:
                if coll.process_group:
                    pg_ids.add(coll.process_group.comm_id)
                else:
                    # Use empty string for default/unknown process group
                    pg_ids.add("")

        return pg_ids

    def _match_within_process_group(self, pg_id: str) -> list[CollectiveMatch]:
        """Match collectives within a single process group.

        Args:
            pg_id: Process group ID to match within

        Returns:
            List of CollectiveMatch objects for this process group
        """
        # Collect collectives by rank for this process group
        colls_by_rank: dict[int, list[Collective]] = {}

        for rank, timeline in self.session.timelines.items():
            colls = []
            for c in timeline.collectives:
                coll_pg_id = c.process_group.comm_id if c.process_group else ""
                if coll_pg_id == pg_id:
                    colls.append(c)

            # Sort by start time
            colls.sort(key=lambda c: c.start_time_ns)
            colls_by_rank[rank] = colls

        if not colls_by_rank:
            return []

        # Match using sequence numbers if available
        if self._has_sequence_numbers(colls_by_rank):
            return self._match_by_sequence(colls_by_rank, pg_id)
        else:
            return self._match_by_ordering(colls_by_rank, pg_id)

    def _has_sequence_numbers(self, colls_by_rank: dict[int, list[Collective]]) -> bool:
        """Check if collectives have sequence numbers."""
        for colls in colls_by_rank.values():
            for c in colls:
                if c.sequence_number is not None:
                    return True
        return False

    def _match_by_sequence(
        self,
        colls_by_rank: dict[int, list[Collective]],
        pg_id: str,
    ) -> list[CollectiveMatch]:
        """Match collectives by sequence number.

        Args:
            colls_by_rank: Collectives grouped by rank
            pg_id: Process group ID

        Returns:
            List of matched collectives
        """
        # Group by (collective_type, sequence_number)
        groups: dict[tuple[CollectiveType, int], dict[int, Collective]] = {}

        for rank, colls in colls_by_rank.items():
            for c in colls:
                if c.sequence_number is None:
                    continue

                key = (c.collective_type, c.sequence_number)
                if key not in groups:
                    groups[key] = {}
                groups[key][rank] = c

        # Create matches from groups
        matches = []
        for (ctype, seq), rank_colls in sorted(groups.items(), key=lambda x: x[0][1]):
            if len(rank_colls) < 2:
                continue

            # Validate message size similarity
            if not self._validate_sizes(list(rank_colls.values())):
                continue

            # Check time difference if configured
            if not self._validate_time_diff(list(rank_colls.values())):
                continue

            # Check require_all_ranks
            if self.config.require_all_ranks:
                if len(rank_colls) != self.session.num_ranks:
                    continue

            # Get representative message size
            msg_size = max(c.message_size_bytes for c in rank_colls.values())

            match = CollectiveMatch(
                collective_type=ctype,
                sequence_number=seq,
                process_group_id=pg_id,
                collectives=rank_colls,
                message_size_bytes=msg_size,
            )
            matches.append(match)

        return matches

    def _match_by_ordering(
        self,
        colls_by_rank: dict[int, list[Collective]],
        pg_id: str,
    ) -> list[CollectiveMatch]:
        """Match collectives by their ordering within each rank.

        Fallback method when sequence numbers are not available.
        Assumes collectives of the same type appear in the same order
        across all ranks.

        Args:
            colls_by_rank: Collectives grouped by rank
            pg_id: Process group ID

        Returns:
            List of matched collectives
        """
        # Group by collective type first
        by_type: dict[CollectiveType, dict[int, list[Collective]]] = {}

        for rank, colls in colls_by_rank.items():
            for c in colls:
                if c.collective_type not in by_type:
                    by_type[c.collective_type] = {}
                if rank not in by_type[c.collective_type]:
                    by_type[c.collective_type][rank] = []
                by_type[c.collective_type][rank].append(c)

        matches = []
        seq_counter = 0

        for ctype, rank_colls in by_type.items():
            # Find minimum count across ranks for this type
            min_count = min(len(colls) for colls in rank_colls.values())

            # Match by position
            for i in range(min_count):
                matched = {}
                for rank, colls in rank_colls.items():
                    if i < len(colls):
                        matched[rank] = colls[i]

                if len(matched) < 2:
                    continue

                # Validate sizes and time
                if not self._validate_sizes(list(matched.values())):
                    continue
                if not self._validate_time_diff(list(matched.values())):
                    continue

                if self.config.require_all_ranks:
                    if len(matched) != self.session.num_ranks:
                        continue

                msg_size = max(c.message_size_bytes for c in matched.values())

                match = CollectiveMatch(
                    collective_type=ctype,
                    sequence_number=seq_counter,
                    process_group_id=pg_id,
                    collectives=matched,
                    message_size_bytes=msg_size,
                )
                matches.append(match)
                seq_counter += 1

        return matches

    def _validate_sizes(self, collectives: list[Collective]) -> bool:
        """Validate that message sizes are similar."""
        sizes = [c.message_size_bytes for c in collectives if c.message_size_bytes > 0]

        if len(sizes) < 2:
            return True  # Can't validate, assume OK

        min_size = min(sizes)
        max_size = max(sizes)

        if min_size == 0:
            return True

        ratio = max_size / min_size
        return ratio <= (1 + self.config.size_tolerance)

    def _validate_time_diff(self, collectives: list[Collective]) -> bool:
        """Validate that time differences are within limit."""
        if self.config.max_time_diff_ns is None:
            return True

        starts = [c.start_time_ns for c in collectives]
        time_diff = max(starts) - min(starts)

        return time_diff <= self.config.max_time_diff_ns


def match_collectives_across_ranks(
    session: TraceSession,
    config: Optional[MatchingConfig] = None,
) -> list[CollectiveMatch]:
    """Convenience function to match collectives in a session.

    Args:
        session: TraceSession with timelines
        config: Optional matching configuration

    Returns:
        List of CollectiveMatch objects
    """
    matcher = CollectiveMatcher(session, config)
    return matcher.match_all()
