#
# Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC
# All rights reserved.

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 X.AI Corp.
"""
Segment-Bound Masking for Ranker Attention.

Segment IDs: HISTORY=1, CANDIDATE=-1, PADDING=0
Ranker mask: history attends to all, candidates attend only to self.
"""

from dataclasses import dataclass

import jax
import jax.numpy as jnp

HISTORY_SEGMENT_ID = 1
CANDIDATE_SEGMENT_ID = -1
PADDING_SEGMENT_ID = 0


@dataclass(frozen=True)
class SegmentBounds:
    history_lower: jax.Array
    history_upper: jax.Array
    candidate_lower: jax.Array
    candidate_upper: jax.Array

    @classmethod
    def from_segment_ids(cls, segment_ids: jax.Array) -> "SegmentBounds":
        segment_ids = jnp.asarray(segment_ids)
        if segment_ids.ndim == 1:
            segment_ids = segment_ids[None, :]
        B, S = segment_ids.shape

        history_mask = segment_ids == HISTORY_SEGMENT_ID
        candidate_mask = segment_ids == CANDIDATE_SEGMENT_ID

        def _bounds(mask):
            any_mask = jnp.any(mask, axis=1)
            first = jnp.argmax(mask, axis=1)
            last = (S - 1) - jnp.argmax(mask[:, ::-1], axis=1)
            lower = jnp.where(any_mask, first, 0)
            upper = jnp.where(any_mask, last + 1, 0)
            return lower.astype(jnp.int32), upper.astype(jnp.int32)

        hl, hu = _bounds(history_mask)
        cl, cu = _bounds(candidate_mask)
        return cls(hl, hu, cl, cu)

    def to_array(self) -> jax.Array:
        return jnp.stack([self.history_lower, self.history_upper,
                          self.candidate_lower, self.candidate_upper], axis=1)


def ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds) -> jax.Array:
    q_is_hist = (q_pos >= bounds.history_lower[..., None, None]) & \
                (q_pos < bounds.history_upper[..., None, None])
    q_is_cand = (q_pos >= bounds.candidate_lower[..., None, None]) & \
                (q_pos < bounds.candidate_upper[..., None, None])
    kv_is_hist = (kv_pos >= bounds.history_lower[..., None, None]) & \
                 (kv_pos < bounds.history_upper[..., None, None])
    kv_is_cand = (kv_pos >= bounds.candidate_lower[..., None, None]) & \
                 (kv_pos < bounds.candidate_upper[..., None, None])

    history_mask = kv_is_hist & (q_is_hist | q_is_cand)
    candidate_self_mask = q_is_cand & kv_is_cand & (q_pos == kv_pos)
    return history_mask | candidate_self_mask


def causal_ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds,
                       window_len: int = -1) -> jax.Array:
    base = ranker_mask(q_pos, kv_pos, bounds)
    causal = q_pos >= kv_pos
    if window_len > 0:
        causal = causal & (kv_pos > q_pos - window_len)
    return base & causal