import abc
import torch
import torch.nn as nn
from typing import Dict, Any, Optional

class VLABackbone(nn.Module, abc.ABC):
    """
    Contract for Vision/Language Backbones.
    Must return a feature tensor of shape (B, Seq, Embed_Dim).
    """
    @abc.abstractmethod
    def forward(self, obs: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Args:
            obs: Dictionary containing 'image' and optionally 'text'.
        Returns:
            features: (B, S, D) embedding.
        """
        pass


class VLAProjector(nn.Module, abc.ABC):
    """
    Contract for the adaptation layer (Projector).
    Connects Backbone features to the Policy Head.
    """
    @abc.abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass


class VLAHead(nn.Module, abc.ABC):
    """
    Contract for Action Generation Heads (Policies).
    Handles both training (loss calculation) and inference (action generation).
    """
    @abc.abstractmethod
    def forward(self, embeddings: torch.Tensor, actions: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Args:
            embeddings: (B, S, Hidden) from Projector.
            actions: (B, Pred_Horizon, Action_Dim) - Ground truth for training.
        Returns:
            Dict containing 'loss' (if actions provided) or 'pred_actions'.
        """
        pass