from a2a.agent.base import BasePlugin, Capability, ModelInfo

class BasePlugin(ABC):
    """Superclass for all A2A plugins."""

    # ── Identity ──
    plugin_id: str            # Unique ID, e.g. "code-fixer"
    plugin_name: str          # Human-readable name
    version: str              # Semantic version, e.g. "1.0.0"

    # ── Lifecycle ──
    async def initialize(self, model, tokenizer, plugin_config: dict, global_config: A2AConfig): ...
    async def on_load(self, runtime: "A2ARuntime"): ...
    async def on_unload(self): ...

    # ── Capabilities ──
    @abstractmethod
    def get_capabilities(self) -> list[Capability]: ...
    @abstractmethod
    def get_model_info(self) -> ModelInfo: ...
    @abstractmethod
    def listens_to(self) -> list[str]: ...
    @abstractmethod
    def emits(self) -> list[str]: ...

    # ── Tensor Processing ──
    @abstractmethod
    async def extract_tensor(self, input_data: Any, semantic_label: str) -> torch.Tensor: ...
    @abstractmethod
    async def on_receive_tensor(self, tensor: torch.Tensor, metadata: A2AMetadata) -> Optional[torch.Tensor]: ...

@dataclass
class Capability:
    name: str              # "log_analysis", "code_generation"
    description: str
    input_labels: list[str]   # Labels this capability accepts
    output_labels: list[str]  # Labels this capability produces
    model_id: str
    hidden_dim: int
    dtype: str             # "float16", "bfloat16", "float32"
    max_batch_size: int
    priority: int = 0

@dataclass
class ModelInfo:
    model_id: str          # "llama-3-8b"
    architecture: str      # "llama", "mistral", "gpt"
    hidden_dim: int        # 4096
    num_layers: int        # 32
    dtype: str             # "float16"
    supported_projection: list[str] = []

class PluginManager:
    def register(self, plugin: BasePlugin): ...
    def unregister(self, plugin_id: str): ...

    async def load_plugin(self, entry: PluginEntry, global_config: A2AConfig) -> BasePlugin: ...
    def discover_from_directory(self, path: str): ...
    def discover_from_entry_points(self, group: str = "a2a.plugins"): ...

    async def route_tensor(self, tensor: torch.Tensor, metadata: A2AMetadata) -> list[BasePlugin]: ...
    def get_matching_plugins(self, required_labels: list[str]) -> list[BasePlugin]: ...

# 1. Direct code registration
manager.register(LogReaderPlugin(model, tokenizer))

# 2. Filesystem discovery
manager.discover_from_directory("./my_plugins/")

# 3. Python entry_points (pip-installable)
manager.discover_from_entry_points("a2a.plugins")

# 4. YAML config-driven
manager.load_from_config("config.yaml")

class TensorExtractor:
    def __init__(self, model: nn.Module, layer_idx: int = -1):
        """
        Args:
            model: HuggingFace model
            layer_idx: Which layer's hidden state (-1 = last layer)
        """

    def extract(self, text: str, pooling: str = "last") -> torch.Tensor:
        """
        Forward pass → extract hidden state at layer_idx.

        Args:
            text: Input text
            pooling: "last" = last token, "mean" = mean pooling, "max" = max pooling

        Returns:
            Tensor of shape (1, hidden_dim)
        """

    def set_layer(self, layer_idx: int): ...
    def get_registered_hooks(self) -> list: ...

class TensorInjector:
    def __init__(self, model: nn.Module, target_layer: int = 0):
        """
        Args:
            model: HuggingFace model
            target_layer: Layer index for injection (0 = prefix, -1 = last)
        """

    def inject_prefix(self, tensor: torch.Tensor, prompt: str) -> dict:
        """
        Prepend tensor as embedding prefix before the prompt.
        Returns the modified inputs dict for model.generate().
        """

    def inject_cross_attention(self, tensor: torch.Tensor, prompt: str) -> dict:
        """
        Use tensor as cross-attention key/value pairs.
        Alternative injection method for models supporting cross-attention.
        """

    async def inject_and_generate(
        self, tensor: torch.Tensor, prompt: str,
        max_tokens: int = 512, temperature: float = 0.7
    ) -> str:
        """Convenience: inject + generate in one call."""

# Safetensors-based save/load
def save_tensor_to_bytes(tensor: torch.Tensor) -> bytes: ...
def load_tensor_from_bytes(data: bytes) -> torch.Tensor: ...

# FlatBuffers codec (zero-copy)
def encode_tensor(tensor: torch.Tensor) -> bytes:
    """PyTorch tensor → FlatBuffers bytes (shape + dtype preserved)"""

def decode_tensor(data: bytes) -> torch.Tensor:
    """FlatBuffers bytes → PyTorch tensor"""

# dtype utilities
def convert_dtype(tensor: torch.Tensor, target_dtype: str) -> torch.Tensor:
    """Convert between FP32, FP16, BF16"""

def validate_tensor(tensor: torch.Tensor, max_l2: float = 1000.0, min_l2: float = 0.001) -> bool:
    """Check for NaN, Inf, bounds"""

class A2AConfig(BaseModel):
    version: str
    runtime: RuntimeConfig
    models: dict[str, ModelConfig]
    transport: TransportConfig
    discovery: DiscoveryConfig
    projection: ProjectionConfig
    plugins: dict[str, PluginEntry]
    routes: dict[str, list[str]]
    security: SecurityConfig
    rate_limit: RateLimitConfig

    @classmethod
    def from_yaml(cls, path: Path | str) -> "A2AConfig": ...
    def validate(self) -> bool: ...


class ModelConfig(BaseModel):
    provider: str           # "openai" | "huggingface" | "ollama" | "anthropic"
    model_id: str
    api_key_env: str | None  # e.g. "DEEPSEEK_API_KEY"
    api_base: str | None
    device: str = "cuda:0"
    dtype: str = "float16"
    max_tokens: int
    load_in_8bit: bool = False


def load_config(path: Path | str | None = None) -> A2AConfig:
    """Searches: A2A_CONFIG env → ./a2a.yaml → ~/.config/a2a/ → /etc/a2a/"""

def load_plugin_config(plugin_path: Path) -> dict:
    """Loads plugin-local config.yaml from plugin directory."""

class ProjectionModel(nn.Module):
    """
    Args:
        src_dim: Source model hidden dimension (e.g. 4096 for Llama-8B)
        tgt_dim: Target model hidden dimension (e.g. 4096 for Mistral-7B)
        hidden_dim: Internal MLP dimension (default: 2048)
    """
    def __init__(self, src_dim: int, tgt_dim: int, hidden_dim: int = 2048): ...
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...


class ProjectionTrainer:
    """Train a projection model using contrastive loss (InfoNCE)."""
    def __init__(self, model: ProjectionModel, lr: float = 0.001, temperature: float = 0.07): ...
    def train(self, dataset: "ProjectionPairDataset", epochs: int = 50) -> dict[str, list[float]]: ...
    def save(self, path: Path): ...


class ProjectionPairDataset(Dataset):
    """Creates (src_hidden, tgt_hidden) pairs from a shared text corpus."""
    @classmethod
    def from_corpus(cls, corpus_path: Path, src_model, src_tokenizer, src_extractor,
                    tgt_model, tgt_tokenizer, tgt_extractor, device="cuda") -> "ProjectionPairDataset": ...


class ProjectionRegistry:
    """Cache and retrieve trained projection models by src__tgt key."""
    def get(self, src_model: str, tgt_model: str) -> ProjectionModel | None: ...
    def load(self, src_model: str, tgt_model: str, path: Path, src_dim: int, tgt_dim: int): ...
    def save(self, src_model: str, tgt_model: str, model: ProjectionModel, path: Path): ...


class AutoTrainer:
    """Runtime auto-training: detect unknown model pair → collect data → train → cache."""
    async def auto_train(self, src_plugin: BasePlugin, tgt_plugin: BasePlugin) -> ProjectionModel: ...

class A2ARuntime:
    def __init__(self, config: A2AConfig): ...
    async def start(self): ...
    async def stop(self): ...
    async def send_tensor(self, tensor: torch.Tensor, metadata: A2AMetadata): ...

async def start(self):
    config = load_config()                    # 1. Load a2a.yaml
    setup_logging(config.runtime)             # 2. Structured logging
    for entry in config.plugins.values():      # 3. Load plugins
        if entry.enabled:
            await plugin_manager.load_plugin(entry, config)
    for pair_key, path in config.projection.pretrained.items():  # 4. Load projections
        projection_registry.load(*pair_key.split("__"), path)
    self.server = A2AServer(config)            # 5. Start gRPC server
    await self.server.start()
    start_metrics_server(config.runtime.metrics_port)  # 6. Health + metrics

# Protobuf metadata
message A2AMetadata:
  source_model: str       # "llama-3-8b-finetuned"
  target_model: str       # "mistral-7b"
  source_layer: uint32    # hidden state extraction layer
  target_layer: uint32    # hidden state injection layer
  tensor_dtype: str       # "float16", "bfloat16", "float32"
  tensor_shape: [uint32]  # [seq_len, hidden_dim]
  semantic_label: str     # "error_context", "code_intent"
  confidence: float       # 0.0 - 1.0
  timestamp: uint64
  session_id: str
  projection_id: str
  requires_projection: bool

# Message types (msg_type)
MSG_UNICAST       = 0x01   # Single tensor send
MSG_MULTICAST     = 0x02   # Broadcast to multiple agents
MSG_DISCOVER      = 0x03   # Find agents on network
MSG_CAPABILITY    = 0x04   # Announce capabilities
MSG_PROJECT_REQ   = 0x05   # Request projection
MSG_PROJECT_RESP  = 0x06   # Projection response
MSG_PROJECT_AUTO  = 0x07   # Trigger auto-training
MSG_PROJECT_READY = 0x08   # Auto-training complete
MSG_KEEPALIVE     = 0x09   # Connection keepalive
MSG_ACK           = 0x0A   # Acknowledgment
MSG_ERROR         = 0x0B   # Error response
MSG_BACKPRESSURE  = 0x0C   # Flow control signal

# Error codes
TENSOR_SHAPE_MISMATCH     = 100
TENSOR_DTYPE_MISMATCH     = 101
TENSOR_VALUE_INVALID      = 102
PROJECTION_NOT_FOUND      = 200
PROJECTION_TRAINING_FAILED = 201
MODEL_NOT_FOUND           = 300
MODEL_OVERLOADED          = 301
AUTH_FAILED               = 400
PERMISSION_DENIED         = 401
MESH_MISMATCH             = 402
INTERNAL_ERROR            = 500
PLUGIN_CRASHED            = 501
GENERATION_FAILED         = 502

class A2AServer:
    def __init__(self, config: TransportConfig, plugin_manager: PluginManager): ...
    async def start(self): ...
    async def stop(self): ...

class A2AClient:
    def __init__(self, address: str, credentials: grpc.ChannelCredentials | None = None): ...
    async def send_tensor(self, tensor: torch.Tensor, metadata: A2AMetadata) -> TensorResponse: ...
    async def stream_tensors(self, tensors: AsyncIterable[torch.Tensor]) -> AsyncIterable[TensorResponse]: ...
    async def health_check(self) -> HealthResponse: ...
    async def discover_agents(self) -> list[AgentInfo]: ...

class PluginManager:
    async def route_tensor(self, tensor: torch.Tensor, metadata: A2AMetadata):
        """
        Routing logic:
        1. Extract semantic_label from metadata
        2. Find all plugins where label is in listens_to()
        3. For each target, check model compatibility
        4. If different model → resolve projection (cached or auto-train)
        5. Forward projected tensor to plugin.on_receive_tensor()
        """

# Example route in a2a.yaml:
# routes:
#   "error_context":
#     - "code-fixer"
#     - "security-analyzer"
#   "code_patch":
#     - "code-reviewer"

# JWT Authentication
def create_token(agent_id: str, mesh_id: str, secret: str, permissions: list[str]) -> str: ...
def validate_token(token: str, secret: str) -> dict | None: ...

# mTLS
class TlsConfig:
    cert_file: str
    key_file: str
    ca_file: str

def create_server_credentials(tls: TlsConfig) -> grpc.ServerCredentials: ...
def create_client_credentials(tls: TlsConfig) -> grpc.ChannelCredentials: ...

class TokenBucket:
    def __init__(self, bucket_size: int, refill_rate: float): ...
    def allow(self) -> bool: ...
    def consume(self, n: int = 1) -> bool: ...

class RateLimiter:
    def get_bucket(self, agent_id: str | None = None, route: str | None = None) -> TokenBucket: ...
    def allow(self, agent_id: str, route: str) -> bool: ...

# Prometheus metrics
a2a_tensors_sent_total          # Counter: tensors sent by agent/label
a2a_tensors_received_total      # Counter: tensors received
a2a_tensor_latency_seconds      # Histogram: end-to-end latency
a2a_tensor_size_bytes           # Histogram: tensor payload size
a2a_projection_requests_total   # Counter: projection lookups
a2a_projection_auto_train_total # Counter: auto-train events
a2a_plugin_active               # Gauge: active plugin count
a2a_plugin_crash_total          # Counter: plugin failures
a2a_rate_limit_hits_total       # Counter: rate limit triggers
a2a_backpressure_events_total   # Counter: flow control signals

# Health endpoints
GET /health           → {"status": "ok", "uptime": 3600, "version": "1.0"}
GET /health/ready     → {"ready": true, "plugins_loaded": 5}
GET /health/live      → {"alive": true}
GET /metrics          → Prometheus scrape endpoint