bzzz/pkg/mcp/server.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"chorus.services/bzzz/logging"
	"chorus.services/bzzz/p2p"
	"chorus.services/bzzz/pubsub"
	"github.com/gorilla/websocket"
	"github.com/sashabaranov/go-openai"
)

// McpServer integrates BZZZ P2P network with MCP protocol for GPT-4 agents
type McpServer struct {
	// Core components
	p2pNode      *p2p.Node
	pubsub       *pubsub.PubSub
	hlog         *logging.HypercoreLog
	openaiClient *openai.Client
	
	// Agent management
	agents       map[string]*GPTAgent
	agentsMutex  sync.RWMutex
	
	// Server configuration
	httpServer   *http.Server
	wsUpgrader   websocket.Upgrader
	
	// Context and lifecycle
	ctx          context.Context
	cancel       context.CancelFunc
	
	// Statistics and monitoring
	stats        *ServerStats
}

// ServerStats tracks MCP server performance metrics
type ServerStats struct {
	StartTime           time.Time
	TotalRequests       int64
	ActiveAgents        int
	MessagesProcessed   int64
	TokensConsumed      int64
	AverageCostPerTask  float64
	ErrorRate           float64
	mutex               sync.RWMutex
}

// GPTAgent represents a GPT-4 agent integrated with BZZZ network
type GPTAgent struct {
	ID              string
	Role            AgentRole
	Model           string
	SystemPrompt    string
	Capabilities    []string
	Specialization  string
	MaxTasks        int
	
	// State management
	Status          AgentStatus
	CurrentTasks    map[string]*AgentTask
	Memory          *AgentMemory
	
	// Cost tracking
	TokenUsage      *TokenUsage
	CostLimits      *CostLimits
	
	// P2P Integration
	NodeID          string
	LastAnnouncement time.Time
	
	// Conversation participation
	ActiveThreads   map[string]*ConversationThread
	
	mutex           sync.RWMutex
}

// AgentRole defines the role and responsibilities of an agent
type AgentRole string

const (
	RoleArchitect      AgentRole = "architect"
	RoleReviewer       AgentRole = "reviewer"
	RoleDocumentation  AgentRole = "documentation"
	RoleDeveloper      AgentRole = "developer"
	RoleTester         AgentRole = "tester"
	RoleSecurityExpert AgentRole = "security_expert"
	RoleDevOps         AgentRole = "devops"
)

// AgentStatus represents the current state of an agent
type AgentStatus string

const (
	StatusIdle         AgentStatus = "idle"
	StatusActive       AgentStatus = "active"
	StatusCollaborating AgentStatus = "collaborating"
	StatusEscalating   AgentStatus = "escalating"
	StatusTerminating  AgentStatus = "terminating"
)

// AgentTask represents a task being worked on by an agent
type AgentTask struct {
	ID          string
	Title       string
	Repository  string
	Number      int
	StartTime   time.Time
	Status      string
	ThreadID    string
	Context     map[string]interface{}
}

// AgentMemory manages agent memory and learning
type AgentMemory struct {
	WorkingMemory    map[string]interface{}
	EpisodicMemory   []ConversationEpisode
	SemanticMemory   *KnowledgeGraph
	ThreadMemories   map[string]*ThreadMemory
	mutex            sync.RWMutex
}

// ConversationEpisode represents a past interaction
type ConversationEpisode struct {
	Timestamp    time.Time
	Participants []string
	Topic        string
	Summary      string
	Outcome      string
	Lessons      []string
	TokensUsed   int
}

// ConversationThread represents an active conversation
type ConversationThread struct {
	ID           string
	Topic        string
	Participants []AgentParticipant
	Messages     []ThreadMessage
	State        ThreadState
	SharedContext map[string]interface{}
	DecisionLog   []Decision
	CreatedAt     time.Time
	LastActivity  time.Time
	mutex         sync.RWMutex
}

// AgentParticipant represents an agent participating in a conversation
type AgentParticipant struct {
	AgentID string
	Role    AgentRole
	Status  ParticipantStatus
}

// ParticipantStatus represents the status of a participant in a conversation
type ParticipantStatus string

const (
	ParticipantStatusInvited ParticipantStatus = "invited"
	ParticipantStatusActive  ParticipantStatus = "active"
	ParticipantStatusIdle    ParticipantStatus = "idle"
	ParticipantStatusLeft    ParticipantStatus = "left"
)

// ThreadMessage represents a message in a conversation thread
type ThreadMessage struct {
	ID          string
	From        string
	Role        AgentRole
	Content     string
	MessageType pubsub.MessageType
	Timestamp   time.Time
	ReplyTo     string
	TokenCount  int
	Model       string
}

// ThreadState represents the state of a conversation thread
type ThreadState string

const (
	ThreadStateActive    ThreadState = "active"
	ThreadStateCompleted ThreadState = "completed"
	ThreadStateEscalated ThreadState = "escalated"
	ThreadStateClosed    ThreadState = "closed"
)

// Decision represents a decision made in a conversation
type Decision struct {
	ID          string
	Description string
	DecidedBy   []string
	Timestamp   time.Time
	Rationale   string
	Confidence  float64
}

// NewMcpServer creates a new MCP server instance
func NewMcpServer(
	ctx context.Context,
	node *p2p.Node,
	ps *pubsub.PubSub,
	hlog *logging.HypercoreLog,
	openaiAPIKey string,
) *McpServer {
	serverCtx, cancel := context.WithCancel(ctx)
	
	server := &McpServer{
		p2pNode:      node,
		pubsub:       ps,
		hlog:         hlog,
		openaiClient: openai.NewClient(openaiAPIKey),
		agents:       make(map[string]*GPTAgent),
		ctx:          serverCtx,
		cancel:       cancel,
		wsUpgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool { return true },
		},
		stats: &ServerStats{
			StartTime: time.Now(),
		},
	}
	
	return server
}

// Start initializes and starts the MCP server
func (s *McpServer) Start(port int) error {
	// Set up HTTP handlers
	mux := http.NewServeMux()
	
	// MCP WebSocket endpoint
	mux.HandleFunc("/mcp", s.handleMCPWebSocket)
	
	// REST API endpoints
	mux.HandleFunc("/api/agents", s.handleAgentsAPI)
	mux.HandleFunc("/api/conversations", s.handleConversationsAPI)
	mux.HandleFunc("/api/stats", s.handleStatsAPI)
	mux.HandleFunc("/health", s.handleHealthCheck)
	
	// Start HTTP server
	s.httpServer = &http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: mux,
	}
	
	go func() {
		if err := s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			fmt.Printf("❌ MCP HTTP server error: %v\n", err)
		}
	}()
	
	// Start message handlers
	go s.handleBzzzMessages()
	go s.handleHmmmMessages()
	
	// Start periodic tasks
	go s.periodicTasks()
	
	fmt.Printf("🚀 MCP Server started on port %d\n", port)
	return nil
}

// Stop gracefully shuts down the MCP server
func (s *McpServer) Stop() error {
	s.cancel()
	
	// Stop all agents
	s.agentsMutex.Lock()
	for _, agent := range s.agents {
		s.stopAgent(agent)
	}
	s.agentsMutex.Unlock()
	
	// Stop HTTP server
	if s.httpServer != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return s.httpServer.Shutdown(ctx)
	}
	
	return nil
}

// CreateGPTAgent creates a new GPT-4 agent
func (s *McpServer) CreateGPTAgent(config *AgentConfig) (*GPTAgent, error) {
	agent := &GPTAgent{
		ID:              config.ID,
		Role:            config.Role,
		Model:           config.Model,
		SystemPrompt:    config.SystemPrompt,
		Capabilities:    config.Capabilities,
		Specialization:  config.Specialization,
		MaxTasks:        config.MaxTasks,
		Status:          StatusIdle,
		CurrentTasks:    make(map[string]*AgentTask),
		Memory:          NewAgentMemory(),
		TokenUsage:      NewTokenUsage(),
		CostLimits:      config.CostLimits,
		NodeID:          s.p2pNode.ID().ShortString(),
		ActiveThreads:   make(map[string]*ConversationThread),
	}
	
	s.agentsMutex.Lock()
	s.agents[agent.ID] = agent
	s.agentsMutex.Unlock()
	
	// Announce agent to BZZZ network
	if err := s.announceAgent(agent); err != nil {
		return nil, fmt.Errorf("failed to announce agent: %w", err)
	}
	
	s.hlog.Append(logging.PeerJoined, map[string]interface{}{
		"agent_id":      agent.ID,
		"role":          string(agent.Role),
		"capabilities":  agent.Capabilities,
		"specialization": agent.Specialization,
	})
	
	fmt.Printf("✅ Created GPT-4 agent: %s (%s)\n", agent.ID, agent.Role)
	return agent, nil
}

// ProcessCollaborativeTask handles a task that requires multi-agent collaboration
func (s *McpServer) ProcessCollaborativeTask(
	task *AgentTask,
	requiredRoles []AgentRole,
) (*ConversationThread, error) {
	
	// Create conversation thread
	thread := &ConversationThread{
		ID:    fmt.Sprintf("task-%s-%d", task.Repository, task.Number),
		Topic: fmt.Sprintf("Collaborative Task: %s", task.Title),
		State: ThreadStateActive,
		SharedContext: map[string]interface{}{
			"task": task,
			"required_roles": requiredRoles,
		},
		CreatedAt:    time.Now(),
		LastActivity: time.Now(),
	}
	
	// Find and invite agents
	for _, role := range requiredRoles {
		agents := s.findAgentsByRole(role)
		if len(agents) == 0 {
			return nil, fmt.Errorf("no available agents for role: %s", role)
		}
		
		// Select best agent for this role
		selectedAgent := s.selectBestAgent(agents, task)
		
		thread.Participants = append(thread.Participants, AgentParticipant{
			AgentID: selectedAgent.ID,
			Role:    role,
			Status:  ParticipantStatusInvited,
		})
		
		// Add thread to agent
		selectedAgent.mutex.Lock()
		selectedAgent.ActiveThreads[thread.ID] = thread
		selectedAgent.mutex.Unlock()
	}
	
	// Send initial collaboration request
	if err := s.initiateCollaboration(thread); err != nil {
		return nil, fmt.Errorf("failed to initiate collaboration: %w", err)
	}
	
	return thread, nil
}

// handleMCPWebSocket handles WebSocket connections for MCP protocol
func (s *McpServer) handleMCPWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.wsUpgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Printf("❌ WebSocket upgrade failed: %v\n", err)
		return
	}
	defer conn.Close()
	
	fmt.Printf("📡 MCP WebSocket connection established\n")
	
	// Handle MCP protocol messages
	for {
		var message map[string]interface{}
		if err := conn.ReadJSON(&message); err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				fmt.Printf("❌ WebSocket error: %v\n", err)
			}
			break
		}
		
		// Process MCP message
		response, err := s.processMCPMessage(message)
		if err != nil {
			fmt.Printf("❌ MCP message processing error: %v\n", err)
			response = map[string]interface{}{
				"error": err.Error(),
			}
		}
		
		if err := conn.WriteJSON(response); err != nil {
			fmt.Printf("❌ WebSocket write error: %v\n", err)
			break
		}
	}
}

// processMCPMessage processes incoming MCP protocol messages
func (s *McpServer) processMCPMessage(message map[string]interface{}) (map[string]interface{}, error) {
	method, ok := message["method"].(string)
	if !ok {
		return nil, fmt.Errorf("missing or invalid method")
	}
	
	params, _ := message["params"].(map[string]interface{})
	
	switch method {
	case "tools/list":
		return s.listTools(), nil
	case "tools/call":
		return s.callTool(params)
	case "resources/list":
		return s.listResources(), nil
	case "resources/read":
		return s.readResource(params)
	default:
		return nil, fmt.Errorf("unknown method: %s", method)
	}
}

// callTool handles tool execution requests
func (s *McpServer) callTool(params map[string]interface{}) (map[string]interface{}, error) {
	toolName, ok := params["name"].(string)
	if !ok {
		return nil, fmt.Errorf("missing tool name")
	}
	
	args, _ := params["arguments"].(map[string]interface{})
	
	switch toolName {
	case "bzzz_announce":
		return s.handleBzzzAnnounce(args)
	case "bzzz_lookup":
		return s.handleBzzzLookup(args)
	case "bzzz_get":
		return s.handleBzzzGet(args)
	case "bzzz_post":
		return s.handleBzzzPost(args)
	case "bzzz_thread":
		return s.handleBzzzThread(args)
	case "bzzz_subscribe":
		return s.handleBzzzSubscribe(args)
	default:
		return nil, fmt.Errorf("unknown tool: %s", toolName)
	}
}

// handleBzzzAnnounce implements the bzzz_announce tool
func (s *McpServer) handleBzzzAnnounce(args map[string]interface{}) (map[string]interface{}, error) {
	agentID, ok := args["agent_id"].(string)
	if !ok {
		return nil, fmt.Errorf("agent_id is required")
	}
	
	role, ok := args["role"].(string)
	if !ok {
		return nil, fmt.Errorf("role is required")
	}
	
	// Create announcement message
	announcement := map[string]interface{}{
		"agent_id":      agentID,
		"role":          role,
		"capabilities":  args["capabilities"],
		"specialization": args["specialization"],
		"max_tasks":     args["max_tasks"],
		"announced_at":  time.Now(),
		"node_id":       s.p2pNode.ID().ShortString(),
	}
	
	// Publish to BZZZ network
	if err := s.pubsub.PublishBzzzMessage(pubsub.CapabilityBcast, announcement); err != nil {
		return nil, fmt.Errorf("failed to announce: %w", err)
	}
	
	return map[string]interface{}{
		"success": true,
		"message": fmt.Sprintf("Agent %s (%s) announced to network", agentID, role),
	}, nil
}

// Additional tool handlers would be implemented here...

// Helper methods

// announceAgent announces an agent to the BZZZ network
func (s *McpServer) announceAgent(agent *GPTAgent) error {
	announcement := map[string]interface{}{
		"type":           "gpt_agent_announcement",
		"agent_id":       agent.ID,
		"role":           string(agent.Role),
		"capabilities":   agent.Capabilities,
		"specialization": agent.Specialization,
		"max_tasks":      agent.MaxTasks,
		"model":          agent.Model,
		"node_id":        agent.NodeID,
		"timestamp":      time.Now(),
	}
	
	return s.pubsub.PublishBzzzMessage(pubsub.CapabilityBcast, announcement)
}

// findAgentsByRole finds all agents with a specific role
func (s *McpServer) findAgentsByRole(role AgentRole) []*GPTAgent {
	s.agentsMutex.RLock()
	defer s.agentsMutex.RUnlock()
	
	var agents []*GPTAgent
	for _, agent := range s.agents {
		if agent.Role == role && agent.Status == StatusIdle {
			agents = append(agents, agent)
		}
	}
	
	return agents
}

// selectBestAgent selects the best agent for a task
func (s *McpServer) selectBestAgent(agents []*GPTAgent, task *AgentTask) *GPTAgent {
	if len(agents) == 0 {
		return nil
	}
	
	// Simple selection: least busy agent
	bestAgent := agents[0]
	for _, agent := range agents[1:] {
		if len(agent.CurrentTasks) < len(bestAgent.CurrentTasks) {
			bestAgent = agent
		}
	}
	
	return bestAgent
}

// Additional helper methods would be implemented here...

// AgentConfig holds configuration for creating a new agent
type AgentConfig struct {
	ID             string
	Role           AgentRole
	Model          string
	SystemPrompt   string
	Capabilities   []string
	Specialization string
	MaxTasks       int
	CostLimits     *CostLimits
}

// CostLimits defines spending limits for an agent
type CostLimits struct {
	DailyLimit   float64
	MonthlyLimit float64
	PerTaskLimit float64
}

// TokenUsage tracks token consumption
type TokenUsage struct {
	TotalTokens      int64
	PromptTokens     int64
	CompletionTokens int64
	TotalCost        float64
	mutex            sync.RWMutex
}

// NewTokenUsage creates a new token usage tracker
func NewTokenUsage() *TokenUsage {
	return &TokenUsage{}
}

// NewAgentMemory creates a new agent memory instance
func NewAgentMemory() *AgentMemory {
	return &AgentMemory{
		WorkingMemory:  make(map[string]interface{}),
		EpisodicMemory: make([]ConversationEpisode, 0),
		ThreadMemories: make(map[string]*ThreadMemory),
	}
}

// ThreadMemory represents memory for a specific conversation thread
type ThreadMemory struct {
	ThreadID    string
	Summary     string
	KeyPoints   []string
	Decisions   []Decision
	LastUpdated time.Time
}

// KnowledgeGraph represents semantic knowledge
type KnowledgeGraph struct {
	Concepts   map[string]*Concept
	Relations  map[string]*Relation
	mutex      sync.RWMutex
}

// Concept represents a knowledge concept
type Concept struct {
	ID          string
	Name        string
	Description string
	Category    string
	Confidence  float64
}

// Relation represents a relationship between concepts
type Relation struct {
	From       string
	To         string
	Type       string
	Strength   float64
	Evidence   []string
}