// src/worker-manager.ts
import { Worker } from './worker.js';
import { Repository } from './db/repository.js';
import { ConfigManager } from './config-manager.js';
import { type AppConfig, type WorkerDef } from './config.js';
import { logger } from './logger.js';
import type { McpTokenManager } from './mcp/token-manager.js';
import type { WorkerMetrics } from './metrics/worker-metrics.js';
import type { SkillCatalog } from './engine/skills.js';
import type { PushService } from './push-service.js';

export class WorkerManager {
  private workers: Worker[] = [];
  /**
   * Workers removed by a differential rebuild that still have an in-flight
   * job. They have stopped polling (claim no new work) but their detached
   * agent loop is allowed to finish — we deliberately do NOT requeue the job,
   * which is what previously caused the same job to run twice. Pruned on the
   * next rebuild and drained on shutdown.
   */
  private retiring: Worker[] = [];
  /** Config the current `workers` were built from — used to diff worker defs. */
  private lastConfig: AppConfig | null = null;
  private repo: Repository;
  private configManager: ConfigManager;
  private rebuildPromise: Promise<any> | null = null;
  private mcpTokenManager: McpTokenManager | null = null;
  /**
   * Phase 3b: Prometheus metrics handle. Wired post-construction by
   * worker-bootstrap once createCoreServer has built the registry.
   * Fanned out to every Worker on construct + on rebuild.
   */
  private workerMetrics: WorkerMetrics | null = null;
  private skillCatalog: SkillCatalog | null = null;
  private pushService: PushService | null = null;

  constructor(repo: Repository, configManager: ConfigManager) {
    this.repo = repo;
    this.configManager = configManager;
    this.configManager.onConfigChanged((config) => {
      this.rebuild(config).catch(err =>
        logger.error(`WorkerManager rebuild failed: ${err}`)
      );
    });
  }

  start(): void {
    const config = this.configManager.getConfig();
    this.workers = this.createWorkers(config);
    this.lastConfig = config;
    for (const w of this.workers) w.start();
    logger.info(`WorkerManager: started ${this.workers.length} workers`);
  }

  /**
   * Full shutdown of every worker (active + retiring). The global
   * `requeueRunningJobs()` is only valid here, on true process shutdown:
   * the process is exiting so the in-flight agent loops die with it, and
   * requeuing lets the next process pick the jobs up. It must NEVER be used
   * on a config rebuild (see doRebuild) — there the old loops keep running,
   * so requeuing would let the same job execute twice.
   */
  async stop(timeoutMs = 30000): Promise<{ drained: string[]; requeued: string[] }> {
    const drained: string[] = [];
    const requeued: string[] = [];

    const all = [...this.workers, ...this.retiring];

    // Stop polling for all workers
    for (const w of all) w.stop();

    // Wait for completion with timeout
    const results = await Promise.all(
      all.map(w => w.waitForCompletion(timeoutMs))
    );

    // Track drained (completed normally) and requeue (timed out)
    for (let i = 0; i < all.length; i++) {
      const w = all[i]!;
      if (results[i]) {
        drained.push(w.id);
      } else {
        // Timed out — requeue running jobs for this worker
        // Note: Repository doesn't have a workerId filter for running jobs,
        // so we requeue all running jobs. This is acceptable because the
        // process is shutting down and all worker loops are being torn down.
        requeued.push(w.id);
      }
    }

    // If any workers timed out, requeue all running jobs (shutdown only)
    if (requeued.length > 0) {
      this.repo.requeueRunningJobs();
      logger.info(`WorkerManager: requeued running jobs for timed-out workers: ${requeued.join(', ')}`);
    }

    this.workers = [];
    this.retiring = [];
    return { drained, requeued };
  }

  async rebuild(newConfig: AppConfig): Promise<{ drained: string[]; requeued: string[] }> {
    // Serialize rebuilds to prevent re-entrancy issues
    if (this.rebuildPromise) {
      await this.rebuildPromise;
    }
    this.rebuildPromise = this.doRebuild(newConfig);
    try {
      return await this.rebuildPromise;
    } finally {
      this.rebuildPromise = null;
    }
  }

  /**
   * Differential rebuild on config change.
   *
   * Workers whose own def is unchanged are KEPT running (config hot-swapped
   * in place) — so a change to, say, a tool size limit never restarts a busy
   * worker. Workers whose def changed (or were removed/disabled) are retired:
   * they stop polling but their in-flight agent loop finishes naturally and
   * the job is NOT requeued. The job stays `running` + issue-locked, so a
   * freshly started worker with the same id can't re-claim it.
   *
   * This closes the bug where any config edit during a long-running job
   * caused the job to be requeued under a still-running loop and executed a
   * second time concurrently.
   */
  private async doRebuild(newConfig: AppConfig): Promise<{ drained: string[]; requeued: string[] }> {
    logger.info('WorkerManager: rebuilding workers for new config');

    const oldSigs = this.defSignatures(this.lastConfig ?? this.configManager.getConfig());
    const newSigs = this.defSignatures(newConfig);

    const kept: Worker[] = [];
    const retired: Worker[] = [];
    for (const w of this.workers) {
      const newSig = newSigs.get(w.id);
      if (newSig !== undefined && newSig === oldSigs.get(w.id)) {
        w.updateConfig(newConfig); // def unchanged → keep alive, refresh global config
        kept.push(w);
      } else {
        retired.push(w); // removed, disabled, or def changed
      }
    }

    // Retire: stop polling only. Do NOT requeue — the detached loop finishes
    // on its own and the job remains running + locked, so no double execution.
    for (const w of retired) w.stop();
    this.retiring.push(...retired);

    // Create fresh workers for ids not already kept.
    const keptIds = new Set(kept.map(w => w.id));
    const fresh = newConfig.provider.workers
      .filter(def => def.enabled !== false && !keptIds.has(def.id))
      .map(def => this.createWorker(def, newConfig));
    for (const w of fresh) w.start();

    this.workers = [...kept, ...fresh];
    this.lastConfig = newConfig;

    const drained = this.pruneRetiring();
    if (this.retiring.length > 0) {
      logger.info(
        `WorkerManager: ${this.retiring.length} retired worker(s) draining in-flight jobs, no requeue: ${this.retiring.map(w => w.id).join(', ')}`
      );
    }
    logger.info(
      `WorkerManager: rebuilt kept=${kept.length} fresh=${fresh.length} retired=${retired.length} retiring=${this.retiring.length}`
    );
    return { drained, requeued: [] };
  }

  /** Stable per-id signature of each enabled worker def, for rebuild diffing. */
  private defSignatures(config: AppConfig): Map<string, string> {
    const sigs = new Map<string, string>();
    for (const def of config.provider.workers) {
      if (def.enabled === false) continue;
      sigs.set(def.id, JSON.stringify(def));
    }
    return sigs;
  }

  /** Drop retired workers whose in-flight job has finished; returns their ids. */
  private pruneRetiring(): string[] {
    const done = this.retiring.filter(w => w.inflightCount === 0);
    this.retiring = this.retiring.filter(w => w.inflightCount > 0);
    return done.map(w => w.id);
  }

  getWorkers(): Worker[] {
    return this.workers;
  }

  setMcpDeps(deps: { tokenManager: McpTokenManager }): void {
    this.mcpTokenManager = deps.tokenManager;
    for (const w of this.workers) {
      w.setMcpTokenManager(deps.tokenManager);
    }
  }

  /**
   * Phase 3b: install (or remove) the Prometheus metrics handle on
   * every Worker. Called by worker-bootstrap after createCoreServer
   * builds the registry. Setting null disables emission.
   */
  setWorkerMetrics(metrics: WorkerMetrics | null): void {
    this.workerMetrics = metrics;
    for (const w of this.workers) {
      w.setWorkerMetrics(metrics);
    }
  }

  setSkillCatalog(catalog: SkillCatalog): void {
    this.skillCatalog = catalog;
    for (const w of this.workers) {
      w.setSkillCatalog(catalog);
    }
  }

  /**
   * V2 Web Push notification service. Null disables push hooks across
   * all workers (rebuild safe). Fanned out to every Worker on construct
   * + on rebuild.
   * Spec: docs/superpowers/specs/2026-05-28-browser-notifications-v2-webpush.md.
   */
  setPushService(service: PushService | null): void {
    this.pushService = service;
    for (const w of this.workers) {
      w.setPushService(service);
    }
  }

  private createWorkers(config: AppConfig): Worker[] {
    return config.provider.workers
      .filter(def => def.enabled !== false)
      .map(def => this.createWorker(def, config));
  }

  /** Build a single worker and wire up the optional collaborators. */
  private createWorker(def: WorkerDef, config: AppConfig): Worker {
    const w = new Worker(def.id, def.endpoint, def.model, this.repo, config);
    if (this.mcpTokenManager) w.setMcpTokenManager(this.mcpTokenManager);
    if (this.workerMetrics) w.setWorkerMetrics(this.workerMetrics);
    if (this.skillCatalog) w.setSkillCatalog(this.skillCatalog);
    if (this.pushService) w.setPushService(this.pushService);
    return w;
  }
}
