import { ChatOpenAI } from "@langchain/openai";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";

const llm = new ChatOpenAI({ temperature: 0.7 });

// A2A Message Types
const A2AMessageType = {
  REQUEST: 'request',
  RESPONSE: 'response',
  NOTIFICATION: 'notification',
  ERROR: 'error',
  HEARTBEAT: 'heartbeat'
};

// Message Priority
const MessagePriority = {
  LOW: 1,
  NORMAL: 5,
  HIGH: 10,
  CRITICAL: 20
};

// A2A Message Structure
class A2AMessage {
  constructor(id, type, sender, receiver, payload, options = {}) {
    this.id = id || `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
    this.type = type;
    this.sender = sender;
    this.receiver = receiver;
    this.payload = payload;
    this.priority = options.priority || MessagePriority.NORMAL;
    this.correlationId = options.correlationId || null;
    this.replyTo = options.replyTo || null;
    this.timestamp = Date.now();
    this.ttl = options.ttl || 300000; // 5 minutes default
    this.metadata = options.metadata || {};
  }

  // Create request message
  static request(sender, receiver, payload, options = {}) {
    return new A2AMessage(
      null,
      A2AMessageType.REQUEST,
      sender,
      receiver,
      payload,
      options
    );
  }

  // Create response message
  static response(sender, receiver, payload, correlationId, options = {}) {
    return new A2AMessage(
      null,
      A2AMessageType.RESPONSE,
      sender,
      receiver,
      payload,
      { ...options, correlationId }
    );
  }

  // Create notification
  static notification(sender, receiver, payload, options = {}) {
    return new A2AMessage(
      null,
      A2AMessageType.NOTIFICATION,
      sender,
      receiver,
      payload,
      options
    );
  }

  // Create error message
  static error(sender, receiver, error, correlationId, options = {}) {
    return new A2AMessage(
      null,
      A2AMessageType.ERROR,
      sender,
      receiver,
      { error: error.message || error, code: error.code || 'UNKNOWN' },
      options
    );
  }

  // Check if message is expired
  isExpired() {
    return Date.now() - this.timestamp > this.ttl;
  }

  // Serialize to JSON
  toJSON() {
    return {
      id: this.id,
      type: this.type,
      sender: this.sender,
      receiver: this.receiver,
      payload: this.payload,
      priority: this.priority,
      correlationId: this.correlationId,
      replyTo: this.replyTo,
      timestamp: this.timestamp,
      ttl: this.ttl,
      metadata: this.metadata
    };
  }

  // Deserialize from JSON
  static fromJSON(json) {
    return new A2AMessage(
      json.id,
      json.type,
      json.sender,
      json.receiver,
      json.payload,
      {
        priority: json.priority,
        correlationId: json.correlationId,
        replyTo: json.replyTo,
        ttl: json.ttl,
        metadata: json.metadata
      }
    );
  }
}

// A2A Agent Endpoint
class A2AEndpoint {
  constructor(agentId, agentName) {
    this.agentId = agentId;
    this.agentName = agentName;
    this.messageQueue = [];
    this.subscribers = new Map();
    this.handlers = new Map();
    this.connected = false;
  }

  // Register message handler
  on(messageType, handler) {
    this.handlers.set(messageType, handler);
  }

  // Subscribe to messages from specific agent
  subscribe(agentId, callback) {
    if (!this.subscribers.has(agentId)) {
      this.subscribers.set(agentId, []);
    }
    this.subscribers.get(agentId).push(callback);
  }

  // Send message to another agent
  async send(receiver, payload, options = {}) {
    const message = A2AMessage.request(this.agentId, receiver, payload, options);
    this.messageQueue.push(message);
    console.log(`[${this.agentName}] Sent ${message.type} to ${receiver}`);

    // Simulate async delivery
    return message;
  }

  // Handle incoming message
  async handleMessage(message) {
    console.log(`[${this.agentName}] Received ${message.type} from ${message.sender}`);

    // Check if expired
    if (message.isExpired()) {
      console.log(`[${this.agentName}] Message expired, ignoring`);
      return null;
    }

    // Find handler
    const handler = this.handlers.get(message.type);
    if (handler) {
      try {
        const result = await handler(message.payload, message);
        return result;
      } catch (error) {
        console.error(`[${this.agentName}] Handler error:`, error);
        return A2AMessage.error(this.agentId, message.sender, error, message.id);
      }
    }

    return null;
  }

  // Get message history
  getHistory(limit = 50) {
    return this.messageQueue.slice(-limit);
  }

  // Get connection status
  getStatus() {
    return {
      agentId: this.agentId,
      agentName: this.agentName,
      connected: this.connected,
      queuedMessages: this.messageQueue.length,
      handlers: Array.from(this.handlers.keys()),
      subscribers: Array.from(this.subscribers.keys())
    };
  }
}

// Usage
async function demoA2AProtocol() {
  // Create agent endpoints
  const agentA = new A2AEndpoint('agent-alpha', 'Alpha Agent');
  const agentB = new A2AEndpoint('agent-beta', 'Beta Agent');

  // Register handlers
  agentB.on(A2AMessageType.REQUEST, async (payload, message) => {
    console.log(`[Beta] Processing request: ${payload.task}`);

    // Simulate processing
    const result = { status: 'completed', output: `Processed: ${payload.task}` };

    // Send response
    return A2AMessage.response(
      agentB.agentId,
      message.sender,
      result,
      message.id
    );
  });

  // Subscribe to notifications
  agentA.subscribe(agentB.agentId, (notification) => {
    console.log(`[Alpha] Received notification: ${notification.event}`);
  });

  // Send request
  console.log('\n--- Sending Request ---');
  const request = await agentA.send(agentB.agentId, {
    task: 'analyze_data',
    data: { symbol: 'AAPL', period: '1Y' }
  }, { priority: MessagePriority.HIGH });

  // Handle the request at receiver
  const response = await agentB.handleMessage(request);

  if (response) {
    console.log('\n--- Response ---');
    console.log(response.toJSON());
  }

  // Show status
  console.log('\n--- Agent Status ---');
  console.log(agentA.getStatus());
}

demoA2AProtocol();

import { ChatOpenAI } from "@langchain/openai";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";

const llm = new ChatOpenAI({ temperature: 0.7 });

// Task Status
const TaskStatus = {
  PENDING: 'pending',
  ASSIGNED: 'assigned',
  IN_PROGRESS: 'in_progress',
  COMPLETED: 'completed',
  FAILED: 'failed',
  CANCELLED: 'cancelled'
};

// Task Priority
const TaskPriority = {
  LOW: 1,
  NORMAL: 5,
  HIGH: 10,
  URGENT: 20
};

// Task Definition
class Task {
  constructor(id, type, description, config = {}) {
    this.id = id || `task_${Date.now()}`;
    this.type = type;
    this.description = description;
    this.config = config;
    this.status = TaskStatus.PENDING;
    this.priority = config.priority || TaskPriority.NORMAL;
    this.assignee = null;
    this.requester = null;
    this.dependencies = config.dependencies || [];
    this.result = null;
    this.error = null;
    this.createdAt = Date.now();
    this.startedAt = null;
    this.completedAt = null;
    this.metadata = {};
  }

  assign(agentId) {
    this.assignee = agentId;
    this.status = TaskStatus.ASSIGNED;
  }

  start() {
    this.status = TaskStatus.IN_PROGRESS;
    this.startedAt = Date.now();
  }

  complete(result) {
    this.status = TaskStatus.COMPLETED;
    this.result = result;
    this.completedAt = Date.now();
  }

  fail(error) {
    this.status = TaskStatus.FAILED;
    this.error = error;
    this.completedAt = Date.now();
  }

  cancel() {
    this.status = TaskStatus.CANCELLED;
    this.completedAt = Date.now();
  }

  getDuration() {
    if (!this.startedAt || !this.completedAt) return null;
    return this.completedAt - this.startedAt;
  }
}

// Task Delegation Manager
class TaskDelegationManager {
  constructor(agentId) {
    this.agentId = agentId;
    this.tasks = new Map();
    this.agents = new Map(); // Available agents and their capabilities
    this.taskQueue = [];
  }

  // Register an agent with capabilities
  registerAgent(agentId, capabilities, load = 0) {
    this.agents.set(agentId, {
      id: agentId,
      capabilities,
      load,
      status: 'available',
      currentTasks: []
    });
    console.log(`[Delegation] Registered agent: ${agentId} with capabilities: ${capabilities.join(', ')}`);
  }

  // Create and delegate a task
  delegateTask(type, description, config = {}) {
    const task = new Task(null, type, description, config);
    task.requester = this.agentId;

    // Find best agent for the task
    const suitableAgents = this.findSuitableAgents(task);

    if (suitableAgents.length === 0) {
      console.log(`[Delegation] No suitable agent found for task: ${type}`);
      this.taskQueue.push(task);
      return { status: 'queued', taskId: task.id };
    }

    // Select agent (lowest load)
    const selectedAgent = suitableAgents.sort((a, b) => a.load - b.load)[0];

    // Assign task
    task.assign(selectedAgent.id);
    this.tasks.set(task.id, task);
    selectedAgent.currentTasks.push(task.id);
    selectedAgent.load++;

    console.log(`[Delegation] Delegated task ${task.id} to ${selectedAgent.id}`);
    return { status: 'assigned', taskId: task.id, assignee: selectedAgent.id };
  }

  // Find agents suitable for task
  findSuitableAgents(task) {
    return Array.from(this.agents.values())
      .filter(agent => {
        // Check capabilities
        if (task.config.requiredCapabilities) {
          return task.config.requiredCapabilities.every(cap =>
            agent.capabilities.includes(cap)
          );
        }
        return true;
      })
      .filter(agent => agent.status !== 'offline');
  }

  // Complete a task
  completeTask(taskId, result) {
    const task = this.tasks.get(taskId);
    if (!task) return;

    task.complete(result);

    // Update agent load
    const agent = this.agents.get(task.assignee);
    if (agent) {
      agent.load = Math.max(0, agent.load - 1);
      agent.currentTasks = agent.currentTasks.filter(id => id !== taskId);
    }

    // Check queued tasks
    this.processQueue();

    return task;
  }

  // Fail a task
  failTask(taskId, error) {
    const task = this.tasks.get(taskId);
    if (!task) return;

    task.fail(error);

    // Update agent load
    const agent = this.agents.get(task.assignee);
    if (agent) {
      agent.load = Math.max(0, agent.load - 1);
    }

    return task;
  }

  // Process queued tasks
  processQueue() {
    if (this.taskQueue.length === 0) return;

    const queued = [...this.taskQueue];
    this.taskQueue = [];

    for (const task of queued) {
      const result = this.delegateTask(task.type, task.description, task.config);
      if (result.status !== 'queued') {
        console.log(`[Delegation] Processed queued task: ${task.id}`);
      }
    }
  }

  // Get task status
  getTaskStatus(taskId) {
    const task = this.tasks.get(taskId);
    if (!task) return null;

    return {
      id: task.id,
      type: task.type,
      status: task.status,
      assignee: task.assignee,
      duration: task.getDuration(),
      result: task.result
    };
  }

  // Get agent workload
  getAgentWorkload() {
    const workload = {};
    for (const [id, agent] of this.agents.entries()) {
      workload[id] = {
        load: agent.load,
        currentTasks: agent.currentTasks.length,
        status: agent.status
      };
    }
    return workload;
  }
}

// Usage
function demoTaskDelegation() {
  const coordinator = new TaskDelegationManager('coordinator');

  // Register specialized agents
  coordinator.registerAgent('researcher', ['research', 'analysis', 'information_gathering']);
  coordinator.registerAgent('writer', ['writing', 'summarization', 'editing']);
  coordinator.registerAgent('coder', ['coding', 'debugging', 'testing']);

  console.log('\n--- Delegating Tasks ---');

  // Delegate research task
  const result1 = coordinator.delegateTask(
    'research',
    'Research latest AI trends',
    { priority: TaskPriority.HIGH, requiredCapabilities: ['research'] }
  );
  console.log('Research task:', result1);

  // Delegate writing task
  const result2 = coordinator.delegateTask(
    'writing',
    'Write a blog post about AI',
    { priority: TaskPriority.NORMAL, requiredCapabilities: ['writing', 'summarization'] }
  );
  console.log('Writing task:', result2);

  // Delegate coding task
  const result3 = coordinator.delegateTask(
    'coding',
    'Implement a new feature',
    { priority: TaskPriority.URGENT, requiredCapabilities: ['coding'] }
  );
  console.log('Coding task:', result3);

  // Simulate task completion
  console.log('\n--- Completing Tasks ---');
  coordinator.completeTask(result1.taskId, { report: 'AI trends report ready' });
  coordinator.completeTask(result2.taskId, { article: 'Blog post published' });

  // Check workload
  console.log('\n--- Agent Workload ---');
  console.log(coordinator.getAgentWorkload());
}

demoTaskDelegation();

import { ChatOpenAI } from "@langchain/openai";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";

const llm = new ChatOpenAI({ temperature: 0.7 });

// Aggregation Strategy
const AggregationStrategy = {
  FIRST: 'first',           // Return first successful result
  ALL: 'all',               // Aggregate all results
  BEST: 'best',             // Select best result based on score
  WEIGHTED: 'weighted',    // Weighted average
  CONSENSUS: 'consensus'    // Find consensus among results
};

// Result Item
class AgentResult {
  constructor(agentId, result, metadata = {}) {
    this.agentId = agentId;
    this.result = result;
    this.metadata = metadata;
    this.timestamp = Date.now();
    this.score = metadata.score || 0;
    this.error = metadata.error || null;
    this.success = !metadata.error;
  }
}

// Aggregation Engine
class ResultAggregator {
  constructor(strategy = AggregationStrategy.ALL) {
    this.strategy = strategy;
    this.results = [];
    this.listeners = [];
  }

  // Add result from an agent
  addResult(agentId, result, metadata = {}) {
    const agentResult = new AgentResult(agentId, result, metadata);
    this.results.push(agentResult);
    console.log(`[Aggregator] Received result from ${agentId}:`, result);

    // Notify listeners
    this.notifyListeners(agentResult);

    return agentResult;
  }

  // Add error from an agent
  addError(agentId, error) {
    return this.addResult(agentId, null, { error: error.message || error });
  }

  // Aggregate results based on strategy
  aggregate() {
    const successfulResults = this.results.filter(r => r.success);

    switch (this.strategy) {
      case AggregationStrategy.FIRST:
        return this.aggregateFirst(successfulResults);

      case AggregationStrategy.ALL:
        return this.aggregateAll(successfulResults);

      case AggregationStrategy.BEST:
        return this.aggregateBest(successfulResults);

      case AggregationStrategy.WEIGHTED:
        return this.aggregateWeighted(successfulResults);

      case AggregationStrategy.CONSENSUS:
        return this.aggregateConsensus(successfulResults);

      default:
        return this.aggregateAll(successfulResults);
    }
  }

  // Return first successful result
  aggregateFirst(results) {
    if (results.length === 0) return { error: 'No successful results' };
    return {
      strategy: 'first',
      result: results[0].result,
      agentId: results[0].agentId
    };
  }

  // Return all results
  aggregateAll(results) {
    return {
      strategy: 'all',
      count: results.length,
      results: results.map(r => ({
        agentId: r.agentId,
        result: r.result,
        score: r.score
      }))
    };
  }

  // Return best result by score
  aggregateBest(results) {
    if (results.length === 0) return { error: 'No successful results' };

    const best = results.reduce((a, b) => a.score > b.score ? a : b);
    return {
      strategy: 'best',
      result: best.result,
      agentId: best.agentId,
      score: best.score
    };
  }

  // Weighted average (for numeric results)
  aggregateWeighted(results) {
    if (results.length === 0) return { error: 'No successful results' };

    const weightedResults = results.filter(r => typeof r.result === 'number');
    if (weightedResults.length === 0) {
      return this.aggregateAll(results);
    }

    let totalWeight = 0;
    let weightedSum = 0;

    for (const r of weightedResults) {
      const weight = r.metadata.weight || 1;
      weightedSum += r.result * weight;
      totalWeight += weight;
    }

    return {
      strategy: 'weighted',
      result: totalWeight > 0 ? weightedSum / totalWeight : 0,
      totalWeight
    };
  }

  // Find consensus (using LLM)
  async aggregateConsensus(results) {
    if (results.length === 0) return { error: 'No results to aggregate' };
    if (results.length === 1) return { strategy: 'consensus', ...results[0] };

    // Use LLM to find consensus
    const prompt = ChatPromptTemplate.fromTemplate(
      `Analyze the following results from multiple agents and find the consensus or best answer:

Results:
{results}

Provide a consolidated answer and explain the reasoning:`
    );

    const chain = prompt.pipe(llm).pipe(new StringOutputParser());

    try {
      const consolidated = await chain.invoke({
        results: results.map(r => `- ${r.agentId}: ${JSON.stringify(r.result)}`).join('\n')
      });

      return {
        strategy: 'consensus',
        result: consolidated,
        sourceCount: results.length
      };
    } catch (error) {
      // Fallback to simple voting
      return this.aggregateBest(results);
    }
  }

  // Subscribe to new results
  subscribe(callback) {
    this.listeners.push(callback);
  }

  // Notify listeners
  notifyListeners(result) {
    this.listeners.forEach(cb => cb(result));
  }

  // Get statistics
  getStats() {
    const successful = this.results.filter(r => r.success);
    const failed = this.results.filter(r => !r.success);

    return {
      total: this.results.length,
      successful: successful.length,
      failed: failed.length,
      avgScore: successful.length > 0
        ? successful.reduce((sum, r) => sum + r.score, 0) / successful.length
        : 0
    };
  }

  // Clear results
  clear() {
    this.results = [];
  }
}

// Multi-Agent Query Executor
class MultiAgentExecutor {
  constructor() {
    this.agents = new Map();
    this.aggregator = new ResultAggregator(AggregationStrategy.ALL);
  }

  // Register an agent
  registerAgent(agentId, handler, capabilities = []) {
    this.agents.set(agentId, { handler, capabilities });
    console.log(`[Executor] Registered agent: ${agentId}`);
  }

  // Execute query across multiple agents
  async executeQuery(query, options = {}) {
    const {
      targetAgents = null, // null = all agents
      strategy = AggregationStrategy.ALL,
      timeout = 30000,
      minResults = 1
    } = options;

    // Set aggregation strategy
    this.aggregator.strategy = strategy;
    this.aggregator.clear();

    // Determine which agents to query
    const agentsToQuery = targetAgents
      ? Array.from(this.agents.entries()).filter(([id]) => targetAgents.includes(id))
      : Array.from(this.agents.entries());

    console.log(`[Executor] Querying ${agentsToQuery.length} agents...`);

    // Execute in parallel with timeout
    const promises = agentsToQuery.map(async ([agentId, agent]) => {
      try {
        const result = await this.executeWithTimeout(agent.handler, query, timeout);
        this.aggregator.addResult(agentId, result, { capabilities: agent.capabilities });
      } catch (error) {
        this.aggregator.addError(agentId, error);
      }
    });

    await Promise.all(promises);

    // Check minimum results
    const stats = this.aggregator.getStats();
    if (stats.successful < minResults) {
      return {
        success: false,
        error: `Insufficient results: only ${stats.successful} of ${minResults} required`
      };
    }

    // Aggregate results
    const aggregated = this.aggregator.aggregate();

    return {
      success: true,
      stats,
      ...aggregated
    };
  }

  // Execute with timeout
  async executeWithTimeout(handler, query, timeout) {
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        reject(new Error('Execution timeout'));
      }, timeout);

      Promise.resolve(handler(query))
        .then(result => {
          clearTimeout(timer);
          resolve(result);
        })
        .catch(err => {
          clearTimeout(timer);
          reject(err);
        });
    });
  }
}

// Usage
async function demoResultAggregation() {
  const executor = new MultiAgentExecutor();

  // Register different agents with different capabilities
  executor.registerAgent('researcher', async (query) => {
    // Simulate research agent
    return { type: 'research', findings: 'AI market growing at 25% CAGR' };
  }, ['research']);

  executor.registerAgent('analyst', async (query) => {
    // Simulate analyst agent
    return { type: 'analysis', prediction: 'Tech sector will outperform' };
  }, ['analysis']);

  executor.registerAgent('news', async (query) => {
    // Simulate news agent
    return { type: 'news', headline: 'Major AI partnership announced' };
  }, ['news']);

  // Execute query with different strategies
  console.log('\n--- Strategy: ALL ---');
  const result1 = await executor.executeQuery('What is the AI market outlook?', {
    strategy: AggregationStrategy.ALL
  });
  console.log(JSON.stringify(result1, null, 2));

  console.log('\n--- Strategy: BEST (with scores) ---');
  executor.aggregator.addResult('researcher', { data: 'result1' }, { score: 0.9 });
  executor.aggregator.addResult('analyst', { data: 'result2' }, { score: 0.7 });
  executor.aggregator.addResult('news', { data: 'result3' }, { score: 0.8 });

  const result2 = executor.aggregate();
  console.log(JSON.stringify(result2, null, 2));

  console.log('\n--- Strategy: CONSENSUS ---');
  executor.aggregator.clear();
  executor.aggregator.addResult('researcher', { outlook: 'positive' }, { score: 0.9 });
  executor.aggregator.addResult('analyst', { outlook: 'positive' }, { score: 0.8 });
  executor.aggregator.addResult('news', { outlook: 'neutral' }, { score: 0.7 });

  const result3 = await executor.aggregator.aggregateConsensus(
    executor.aggregator.results
  );
  console.log(JSON.stringify(result3, null, 2));
}

demoResultAggregation();

// Agent Card - Discovery Metadata
class AgentCard {
  constructor(agentId, name, description, capabilities, endpoints = {}) {
    this.agentId = agentId;
    this.name = name;
    this.description = description;
    this.capabilities = capabilities;
    this.endpoints = endpoints;
    this.version = '1.0';
    this.status = 'active';
    this.metadata = {
      createdAt: Date.now(),
      updatedAt: Date.now(),
      provider: 'default',
      authType: 'none'
    };
  }

  toJSON() {
    return {
      agentId: this.agentId,
      name: this.name,
      description: this.description,
      capabilities: this.capabilities,
      version: this.version,
      status: this.status,
      endpoints: this.endpoints,
      metadata: this.metadata
    };
  }

  static fromJSON(json) {
    const card = new AgentCard(
      json.agentId,
      json.name,
      json.description,
      json.capabilities,
      json.endpoints
    );
    card.version = json.version;
    card.status = json.status;
    card.metadata = json.metadata;
    return card;
  }
}

// Agent Registry - Discovery Service
class AgentRegistry {
  constructor() {
    this.agents = new Map();
    this.subscriptions = new Map();
  }

  // Register an agent
  register(agentCard) {
    this.agents.set(agentCard.agentId, agentCard);
    console.log(`[Registry] Registered agent: ${agentCard.name} (${agentCard.agentId})`);
    this.notifySubscribers('register', agentCard);
    return agentCard;
  }

  // Unregister an agent
  unregister(agentId) {
    const agent = this.agents.get(agentId);
    if (agent) {
      this.agents.delete(agentId);
      console.log(`[Registry] Unregistered agent: ${agentId}`);
      this.notifySubscribers('unregister', agent);
      return true;
    }
    return false;
  }

  // Update agent status
  updateStatus(agentId, status) {
    const agent = this.agents.get(agentId);
    if (agent) {
      agent.status = status;
      agent.metadata.updatedAt = Date.now();
      console.log(`[Registry] Updated status for ${agentId}: ${status}`);
      this.notifySubscribers('status', agent);
      return agent;
    }
    return null;
  }

  // Discover agents by capability
  discoverByCapability(capability) {
    return Array.from(this.agents.values())
      .filter(agent => agent.status === 'active')
      .filter(agent => agent.capabilities.includes(capability));
  }

  // Discover agents by multiple capabilities (all must match)
  discoverByCapabilities(requiredCapabilities) {
    return Array.from(this.agents.values())
      .filter(agent => agent.status === 'active')
      .filter(agent => requiredCapabilities.every(cap => agent.capabilities.includes(cap)));
  }

  // Get agent by ID
  getAgent(agentId) {
    return this.agents.get(agentId);
  }

  // Get all active agents
  getAllActive() {
    return Array.from(this.agents.values())
      .filter(agent => agent.status === 'active');
  }

  // Subscribe to registry changes
  subscribe(callback) {
    const subscriptionId = `sub_${Date.now()}`;
    this.subscriptions.set(subscriptionId, callback);
    return subscriptionId;
  }

  // Unsubscribe
  unsubscribe(subscriptionId) {
    return this.subscriptions.delete(subscriptionId);
  }

  // Notify subscribers
  notifySubscribers(event, agent) {
    this.subscriptions.forEach(callback => {
      try {
        callback(event, agent);
      } catch (error) {
        console.error('[Registry] Subscription error:', error);
      }
    });
  }

  // Get registry statistics
  getStats() {
    const agents = Array.from(this.agents.values());
    return {
      total: agents.length,
      active: agents.filter(a => a.status === 'active').length,
      inactive: agents.filter(a => a.status === 'inactive').length,
      byCapability: this.getCapabilityCounts()
    };
  }

  // Get capability distribution
  getCapabilityCounts() {
    const counts = {};
    this.agents.forEach(agent => {
      agent.capabilities.forEach(cap => {
        counts[cap] = (counts[cap] || 0) + 1;
      });
    });
    return counts;
  }
}

// A2A Server - Handles Agent Communication
class A2AServer {
  constructor(registry) {
    this.registry = registry;
    this.messageHandlers = new Map();
  }

  // Handle incoming A2A message
  async handleMessage(messageJson) {
    const message = A2AMessage.fromJSON(messageJson);

    console.log(`[A2AServer] Handling ${message.type} from ${message.sender}`);

    // Route based on message type
    switch (message.type) {
      case A2AMessageType.REQUEST:
        return await this.handleRequest(message);

      case A2AMessageType.NOTIFICATION:
        return await this.handleNotification(message);

      case 'discovery':
        return this.handleDiscovery(message);

      case 'heartbeat':
        return this.handleHeartbeat(message);

      default:
        console.warn(`[A2AServer] Unknown message type: ${message.type}`);
    }
  }

  // Handle discovery request
  handleDiscovery(message) {
    const { capability, agentId } = message.payload;

    let results;
    if (agentId) {
      const agent = this.registry.getAgent(agentId);
      results = agent ? [agent] : [];
    } else if (capability) {
      results = this.registry.discoverByCapability(capability);
    } else {
      results = this.registry.getAllActive();
    }

    return A2AMessage.response(
      'server',
      message.sender,
      { agents: results.map(a => a.toJSON()) },
      message.id
    );
  }

  // Handle agent request
  async handleRequest(message) {
    const { targetAgent, action, payload } = message.payload;

    const agent = this.registry.getAgent(targetAgent);
    if (!agent) {
      return A2AMessage.error(
        'server',
        message.sender,
        { code: 'AGENT_NOT_FOUND', message: `Agent ${targetAgent} not found` },
        message.id
      );
    }

    // Find handler for this agent
    const handler = this.messageHandlers.get(targetAgent);
    if (!handler) {
      return A2AMessage.error(
        'server',
        message.sender,
        { code: 'NO_HANDLER', message: 'Agent not available' },
        message.id
      );
    }

    try {
      const result = await handler(action, payload);
      return A2AMessage.response('server', message.sender, result, message.id);
    } catch (error) {
      return A2AMessage.error('server', message.sender, error, message.id);
    }
  }

  // Handle notification
  async handleNotification(message) {
    console.log(`[A2AServer] Notification: ${message.payload.event}`);
    // Process notification
    return null;
  }

  // Handle heartbeat
  handleHeartbeat(message) {
    const { agentId, status } = message.payload;
    if (status) {
      this.registry.updateStatus(agentId, status);
    }
    return A2AMessage.response('server', message.sender, { status: 'ok' }, message.id);
  }

  // Register message handler for agent
  onAgentMessage(agentId, handler) {
    this.messageHandlers.set(agentId, handler);
  }
}

// Usage
function demoAgentDiscovery() {
  const registry = new AgentRegistry();

  // Register agents
  const card1 = new AgentCard(
    'research-agent',
    'Research Agent',
    'Specialized in market research and analysis',
    ['research', 'analysis', 'data_gathering'],
    { apiEndpoint: '/api/agents/research-agent' }
  );

  const card2 = new AgentCard(
    'writer-agent',
    'Writer Agent',
    'Content creation and copywriting',
    ['writing', 'editing', 'translation'],
    { apiEndpoint: '/api/agents/writer-agent' }
  );

  const card3 = new AgentCard(
    'data-agent',
    'Data Agent',
    'Data processing and visualization',
    ['data_processing', 'analysis', 'visualization'],
    { apiEndpoint: '/api/agents/data-agent' }
  );

  registry.register(card1);
  registry.register(card2);
  registry.register(card3);

  // Subscribe to changes
  const subId = registry.subscribe((event, agent) => {
    console.log(`[Watch] ${event}: ${agent.name}`);
  });

  console.log('\n--- Discovery by Capability: research ---');
  const researchers = registry.discoverByCapability('research');
  console.log(researchers.map(a => a.name));

  console.log('\n--- Discovery: agents with analysis ---');
  const analysts = registry.discoverByCapabilities(['analysis']);
  console.log(analysts.map(a => a.name));

  console.log('\n--- Registry Statistics ---');
  console.log(registry.getStats());

  // Simulate A2A server
  console.log('\n--- A2A Server Discovery ---');
  const server = new A2AServer(registry);

  const discoveryMsg = A2AMessage.request('client', 'server', {
    capability: 'writing'
  });

  const response = server.handleDiscovery(discoveryMsg);
  console.log('Discovery response:', response.payload);

  // Cleanup
  registry.unsubscribe(subId);
}

demoAgentDiscovery();