import torch
import torch.nn.functional as F

try:
    from . import _C
    _HAS_METAL = True
except ImportError:
    _HAS_METAL = False


def causal_conv1d_fn(x, weight, bias=None, activation=None, seq_idx=None):
    """Drop-in for causal_conv1d.causal_conv1d_fn on MPS.

    x: (B, D, T), weight: (D, W), bias: (D,) or None
    Returns: (B, D, T)
    """
    _, _, T = x.shape
    W = weight.shape[1]

    if _HAS_METAL and x.is_mps and x.dtype in (torch.float16, torch.bfloat16, torch.float32):
        bias_ = bias if (bias is not None) else torch.empty(0, device=x.device, dtype=x.dtype)
        # Metal kernel has SiLU fused in — activation is always silu in Qwen3.5
        apply_silu = activation in ("silu", "swish")
        out = _C.causal_conv1d(x.contiguous(), weight.contiguous(), bias_, apply_silu)
        if activation not in (None, "silu", "swish"):
            raise ValueError(f"Unsupported activation: {activation}")
        return out

    # Fallback: pure PyTorch
    out = F.conv1d(x, weight.unsqueeze(1), bias, padding=W - 1, groups=x.shape[1])[..., :T]
    if activation in ("silu", "swish"):
        out = F.silu(out)
    return out


def causal_conv1d_update(hidden_states, conv_state, weight, bias=None, activation=None):
    """Drop-in for causal_conv1d.causal_conv1d_update on MPS.

    hidden_states: (B, D, L), conv_state: (B, D, W)  [mutated in-place]
    Returns: (B, D, L)
    """
    if _HAS_METAL and hidden_states.is_mps and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32):
        bias_ = bias if (bias is not None) else torch.empty(0, device=hidden_states.device, dtype=hidden_states.dtype)
        apply_silu = activation in ("silu", "swish")
        # Metal kernel mutates conv_state in-place and returns output
        out = _C.causal_conv1d_update(
            hidden_states.contiguous().to(weight.dtype),
            conv_state,
            weight.contiguous(),
            bias_,
            apply_silu,
        )
        return out.to(hidden_states.dtype)

    # Fallback: pure PyTorch
    _, _, seq_len = hidden_states.shape
    state_len = conv_state.shape[-1]
    combined = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
    conv_state.copy_(combined[:, :, -state_len:])
    out = F.conv1d(combined, weight.unsqueeze(1), bias, padding=0, groups=hidden_states.shape[1])
    out = out[:, :, -seq_len:]
    if activation in ("silu", "swish"):
        out = F.silu(out)
    return out.to(hidden_states.dtype)
