synapse/internal/agent/p2p/protocol_test.go

package p2p

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"synapse/pkg/types"
)

// =============================================================================
// Mock Task Handler
// =============================================================================

type mockTaskHandler struct {
	canHandle     bool
	missingFields []types.MissingField
	result        any
	err           error
}

func (m *mockTaskHandler) ExecuteTask(ctx context.Context, taskType string, payload map[string]interface{}) (any, error) {
	return m.result, m.err
}

func (m *mockTaskHandler) ValidateTask(taskType string, payload map[string]interface{}) ([]types.MissingField, bool) {
	return m.missingFields, m.canHandle
}

func (m *mockTaskHandler) GetCapabilities() []string {
	return []string{"text_generation", "image_generation"}
}

// =============================================================================
// Mock Callback Sender
// =============================================================================

type mockCallbackSender struct {
	lastMsg *types.UnifiedSynapseMessage
	err     error
}

func (m *mockCallbackSender) SendCallback(ctx context.Context, endpoint string, msg *types.UnifiedSynapseMessage) error {
	m.lastMsg = msg
	return m.err
}

// =============================================================================
// P2PProtocol 测试
// =============================================================================

func TestNewP2PProtocol(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	if proto == nil {
		t.Error("Protocol should not be nil")
	}
	if proto.AgentEndpoint().AgentID != "agent-001" {
		t.Error("Agent ID mismatch")
	}
}

func TestP2PProtocolWithCustomConfig(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	config := &P2PConfig{
		MaxRetries:        10,
		RetryBaseDelay:    5 * time.Second,
		TaskTimeout:       10 * time.Minute,
		HeartbeatInterval: 60 * time.Second,
	}

	proto := NewP2PProtocol(ep, config)
	if proto.config.MaxRetries != 10 {
		t.Errorf("MaxRetries mismatch: got %d, want 10", proto.config.MaxRetries)
	}
}

func TestDefaultP2PConfig(t *testing.T) {
	config := DefaultP2PConfig()
	if config.MaxRetries != 5 {
		t.Errorf("Default MaxRetries should be 5")
	}
	if config.TaskTimeout != 5*time.Minute {
		t.Errorf("Default TaskTimeout should be 5 minutes")
	}
}

// =============================================================================
// 消息处理测试
// =============================================================================

func TestHandleTaskRequest(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	proto.SetTaskHandler(&mockTaskHandler{
		canHandle: true,
		result:    "test result",
	})

	// 构造任务请求消息
	taskReq := types.TaskRequestPayload{}
	taskReq.Task.ID = "task-001"
	taskReq.Task.Type = "text_generation"
	taskReq.Task.Priority = types.PriorityNormal
	taskReq.Task.Payload = map[string]interface{}{"prompt": "test"}

	payloadBytes, _ := json.Marshal(taskReq)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	resp, err := proto.HandleMessage(context.Background(), msg)
	if err != nil {
		t.Errorf("HandleMessage failed: %v", err)
	}
	if resp == nil {
		t.Error("Response should not be nil")
	}
}

func TestHandleTaskReject(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	proto.SetTaskHandler(&mockTaskHandler{
		canHandle:     false,
		missingFields: nil,
	})

	// 构造任务请求消息
	taskReq := types.TaskRequestPayload{}
	taskReq.Task.ID = "task-001"
	taskReq.Task.Type = "unsupported_type"
	taskReq.Task.Priority = types.PriorityNormal

	payloadBytes, _ := json.Marshal(taskReq)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	resp, err := proto.HandleMessage(context.Background(), msg)
	if err != nil {
		t.Errorf("HandleMessage should not return error for rejection: %v", err)
	}
	if resp == nil {
		t.Error("Response should not be nil")
	}
	// 应该返回拒绝消息
	if resp.Type != string(types.P2PTaskReject) {
		t.Errorf("Expected reject response, got: %s", resp.Type)
	}
}

func TestHandleHeartbeat(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PHeartbeat),
		Sender:    types.AgentEndpoint{AgentID: "agent-002"},
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	resp, err := proto.HandleMessage(context.Background(), msg)
	if err != nil {
		t.Errorf("HandleHeartbeat failed: %v", err)
	}
	if resp == nil {
		t.Error("Heartbeat response should not be nil")
	}
}

func TestHandleAck(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PAck),
		Sender:    types.AgentEndpoint{AgentID: "agent-002"},
		ReplyTo:   "msg-001",
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	resp, err := proto.HandleMessage(context.Background(), msg)
	if err != nil {
		t.Errorf("HandleAck failed: %v", err)
	}
	// ACK 不需要响应
	if resp != nil {
		t.Error("ACK should not produce a response")
	}
}

// =============================================================================
// 消息验证测试
// =============================================================================

func TestValidateMessage(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	// 有效消息
	validMsg := &types.UnifiedSynapseMessage{
		ID:        "msg-001",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	if err := proto.validateMessage(validMsg); err != nil {
		t.Errorf("Valid message should pass validation: %v", err)
	}

	// 缺少 ID 的消息
	noIDMsg := &types.UnifiedSynapseMessage{
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	if err := proto.validateMessage(noIDMsg); err == nil {
		t.Error("Message without ID should fail validation")
	}

	// 缺少 Sender AgentID 的消息
	noSenderMsg := &types.UnifiedSynapseMessage{
		ID:        "msg-001",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{Endpoint: "http://localhost:8081"},
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	if err := proto.validateMessage(noSenderMsg); err == nil {
		t.Error("Message without sender AgentID should fail validation")
	}
}

// =============================================================================
// 状态管理测试
// =============================================================================

func TestGetTaskStatus(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	// 查询不存在的任务
	status, ok := proto.GetTaskStatus("nonexistent")
	if ok {
		t.Error("Nonexistent task should not be found")
	}

	// 先通过任务请求添加任务状态
	taskReq := types.TaskRequestPayload{}
	taskReq.Task.ID = "task-001"
	taskReq.Task.Type = "text_generation"
	taskReq.Task.Priority = types.PriorityNormal
	taskReq.Task.Payload = map[string]interface{}{"prompt": "test"}

	payloadBytes, _ := json.Marshal(taskReq)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	proto.HandleMessage(context.Background(), msg)

	// 然后查询状态
	status, ok = proto.GetTaskStatus("task-001")
	if !ok {
		t.Error("Task should be found")
	}
	// 状态应该是 pending 或 confirmed
	if status != types.P2PStatusPending && status != types.P2PStatusConfirmed {
		t.Errorf("Status should be pending or confirmed: got %s", status)
	}
}

func TestListTasks(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	// 初始应该为空
	tasks := proto.ListTasks()
	if len(tasks) != 0 {
		t.Error("Initial task list should be empty")
	}

	// 添加任务（通过消息处理）
	taskReq := types.TaskRequestPayload{}
	taskReq.Task.ID = "task-001"
	taskReq.Task.Type = "text_generation"
	taskReq.Task.Priority = types.PriorityNormal

	payloadBytes, _ := json.Marshal(taskReq)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	proto.HandleMessage(context.Background(), msg)

	tasks = proto.ListTasks()
	if len(tasks) < 1 {
		t.Errorf("Task list should have at least 1 item: got %d", len(tasks))
	}
}

// =============================================================================
// 回调发送测试
// =============================================================================

func TestSendStatusUpdate(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	mockSender := &mockCallbackSender{}
	proto.SetCallbackSender(mockSender)

	// 发送状态更新
	proto.sendStatusUpdate(context.Background(), "task-001", types.P2PStatusExecuting, 50, "Processing")

	if mockSender.lastMsg == nil {
		t.Error("Callback should have been sent")
	}
	if mockSender.lastMsg.Type != string(types.P2PTaskStatusUpdate) {
		t.Errorf("Callback type mismatch: got %s", mockSender.lastMsg.Type)
	}
}

func TestSendTaskComplete(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	mockSender := &mockCallbackSender{}
	proto.SetCallbackSender(mockSender)

	// 发送完成通知
	proto.sendTaskComplete(context.Background(), "task-001", "test result")

	if mockSender.lastMsg == nil {
		t.Error("Callback should have been sent")
	}
	if mockSender.lastMsg.Type != string(types.P2PTaskComplete) {
		t.Errorf("Callback type mismatch: got %s", mockSender.lastMsg.Type)
	}
}

func TestSendTaskFailed(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)
	mockSender := &mockCallbackSender{}
	proto.SetCallbackSender(mockSender)

	// 发送失败通知
	proto.sendTaskFailed(context.Background(), "task-001", "execution_failed", "test error", true, nil)

	if mockSender.lastMsg == nil {
		t.Error("Callback should have been sent")
	}
	if mockSender.lastMsg.Type != string(types.P2PTaskFailed) {
		t.Errorf("Callback type mismatch: got %s", mockSender.lastMsg.Type)
	}
}

// =============================================================================
// 错误处理测试
// =============================================================================

func TestHandleInvalidPayload(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskRequest),
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   json.RawMessage("invalid json"),
		Timestamp: types.NowTimestamp(),
	}

	resp, err := proto.HandleMessage(context.Background(), msg)
	if err != nil {
		t.Errorf("Should handle invalid payload gracefully: %v", err)
	}
	// 应该返回 NACK
	if resp == nil {
		t.Error("Should return NACK for invalid payload")
	}
}

func TestHandleUnknownMessageType(t *testing.T) {
	ep := types.AgentEndpoint{
		AgentID:  "agent-001",
		Endpoint: "http://localhost:8080",
	}

	proto := NewP2PProtocol(ep, nil)

	msg := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      "p2p:unknown_type",
		Sender:    types.AgentEndpoint{AgentID: "agent-002", Endpoint: "http://localhost:8081"},
		Payload:   json.RawMessage("{}"),
		Timestamp: types.NowTimestamp(),
	}

	_, err := proto.HandleMessage(context.Background(), msg)
	if err == nil {
		t.Error("Unknown message type should return error")
	}
}