import { Inject, Logger, WSController, OnWSConnection, OnWSMessage } from '@midwayjs/core';
import { ILogger } from '@midwayjs/logger';
import { Context } from '@midwayjs/socketio';
import { CrewOrchestratorService } from '../service/crew_orchestrator.js';
import { escalateResolvers } from '../tools/builtin/escalate.js';
import { InjectEntityModel } from '@midwayjs/typeorm';
import { Repository } from 'typeorm';
import { NetaClawCrewRunEntity } from '../entity/crew_run.js';
import type { CrewCallbacks } from '../service/crew_types.js';

/**
 * Crew WebSocket 网关
 * 命名空间: /crew
 */
@WSController('/crew')
export class CrewGateway {
  @Inject()
  ctx: Context;

  @Logger()
  logger: ILogger;

  @Inject()
  orchestrator: CrewOrchestratorService;

  @InjectEntityModel(NetaClawCrewRunEntity)
  crewRunRepo: Repository<NetaClawCrewRunEntity>;

  @OnWSConnection()
  async onConnection() {
    this.logger.info('[Crew WS] 客户端连接:', this.ctx.id);
  }

  @OnWSMessage('crew:trigger')
  async onTrigger(data: { crewId: number; triggerInput?: string }) {
    const { crewId, triggerInput = '' } = data;
    const callbacks = this.createCallbacks();
    try {
      const runId = await this.orchestrator.start(crewId, 'manual', triggerInput, callbacks);
      this.ctx.emit('crew:run:status', { runId, status: 'running' });
    } catch (err: any) {
      this.ctx.emit('crew:run:status', { runId: 0, status: 'failed', error: err.message });
    }
  }

  @OnWSMessage('crew:control')
  async onControl(data: { runId: number; action: string; userMessage?: string }) {
    const { runId, action, userMessage } = data;
    switch (action) {
      case 'stop':
        await this.orchestrator.stop(runId);
        this.ctx.emit('crew:run:status', { runId, status: 'failed' });
        break;
      case 'resume': {
        const resolver = escalateResolvers.get(runId);
        if (resolver) {
          resolver(userMessage || '继续执行');
          escalateResolvers.delete(runId);
          await this.crewRunRepo.update(runId, { status: 'running' } as any);
          this.ctx.emit('crew:run:status', { runId, status: 'running' });
        }
        break;
      }
      case 'pause':
        await this.orchestrator.pause(runId);
        this.ctx.emit('crew:run:status', { runId, status: 'paused' });
        break;
      case 'retry': {
        const run = await this.crewRunRepo.findOne({ where: { id: runId } });
        if (run) {
          const callbacks = this.createCallbacks();
          const newRunId = await this.orchestrator.start(run.crewId, 'manual', run.triggerInput || '', callbacks);
          this.ctx.emit('crew:run:status', { runId: newRunId, status: 'running' });
        }
        break;
      }
    }
  }

  private createCallbacks(): CrewCallbacks {
    const ctx = this.ctx;
    return {
      onLog: (runId, taskId, agentName, level, message) => {
        ctx.emit('crew:log', { runId, taskId, agentName, level, message, timestamp: new Date().toISOString() });
      },
      onTaskStatus: (runId, taskId, agentName, status, result, error) => {
        ctx.emit('crew:task:status', { runId, taskId, agentName, status, result, error });
      },
      onRunStatus: (runId, status, progress) => {
        ctx.emit('crew:run:status', { runId, status, progress });
      },
      onEscalation: (runId, taskId, reason, error) => {
        ctx.emit('crew:escalation', { runId, taskId, reason, error });
      },
    };
  }
}
