import { 
  Agent, 
  Task, 
  ConversationThread, 
  Message, 
  OrchestrationSession,
  AgentTemplate,
  OrchestrationStrategy,
  TaskResult,
  Artifact
} from './types.js';

export class AgentOrchestrator {
  private agents: Map<string, Agent> = new Map();
  private sessions: Map<string, OrchestrationSession> = new Map();
  private activeSession: OrchestrationSession | null = null;
  private messageQueue: Message[] = [];
  
  constructor() {
    this.initializeDefaultTemplates();
  }
  
  /**
   * Create a new orchestration session
   */
  createSession(name: string, goal: string): OrchestrationSession {
    const session: OrchestrationSession = {
      id: this.generateId(),
      name,
      goal,
      agents: [],
      tasks: [],
      conversations: [],
      status: 'planning',
      startedAt: new Date(),
    };
    
    this.sessions.set(session.id, session);
    this.activeSession = session;
    return session;
  }
  
  /**
   * Create an agent from a template
   */
  createAgent(template: AgentTemplate, customName?: string): Agent {
    const agent: Agent = {
      id: this.generateId(),
      name: customName || `${template.name}-${Date.now()}`,
      role: template.role,
      capabilities: [...template.capabilities],
      systemPrompt: template.systemPrompt,
      status: 'idle',
      memory: {
        context: [],
        findings: {},
        decisions: [],
        conversations: [],
      },
      createdAt: new Date(),
      lastActiveAt: new Date(),
    };
    
    this.agents.set(agent.id, agent);
    
    if (this.activeSession) {
      this.activeSession.agents.push(agent);
    }
    
    return agent;
  }
  
  /**
   * Send a message from one agent to another
   */
  async sendMessage(
    from: string, 
    to: string | string[], 
    content: string, 
    type: Message['type'] = 'query'
  ): Promise<void> {
    const message: Message = {
      id: this.generateId(),
      from,
      to,
      content,
      type,
      timestamp: new Date(),
    };
    
    this.messageQueue.push(message);
    
    // Find or create conversation thread
    const participants = Array.isArray(to) ? [from, ...to] : [from, to];
    let thread = this.findConversationThread(participants);
    
    if (!thread) {
      thread = this.createConversationThread(participants, content);
    }
    
    thread.messages.push(message);
    thread.lastMessageAt = new Date();
    
    // Update agent memories
    const fromAgent = this.agents.get(from);
    if (fromAgent) {
      fromAgent.memory.conversations.push(thread);
      fromAgent.lastActiveAt = new Date();
      fromAgent.status = 'waiting';
    }
    
    // Process message for recipients
    const recipients = Array.isArray(to) ? to : [to];
    for (const recipientId of recipients) {
      await this.processMessageForAgent(recipientId, message);
    }
  }
  
  /**
   * Process a message for a specific agent
   */
  private async processMessageForAgent(agentId: string, message: Message): Promise<void> {
    const agent = this.agents.get(agentId);
    if (!agent) return;
    
    agent.status = 'working';
    agent.lastActiveAt = new Date();
    
    // Add to agent's context
    agent.memory.context.push(`Message from ${message.from}: ${message.content}`);
    
    // Simulate agent processing (in real implementation, this would call an LLM)
    // For now, we'll create a framework for handling different message types
    switch (message.type) {
      case 'query':
        // Agent should analyze and respond
        await this.handleQuery(agent, message);
        break;
        
      case 'instruction':
        // Agent should execute the instruction
        await this.handleInstruction(agent, message);
        break;
        
      case 'result':
        // Agent should process the result
        await this.handleResult(agent, message);
        break;
    }
    
    agent.status = 'idle';
  }
  
  /**
   * Create a task and assign it to agents
   */
  createTask(
    description: string, 
    requiredCapabilities: string[], 
    dependencies: string[] = []
  ): Task {
    const task: Task = {
      id: this.generateId(),
      description,
      assignedTo: [],
      requiredCapabilities,
      status: 'pending',
      dependencies,
      subtasks: [],
      createdAt: new Date(),
    };
    
    // Find suitable agents
    const suitableAgents = this.findAgentsWithCapabilities(requiredCapabilities);
    if (suitableAgents.length > 0) {
      task.assignedTo = suitableAgents.map(a => a.id);
      task.status = 'assigned';
    }
    
    if (this.activeSession) {
      this.activeSession.tasks.push(task);
    }
    
    return task;
  }
  
  /**
   * Execute a task with assigned agents
   */
  async executeTask(taskId: string): Promise<TaskResult> {
    const task = this.findTask(taskId);
    if (!task) {
      throw new Error(`Task ${taskId} not found`);
    }
    
    if (task.dependencies.length > 0) {
      // Check if dependencies are completed
      const incompleteDeps = task.dependencies.filter(depId => {
        const dep = this.findTask(depId);
        return !dep || dep.status !== 'completed';
      });
      
      if (incompleteDeps.length > 0) {
        throw new Error(`Task has incomplete dependencies: ${incompleteDeps.join(', ')}`);
      }
    }
    
    task.status = 'in-progress';
    task.startedAt = new Date();
    
    // Create a conversation for task coordination
    const thread = this.createConversationThread(
      task.assignedTo,
      `Task: ${task.description}`
    );
    
    // Primary agent starts the task
    const primaryAgent = task.assignedTo[0];
    await this.sendMessage(
      'orchestrator',
      primaryAgent,
      `Please execute the following task: ${task.description}`,
      'instruction'
    );
    
    // Simulate task execution (in real implementation, agents would work)
    const result: TaskResult = {
      success: true,
      output: `Task "${task.description}" completed successfully`,
      artifacts: [],
      performance: {
        duration: Date.now() - task.startedAt.getTime(),
      },
    };
    
    task.status = 'completed';
    task.completedAt = new Date();
    task.results = result;
    
    return result;
  }
  
  /**
   * Apply an orchestration strategy
   */
  async applyStrategy(strategy: OrchestrationStrategy): Promise<void> {
    if (!this.activeSession) {
      throw new Error('No active session');
    }
    
    for (const step of strategy.steps) {
      // Create agents for the step if needed
      const stepAgents: Agent[] = [];
      for (const role of step.agents) {
        const agent = this.findAgentByRole(role as Agent['role']) || 
                     this.createAgentFromRole(role as Agent['role']);
        stepAgents.push(agent);
      }
      
      // Execute actions
      if (step.parallel) {
        // Execute all actions in parallel
        await Promise.all(
          step.actions.map(action => this.executeAction(action, stepAgents))
        );
      } else {
        // Execute actions sequentially
        for (const action of step.actions) {
          await this.executeAction(action, stepAgents);
        }
      }
    }
  }
  
  /**
   * Get conversation history between agents
   */
  getConversationHistory(agentIds: string[]): ConversationThread[] {
    const threads: ConversationThread[] = [];
    
    if (this.activeSession) {
      for (const thread of this.activeSession.conversations) {
        const hasAllParticipants = agentIds.every(id => 
          thread.participants.includes(id)
        );
        if (hasAllParticipants) {
          threads.push(thread);
        }
      }
    }
    
    return threads;
  }
  
  /**
   * Get agent's current state and memory
   */
  getAgentState(agentId: string): Agent | undefined {
    return this.agents.get(agentId);
  }
  
  /**
   * Complete the current session
   */
  completeSession(): void {
    if (!this.activeSession) return;
    
    this.activeSession.status = 'completed';
    this.activeSession.completedAt = new Date();
    
    // Generate session results
    this.activeSession.results = {
      goalAchieved: true,
      summary: this.generateSessionSummary(),
      artifacts: this.collectSessionArtifacts(),
      insights: this.extractSessionInsights(),
      nextSteps: this.suggestNextSteps(),
    };
    
    this.activeSession = null;
  }
  
  // Helper methods
  
  private findConversationThread(participants: string[]): ConversationThread | undefined {
    if (!this.activeSession) return undefined;
    
    return this.activeSession.conversations.find(thread => {
      const threadParticipants = new Set(thread.participants);
      const searchParticipants = new Set(participants);
      return threadParticipants.size === searchParticipants.size &&
             [...threadParticipants].every(p => searchParticipants.has(p));
    });
  }
  
  private createConversationThread(participants: string[], topic: string): ConversationThread {
    const thread: ConversationThread = {
      id: this.generateId(),
      participants,
      messages: [],
      status: 'active',
      topic,
      startedAt: new Date(),
      lastMessageAt: new Date(),
    };
    
    if (this.activeSession) {
      this.activeSession.conversations.push(thread);
    }
    
    return thread;
  }
  
  private findAgentsWithCapabilities(capabilities: string[]): Agent[] {
    return Array.from(this.agents.values()).filter(agent => 
      capabilities.every(cap => agent.capabilities.includes(cap))
    );
  }
  
  private findTask(taskId: string): Task | undefined {
    if (!this.activeSession) return undefined;
    
    const searchInTasks = (tasks: Task[]): Task | undefined => {
      for (const task of tasks) {
        if (task.id === taskId) return task;
        const subtask = searchInTasks(task.subtasks);
        if (subtask) return subtask;
      }
      return undefined;
    };
    
    return searchInTasks(this.activeSession.tasks);
  }
  
  private findAgentByRole(role: Agent['role']): Agent | undefined {
    return Array.from(this.agents.values()).find(agent => agent.role === role);
  }
  
  private createAgentFromRole(role: Agent['role']): Agent {
    const template = this.getTemplateForRole(role);
    return this.createAgent(template);
  }
  
  private async executeAction(action: string, agents: Agent[]): Promise<void> {
    // This would be implemented based on specific action types
    console.log(`Executing action: ${action} with agents:`, agents.map(a => a.name));
  }
  
  private async handleQuery(agent: Agent, message: Message): Promise<void> {
    // Simulate agent processing a query
    const response = `${agent.name} analyzing: "${message.content}"`;
    await this.sendMessage(agent.id, message.from, response, 'response');
  }
  
  private async handleInstruction(agent: Agent, message: Message): Promise<void> {
    // Simulate agent executing an instruction
    agent.memory.decisions.push({
      id: this.generateId(),
      description: `Execute: ${message.content}`,
      reasoning: 'Instructed by orchestrator',
      timestamp: new Date(),
      confidence: 0.9,
    });
  }
  
  private async handleResult(agent: Agent, message: Message): Promise<void> {
    // Store result in agent's findings
    agent.memory.findings[`result-${Date.now()}`] = message.content;
  }
  
  private generateSessionSummary(): string {
    if (!this.activeSession) return '';
    
    const completedTasks = this.activeSession.tasks.filter(t => t.status === 'completed');
    const totalMessages = this.activeSession.conversations.reduce(
      (sum, thread) => sum + thread.messages.length, 0
    );
    
    return `Session "${this.activeSession.name}" completed. ` +
           `Goal: ${this.activeSession.goal}. ` +
           `Completed ${completedTasks.length}/${this.activeSession.tasks.length} tasks. ` +
           `${this.activeSession.agents.length} agents exchanged ${totalMessages} messages.`;
  }
  
  private collectSessionArtifacts(): Artifact[] {
    const artifacts: Artifact[] = [];
    
    if (this.activeSession) {
      for (const task of this.activeSession.tasks) {
        if (task.results?.artifacts) {
          artifacts.push(...task.results.artifacts);
        }
      }
    }
    
    return artifacts;
  }
  
  private extractSessionInsights(): string[] {
    const insights: string[] = [];
    
    // Extract insights from agent decisions and findings
    for (const agent of this.agents.values()) {
      const highConfidenceDecisions = agent.memory.decisions
        .filter(d => d.confidence > 0.8)
        .map(d => d.description);
      
      insights.push(...highConfidenceDecisions);
    }
    
    return insights;
  }
  
  private suggestNextSteps(): string[] {
    // Analyze incomplete tasks and suggest next steps
    const suggestions: string[] = [];
    
    if (this.activeSession) {
      const incompleteTasks = this.activeSession.tasks.filter(
        t => t.status !== 'completed'
      );
      
      for (const task of incompleteTasks) {
        suggestions.push(`Complete task: ${task.description}`);
      }
    }
    
    return suggestions;
  }
  
  private generateId(): string {
    return Date.now().toString(36) + Math.random().toString(36).substr(2);
  }
  
  private initializeDefaultTemplates(): void {
    // These would be the default agent templates
  }
  
  private getTemplateForRole(role: Agent['role']): AgentTemplate {
    const templates: Record<Agent['role'], AgentTemplate> = {
      primary: {
        name: 'Primary Coordinator',
        role: 'primary',
        systemPrompt: 'You are the primary coordinator. Manage the overall task and delegate to specialists.',
        capabilities: ['coordination', 'planning', 'delegation', 'summary'],
      },
      researcher: {
        name: 'Research Specialist',
        role: 'researcher',
        systemPrompt: 'You are a research specialist. Find information and analyze requirements.',
        capabilities: ['research', 'analysis', 'documentation', 'search'],
      },
      coder: {
        name: 'Code Developer',
        role: 'coder',
        systemPrompt: 'You are a code developer. Write clean, efficient, and tested code.',
        capabilities: ['coding', 'debugging', 'refactoring', 'optimization'],
      },
      reviewer: {
        name: 'Code Reviewer',
        role: 'reviewer',
        systemPrompt: 'You are a code reviewer. Review code for quality, security, and best practices.',
        capabilities: ['review', 'security-analysis', 'best-practices', 'feedback'],
      },
      tester: {
        name: 'Test Engineer',
        role: 'tester',
        systemPrompt: 'You are a test engineer. Write and execute comprehensive tests.',
        capabilities: ['testing', 'test-design', 'bug-finding', 'coverage-analysis'],
      },
      documenter: {
        name: 'Documentation Writer',
        role: 'documenter',
        systemPrompt: 'You are a documentation writer. Create clear and comprehensive documentation.',
        capabilities: ['documentation', 'tutorials', 'api-docs', 'examples'],
      },
      custom: {
        name: 'Custom Agent',
        role: 'custom',
        systemPrompt: 'You are a specialized agent with custom capabilities.',
        capabilities: ['custom'],
      },
    };
    
    return templates[role];
  }
}