synapse/internal/api/grpcapi/client.go

package grpcapi

import (
	"context"
	"fmt"
	"io"
	"log"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/types/known/anypb"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// =============================================================================
// gRPC Client SDK for Synapse Protocol
// =============================================================================

// GRPCClient is a high-level client for Synapse gRPC services
type GRPCClient struct {
	conn          *grpc.ClientConn
	taskClient    TaskServiceClient
	messageClient A2AMessageServiceClient
	capsClient    CapabilitiesServiceClient
	groupClient   GroupServiceClient
	healthClient  HealthServiceClient

	// Configuration
	apiKey  string
	timeout time.Duration
}

// GRPCClientConfig client configuration
type GRPCClientConfig struct {
	Address  string
	APIKey   string
	Timeout  time.Duration
	Insecure bool // For development only
}

// NewGRPCClient creates a new gRPC client
func NewGRPCClient(config GRPCClientConfig) (*GRPCClient, error) {
	if config.Timeout == 0 {
		config.Timeout = 30 * time.Second
	}

	var opts []grpc.DialOption

	if config.Insecure {
		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	}

	// Add interceptor for API key
	opts = append(opts, grpc.WithUnaryInterceptor(apiKeyInterceptor(config.APIKey)))
	opts = append(opts, grpc.WithStreamInterceptor(apiKeyStreamInterceptor(config.APIKey)))

	ctx, cancel := context.WithTimeout(context.Background(), config.Timeout)
	defer cancel()

	conn, err := grpc.DialContext(ctx, config.Address, opts...)
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %w", err)
	}

	return &GRPCClient{
		conn:          conn,
		taskClient:    NewTaskServiceClient(conn),
		messageClient: NewA2AMessageServiceClient(conn),
		capsClient:    NewCapabilitiesServiceClient(conn),
		groupClient:   NewGroupServiceClient(conn),
		healthClient:  NewHealthServiceClient(conn),
		apiKey:        config.APIKey,
		timeout:       config.Timeout,
	}, nil
}

// Close closes the client connection
func (c *GRPCClient) Close() error {
	return c.conn.Close()
}

// =============================================================================
// Interceptors
// =============================================================================

func apiKeyInterceptor(apiKey string) grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		if apiKey != "" {
			ctx = metadata.AppendToOutgoingContext(ctx, "x-api-key", apiKey)
		}
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

func apiKeyStreamInterceptor(apiKey string) grpc.StreamClientInterceptor {
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		if apiKey != "" {
			ctx = metadata.AppendToOutgoingContext(ctx, "x-api-key", apiKey)
		}
		return streamer(ctx, desc, cc, method, opts...)
	}
}

// =============================================================================
// Task Operations
// =============================================================================

// SubmitTask submits a new task
func (c *GRPCClient) SubmitTask(ctx context.Context, taskType string, payload map[string]any, opts ...TaskOption) (*SubmitTaskResponse, error) {
	req := &SubmitTaskRequest{
		TaskType: taskType,
		Priority: "normal",
	}

	// Apply options
	for _, opt := range opts {
		opt(req)
	}

	// Convert payload to Any
	if payload != nil {
		// In real implementation, use protojson.Marshal
		req.Payload = &anypb.Any{
			Value: []byte(fmt.Sprintf("%v", payload)),
		}
	}

	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.taskClient.SubmitTask(ctx, req)
}

// TaskOption configures task submission
type TaskOption func(*SubmitTaskRequest)

func WithExternalTaskID(id string) TaskOption {
	return func(r *SubmitTaskRequest) {
		r.ExternalTaskId = id
	}
}

func WithPriority(priority string) TaskOption {
	return func(r *SubmitTaskRequest) {
		r.Priority = priority
	}
}

func WithCallback(url, token string) TaskOption {
	return func(r *SubmitTaskRequest) {
		r.CallbackUrl = url
		r.CallbackToken = token
	}
}

func WithTransportMode(mode TransportMode) TaskOption {
	return func(r *SubmitTaskRequest) {
		r.PreferredTransport = mode
	}
}

// GetTask gets task details
func (c *GRPCClient) GetTask(ctx context.Context, taskID string) (*TaskInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.taskClient.GetTask(ctx, &GetTaskRequest{TaskId: taskID})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("task query failed: %s", resp.ErrorMessage)
	}

	return resp.Task, nil
}

// GetTaskStatus gets simplified task status
func (c *GRPCClient) GetTaskStatus(ctx context.Context, taskID string) (*TaskInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.taskClient.GetTaskStatus(ctx, &GetTaskRequest{TaskId: taskID})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("status query failed: %s", resp.ErrorMessage)
	}

	return resp.Task, nil
}

// ListTasks lists tasks with filtering
func (c *GRPCClient) ListTasks(ctx context.Context, status string, limit, offset int) (*ListTasksResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.taskClient.ListTasks(ctx, &ListTasksRequest{
		Status: status,
		Limit:  int32(limit),
		Offset: int32(offset),
	})
}

// CancelTask cancels a task
func (c *GRPCClient) CancelTask(ctx context.Context, taskID string) (*CancelTaskResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.taskClient.CancelTask(ctx, &CancelTaskRequest{TaskId: taskID})
}

// =============================================================================
// Streaming Operations
// =============================================================================

// StreamTaskStatus streams task status updates
func (c *GRPCClient) StreamTaskStatus(ctx context.Context, taskID string) (<-chan *TaskStatusUpdate, error) {
	stream, err := c.taskClient.StreamTaskStatus(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to start stream: %w", err)
	}

	// Send subscription request
	if err := stream.Send(&TaskStatusStreamRequest{
		TaskId:    taskID,
		Subscribe: true,
	}); err != nil {
		return nil, fmt.Errorf("failed to subscribe: %w", err)
	}

	updateChan := make(chan *TaskStatusUpdate, 100)

	go func() {
		defer close(updateChan)
		for {
			update, err := stream.Recv()
			if err == io.EOF {
				return
			}
			if err != nil {
				log.Printf("[gRPC Stream] Receive error: %v", err)
				return
			}
			updateChan <- update
		}
	}()

	return updateChan, nil
}

// StreamMessages opens a bidirectional message stream
func (c *GRPCClient) StreamMessages(ctx context.Context) (chan<- *A2AMessage, <-chan *A2AMessage, error) {
	stream, err := c.messageClient.StreamMessages(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to start message stream: %w", err)
	}

	sendChan := make(chan *A2AMessage, 100)
	recvChan := make(chan *A2AMessage, 100)

	// Send goroutine
	go func() {
		for msg := range sendChan {
			if err := stream.Send(msg); err != nil {
				log.Printf("[gRPC Stream] Send error: %v", err)
				return
			}
		}
	}()

	// Receive goroutine
	go func() {
		defer close(recvChan)
		for {
			msg, err := stream.Recv()
			if err == io.EOF {
				return
			}
			if err != nil {
				log.Printf("[gRPC Stream] Receive error: %v", err)
				return
			}
			recvChan <- msg
		}
	}()

	return sendChan, recvChan, nil
}

// =============================================================================
// Message Operations
// =============================================================================

// SendMessage sends a single Synapse message
func (c *GRPCClient) SendMessage(ctx context.Context, msg *A2AMessage) (*SendMessageResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.messageClient.SendMessage(ctx, &SendMessageRequest{Message: msg})
}

// =============================================================================
// Capabilities Operations
// =============================================================================

// GetCapabilities gets agent capabilities
func (c *GRPCClient) GetCapabilities(ctx context.Context) (*AgentCapabilities, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.capsClient.GetCapabilities(ctx, &emptypb.Empty{})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("capabilities query failed: %s", resp.Error)
	}

	return resp.Capabilities, nil
}

// ListAgents lists available agents
func (c *GRPCClient) ListAgents(ctx context.Context) ([]*AgentInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.capsClient.ListAgents(ctx, &emptypb.Empty{})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("agent list failed")
	}

	return resp.Agents, nil
}

// =============================================================================
// Group Operations
// =============================================================================

// CreateGroup creates a new group
func (c *GRPCClient) CreateGroup(ctx context.Context, name, description, coordinatorID string, initialMembers []string) (*GroupInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.CreateGroup(ctx, &CreateGroupRequest{
		Name:           name,
		Description:    description,
		CoordinatorId:  coordinatorID,
		InitialMembers: initialMembers,
	})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("group creation failed: %s", resp.Error)
	}

	return resp.Group, nil
}

// JoinGroup joins an existing group
func (c *GRPCClient) JoinGroup(ctx context.Context, groupID string, agent *AgentEndpoint) error {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.JoinGroup(ctx, &JoinGroupRequest{
		GroupId: groupID,
		Agent:   agent,
	})
	if err != nil {
		return err
	}

	if !resp.Success {
		return fmt.Errorf("join failed: %s", resp.Error)
	}

	return nil
}

// LeaveGroup leaves a group
func (c *GRPCClient) LeaveGroup(ctx context.Context, groupID, agentID string) error {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.LeaveGroup(ctx, &LeaveGroupRequest{
		GroupId: groupID,
		AgentId: agentID,
	})
	if err != nil {
		return err
	}

	if !resp.Success {
		return fmt.Errorf("leave failed: %s", resp.Error)
	}

	return nil
}

// ListGroupMembers lists group members
func (c *GRPCClient) ListGroupMembers(ctx context.Context, groupID string) ([]*GroupMember, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.ListGroupMembers(ctx, &ListGroupMembersRequest{GroupId: groupID})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("member list failed")
	}

	return resp.Members, nil
}

// =============================================================================
// Voting Operations (Group Protocol)
// =============================================================================

// InitiateVote initiates a vote in a group
func (c *GRPCClient) InitiateVote(ctx context.Context, groupID, initiatorID, question string, options []string, threshold float64, timeoutSeconds int32) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.InitiateVote(ctx, &VoteRequest{
		GroupId:        groupID,
		InitiatorId:    initiatorID,
		Question:       question,
		Options:        options,
		Threshold:      threshold,
		TimeoutSeconds: timeoutSeconds,
	})
	if err != nil {
		return "", err
	}

	if !resp.Success {
		return "", fmt.Errorf("initiate vote failed: %s", resp.Error)
	}

	return resp.VoteId, nil
}

// SubmitVote submits a vote
func (c *GRPCClient) SubmitVote(ctx context.Context, voteID, agentID, option, reason string) error {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.SubmitVote(ctx, &SubmitVoteRequest{
		VoteId:  voteID,
		AgentId: agentID,
		Option:  option,
		Reason:  reason,
	})
	if err != nil {
		return err
	}

	if !resp.Success {
		return fmt.Errorf("submit vote failed: %s", resp.Error)
	}

	return nil
}

// GetVoteStatus gets the current status of a vote
func (c *GRPCClient) GetVoteStatus(ctx context.Context, voteID string) (*VoteStatusResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.groupClient.GetVoteStatus(ctx, &GetVoteStatusRequest{VoteId: voteID})
}

// WaitForVote waits for a vote to complete
func (c *GRPCClient) WaitForVote(ctx context.Context, voteID string, interval time.Duration) (*VoteStatusResponse, error) {
	if interval == 0 {
		interval = 2 * time.Second
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-ticker.C:
			status, err := c.GetVoteStatus(ctx, voteID)
			if err != nil {
				continue
			}

			if status.Status == "decided" || status.Status == "expired" {
				return status, nil
			}
		}
	}
}

// =============================================================================
// Task Split & Merge Operations (Group Protocol)
// =============================================================================

// SplitTask splits a task into subtasks
func (c *GRPCClient) SplitTask(ctx context.Context, groupID, taskID, taskType, requirements string, strategy SplitStrategy) ([]*SubtaskInfo, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	resp, err := c.groupClient.SplitTask(ctx, &SplitTaskRequest{
		GroupId:      groupID,
		TaskId:       taskID,
		TaskType:     taskType,
		Requirements: requirements,
		Strategy:     strategy,
	})
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("split task failed: %s", resp.Error)
	}

	return resp.Subtasks, nil
}

// MergeResults merges subtask results
func (c *GRPCClient) MergeResults(ctx context.Context, groupID, parentTaskID string, strategy MergeStrategy) (*MergeResultsResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.groupClient.MergeResults(ctx, &MergeResultsRequest{
		GroupId:      groupID,
		ParentTaskId: parentTaskID,
		Strategy:     strategy,
	})
}

// =============================================================================
// Health Operations
// =============================================================================

// HealthCheck checks server health
func (c *GRPCClient) HealthCheck(ctx context.Context) (*HealthCheckResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	return c.healthClient.HealthCheck(ctx, &emptypb.Empty{})
}

// =============================================================================
// High-Level Helpers
// =============================================================================

// CreateAndWait creates a task and waits for completion via streaming
func (c *GRPCClient) CreateAndWait(ctx context.Context, taskType string, payload map[string]any, timeout time.Duration) (*TaskInfo, error) {
	// Submit task
	resp, err := c.SubmitTask(ctx, taskType, payload)
	if err != nil {
		return nil, err
	}

	if !resp.Success {
		return nil, fmt.Errorf("task submission failed: %s", resp.Message)
	}

	// Stream status until completion
	updateChan, err := c.StreamTaskStatus(ctx, resp.TaskId)
	if err != nil {
		return nil, err
	}

	// Wait with timeout
	if timeout == 0 {
		timeout = 5 * time.Minute
	}
	timer := time.NewTimer(timeout)
	defer timer.Stop()

	for {
		select {
		case update := <-updateChan:
			if update == nil {
				return nil, fmt.Errorf("stream closed unexpectedly")
			}

			switch update.Status {
			case TaskStatus_TASK_STATUS_COMPLETED:
				return &TaskInfo{
					Id:       resp.TaskId,
					Status:   update.Status,
					Progress: update.Progress,
					Message:  update.Message,
				}, nil
			case TaskStatus_TASK_STATUS_FAILED, TaskStatus_TASK_STATUS_CANCELLED:
				return nil, fmt.Errorf("task failed: %s", update.Message)
			}

		case <-timer.C:
			return nil, fmt.Errorf("timeout waiting for task completion")
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

// WaitForTask waits for a task to complete using polling
func (c *GRPCClient) WaitForTask(ctx context.Context, taskID string, interval time.Duration) (*TaskInfo, error) {
	if interval == 0 {
		interval = 5 * time.Second
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-ticker.C:
			task, err := c.GetTaskStatus(ctx, taskID)
			if err != nil {
				continue
			}

			switch task.Status {
			case TaskStatus_TASK_STATUS_COMPLETED:
				return task, nil
			case TaskStatus_TASK_STATUS_FAILED, TaskStatus_TASK_STATUS_CANCELLED:
				return task, fmt.Errorf("task failed: %s", task.Message)
			}
		}
	}
}

// =============================================================================
// Builder Pattern for A2AMessage
// =============================================================================

// MessageBuilder builds Synapse messages
type MessageBuilder struct {
	msg *A2AMessage
}

// NewMessage creates a new message builder
func NewMessage() *MessageBuilder {
	return &MessageBuilder{
		msg: &A2AMessage{
			Id:        generateMessageID(),
			Version:   "1.0.0",
			Timestamp: timestamppb.Now(),
		},
	}
}

func (b *MessageBuilder) Type(msgType MessageType) *MessageBuilder {
	b.msg.Type = msgType
	return b
}

func (b *MessageBuilder) From(sender *AgentEndpoint) *MessageBuilder {
	b.msg.Sender = sender
	return b
}

func (b *MessageBuilder) To(receiver *AgentEndpoint) *MessageBuilder {
	b.msg.Receiver = receiver
	return b
}

func (b *MessageBuilder) CorrelationID(id string) *MessageBuilder {
	b.msg.CorrelationId = id
	return b
}

func (b *MessageBuilder) GroupID(id string) *MessageBuilder {
	b.msg.GroupId = id
	return b
}

func (b *MessageBuilder) Payload(payload []byte) *MessageBuilder {
	b.msg.Payload = &anypb.Any{Value: payload}
	return b
}

func (b *MessageBuilder) Build() *A2AMessage {
	return b.msg
}

func generateMessageID() string {
	return fmt.Sprintf("msg-%d", time.Now().UnixNano())
}