import { randomUUID } from 'node:crypto';
import { Worker } from 'node:worker_threads';
import type { BehaviorInput } from '../../../../sdk/behavior-plugin/index.js';
import type {
  TriggerInput,
  TriggerResetReason,
} from '../../../../sdk/trigger-plugin/index.js';
import type { RuntimeDiagnosticSeverity } from '../../../framework/plugins/index.js';
import type {
  RuntimeModuleThreadLoaderConfig,
} from '../../../framework/modules/runtime-module-loader.js';
import type { TriggerThreadModuleSpec } from './trigger-module-loader.js';

const WORKER_BOOTSTRAP_URL = new URL(
  '../../../framework/modules/runtime-module-worker-bootstrap.mjs',
  import.meta.url,
);
const WORKER_HOST_URL = new URL('./trigger-thread-worker.ts', import.meta.url).href;
const MAX_TRIGGER_WORKERS = 10;
const START_TIMEOUT_MS = 20_000;

export interface TriggerThreadPoolOptions {
  readonly modules: readonly TriggerThreadModuleSpec[];
  readonly loader: RuntimeModuleThreadLoaderConfig;
  readonly maxOldGenerationSizeMb?: number;
  readonly diagnosticsEnabled?: boolean;
  readonly diagnostic?: (
    name: string,
    attributes?: Readonly<Record<string, unknown>>,
    severity?: RuntimeDiagnosticSeverity,
  ) => void;
}

export interface TriggerThreadEvaluation {
  readonly epochId: number;
  readonly triggerId: string;
  readonly input: Readonly<TriggerInput>;
  readonly readInput: Readonly<BehaviorInput>;
  readonly resetReason?: TriggerResetReason;
}

export interface TriggerThreadPoolSnapshot {
  readonly state: 'new' | 'starting' | 'running' | 'stopping' | 'stopped';
  readonly workerCount: number;
  readonly busyWorkerCount: number;
  readonly queuedRequestCount: number;
  readonly inFlightRequestCount: number;
  readonly workerHeapUsedBytes: number;
  readonly peakWorkerHeapUsedBytes: number;
}

interface PendingRequest {
  readonly requestId: string;
  readonly evaluation: TriggerThreadEvaluation;
  readonly signal: AbortSignal;
  readonly resolve: (value: unknown) => void;
  readonly reject: (error: Error) => void;
  readonly onAbort: () => void;
  settled: boolean;
}

interface TriggerLane {
  readonly id: number;
  readonly modules: readonly TriggerThreadModuleSpec[];
  readonly queue: PendingRequest[];
  worker?: Worker;
  state: 'starting' | 'idle' | 'busy' | 'restarting' | 'stopped';
  current?: PendingRequest;
  readyResolve?: () => void;
  readyReject?: (error: Error) => void;
  heapUsedBytes?: number;
  generation: number;
  consecutiveWorkerFailures: number;
  restartTask?: Promise<void>;
}

interface WorkerMessageBase {
  readonly type: string;
  readonly poolId: string;
  readonly laneId: number;
  readonly requestId?: string;
  readonly heapUsedBytes?: number;
}

interface WorkerReadyMessage extends WorkerMessageBase {
  readonly type: 'ready';
  readonly triggerIds: readonly string[];
}

interface WorkerResultMessage extends WorkerMessageBase {
  readonly type: 'result';
  readonly requestId: string;
  readonly triggerId: string;
  readonly result: unknown;
  readonly outputBytes: number;
}

interface WorkerErrorMessage extends WorkerMessageBase {
  readonly type: 'request_error' | 'error';
  readonly triggerId?: string;
  readonly code: string;
  readonly message?: string;
}

function asError(value: unknown, fallback: string): Error {
  return value instanceof Error ? value : new Error(typeof value === 'string' ? value : fallback);
}

function abortError(signal: AbortSignal): Error {
  return asError(signal.reason, 'trigger_thread_request_aborted');
}

export class TriggerThreadPool {
  private readonly poolId = randomUUID();
  private readonly maxOldGenerationSizeMb: number;
  private readonly lanes: TriggerLane[];
  private readonly laneByTrigger = new Map<string, TriggerLane>();
  private state: TriggerThreadPoolSnapshot['state'] = 'new';
  private peakWorkerHeapUsedBytes = 0;

  constructor(private readonly options: TriggerThreadPoolOptions) {
    const ids = options.modules.map(({ triggerId }) => triggerId);
    if (new Set(ids).size !== ids.length) throw new Error('trigger_thread_module_duplicate');
    this.maxOldGenerationSizeMb = Math.max(16, Math.min(
      256,
      Math.floor(options.maxOldGenerationSizeMb ?? 64),
    ));
    const laneCount = Math.min(MAX_TRIGGER_WORKERS, options.modules.length);
    const groups = Array.from({ length: laneCount }, () => [] as TriggerThreadModuleSpec[]);
    options.modules.forEach((module, index) => groups[index % Math.max(1, laneCount)]!.push(module));
    this.lanes = groups.map((modules, index) => ({
      id: index + 1,
      modules,
      queue: [],
      state: 'stopped',
      generation: 0,
      consecutiveWorkerFailures: 0,
    }));
    for (const lane of this.lanes) {
      for (const module of lane.modules) this.laneByTrigger.set(module.triggerId, lane);
    }
  }

  async start(): Promise<void> {
    if (this.state === 'running') return;
    if (this.state === 'stopping' || this.state === 'stopped') {
      throw new Error('trigger_thread_pool_stopped');
    }
    this.state = 'starting';
    const starts = await Promise.allSettled(this.lanes.map((lane) => this.spawn(lane)));
    if (this.state !== 'starting') throw new Error('trigger_thread_pool_stopped');
    for (const [index, result] of starts.entries()) {
      if (result.status === 'fulfilled') continue;
      const lane = this.lanes[index]!;
      lane.state = 'stopped';
      const worker = lane.worker;
      lane.worker = undefined;
      await worker?.terminate().catch(() => -1);
      this.diagnostic('trigger.module.disabled', {
        trigger_ids: lane.modules.map(({ triggerId }) => triggerId),
        reason: 'worker_start_failed',
        error_code: asError(result.reason, 'trigger_thread_start_failed').message.split(':', 1)[0],
      }, 'error');
    }
    this.state = 'running';
    this.diagnostic('trigger.thread_pool.started', {
      worker_count: this.lanes.filter(({ state }) => state !== 'stopped').length,
      unavailable_trigger_ids: this.lanes
        .filter(({ state }) => state === 'stopped')
        .flatMap(({ modules }) => modules.map(({ triggerId }) => triggerId)),
      max_worker_count: MAX_TRIGGER_WORKERS,
      worker_max_old_generation_mb: this.maxOldGenerationSizeMb,
    }, 'info');
    this.pumpAll();
  }

  evaluate(evaluation: TriggerThreadEvaluation, signal: AbortSignal): Promise<unknown> {
    if (this.state !== 'running') return Promise.reject(new Error('trigger_thread_pool_not_running'));
    if (signal.aborted) return Promise.reject(abortError(signal));
    const lane = this.laneByTrigger.get(evaluation.triggerId);
    if (!lane || lane.state === 'stopped') {
      return Promise.reject(new Error(`trigger_thread_unavailable:${evaluation.triggerId}`));
    }
    return new Promise((resolve, reject) => {
      const request: PendingRequest = {
        requestId: randomUUID(),
        evaluation,
        signal,
        resolve,
        reject,
        settled: false,
        onAbort: () => this.abortRequest(lane, request),
      };
      signal.addEventListener('abort', request.onAbort, { once: true });
      const insertion = lane.queue.findIndex(({ evaluation: queued }) => (
        queued.epochId > evaluation.epochId
      ));
      if (insertion < 0) lane.queue.push(request);
      else lane.queue.splice(insertion, 0, request);
      this.pump(lane);
    });
  }

  async stop(): Promise<void> {
    if (this.state === 'stopped') return;
    this.state = 'stopping';
    for (const lane of this.lanes) {
      const error = new Error('trigger_thread_pool_stopped');
      if (lane.current) this.rejectRequest(lane.current, error);
      for (const request of [...lane.queue]) this.rejectRequest(request, error);
      lane.queue.length = 0;
      lane.current = undefined;
      lane.state = 'stopped';
      const worker = lane.worker;
      lane.worker = undefined;
      if (worker) await worker.terminate().catch(() => -1);
    }
    this.state = 'stopped';
    this.diagnostic('trigger.thread_pool.stopped', {
      worker_count: 0,
      peak_worker_heap_used_bytes: this.peakWorkerHeapUsedBytes,
    }, 'info');
  }

  snapshot(): TriggerThreadPoolSnapshot {
    const requests = this.lanes.flatMap((lane) => [
      ...(lane.current ? [lane.current] : []),
      ...lane.queue,
    ]);
    return {
      state: this.state,
      workerCount: this.lanes.filter(({ worker }) => !!worker).length,
      busyWorkerCount: this.lanes.filter(({ state }) => state === 'busy').length,
      queuedRequestCount: this.lanes.reduce((sum, lane) => sum + lane.queue.length, 0),
      inFlightRequestCount: requests.length,
      workerHeapUsedBytes: this.currentWorkerHeapUsedBytes(),
      peakWorkerHeapUsedBytes: this.peakWorkerHeapUsedBytes,
    };
  }

  private spawn(lane: TriggerLane): Promise<void> {
    lane.generation += 1;
    const generation = lane.generation;
    lane.state = 'starting';
    const worker = new Worker(WORKER_BOOTSTRAP_URL, {
      execArgv: [],
      resourceLimits: { maxOldGenerationSizeMb: this.maxOldGenerationSizeMb },
      stdout: true,
      stderr: true,
      workerData: {
        runtimeModuleBootstrap: {
          hostUrl: WORKER_HOST_URL,
          hooksUrl: this.options.loader.hooksUrl,
          sdkAliases: this.options.loader.sdkAliases,
          esmModuleRoots: [...new Set(lane.modules.map((module) => (
            new URL('.', module.sourceUrl).href
          )))],
        },
        poolId: this.poolId,
        laneId: lane.id,
        modules: lane.modules,
        diagnosticsEnabled: this.options.diagnosticsEnabled === true,
      },
    });
    lane.worker = worker;
    worker.stdout?.on('data', (chunk: Buffer) => this.diagnostic(
      'trigger.thread.stdout',
      { trigger_ids: lane.modules.map(({ triggerId }) => triggerId), lane_id: lane.id, bytes: chunk.length },
    ));
    worker.stderr?.on('data', (chunk: Buffer) => this.diagnostic(
      'trigger.thread.stderr',
      { trigger_ids: lane.modules.map(({ triggerId }) => triggerId), lane_id: lane.id, bytes: chunk.length },
    ));
    worker.on('message', (message: unknown) => {
      if (lane.worker === worker && lane.generation === generation) this.handleMessage(lane, message);
    });
    worker.once('error', (error) => {
      if (lane.worker === worker && lane.generation === generation) this.failLane(lane, error);
    });
    worker.once('exit', (code) => {
      if (lane.worker === worker && lane.generation === generation && lane.state !== 'stopped') {
        this.failLane(lane, new Error(`trigger_thread_worker_exited:${code}`));
      }
    });
    return new Promise<void>((resolve, reject) => {
      let timer!: ReturnType<typeof setTimeout>;
      const finishResolve = (): void => {
        clearTimeout(timer);
        lane.readyResolve = undefined;
        lane.readyReject = undefined;
        resolve();
      };
      const finishReject = (error: Error): void => {
        clearTimeout(timer);
        lane.readyResolve = undefined;
        lane.readyReject = undefined;
        reject(error);
      };
      lane.readyResolve = finishResolve;
      lane.readyReject = finishReject;
      timer = setTimeout(
        () => finishReject(new Error('trigger_thread_worker_start_timeout')),
        START_TIMEOUT_MS,
      );
      timer.unref?.();
    });
  }

  private handleMessage(lane: TriggerLane, value: unknown): void {
    if (!value || typeof value !== 'object' || Array.isArray(value)) {
      this.failLane(lane, new Error('trigger_thread_protocol_invalid'));
      return;
    }
    const message = value as WorkerReadyMessage | WorkerResultMessage | WorkerErrorMessage;
    if (message.poolId !== this.poolId || message.laneId !== lane.id) {
      this.failLane(lane, new Error('trigger_thread_protocol_identity_mismatch'));
      return;
    }
    if (typeof message.heapUsedBytes === 'number' && Number.isFinite(message.heapUsedBytes)) {
      lane.heapUsedBytes = Math.max(0, message.heapUsedBytes);
      this.peakWorkerHeapUsedBytes = Math.max(
        this.peakWorkerHeapUsedBytes,
        this.currentWorkerHeapUsedBytes(),
      );
    }
    if (message.type === 'ready') {
      if (
        lane.state !== 'starting'
        || message.triggerIds.length !== lane.modules.length
        || message.triggerIds.some((id, index) => id !== lane.modules[index]?.triggerId)
      ) {
        this.failLane(lane, new Error('trigger_thread_ready_invalid'));
        return;
      }
      lane.state = 'idle';
      lane.readyResolve?.();
      this.pump(lane);
      return;
    }
    const request = lane.current;
    if (!request || message.requestId !== request.requestId) {
      this.failLane(lane, new Error('trigger_thread_request_mismatch'));
      return;
    }
    lane.current = undefined;
    lane.state = 'idle';
    if (message.type === 'result') {
      if (message.triggerId !== request.evaluation.triggerId) {
        this.rejectRequest(request, new Error('trigger_thread_result_invalid'));
      } else {
        lane.consecutiveWorkerFailures = 0;
        this.resolveRequest(request, message.result);
      }
    } else {
      const detail = message.message?.trim().slice(0, 512);
      this.rejectRequest(request, new Error(detail ? `${message.code}: ${detail}` : message.code));
    }
    this.pump(lane);
  }

  private pumpAll(): void {
    for (const lane of this.lanes) this.pump(lane);
  }

  private pump(lane: TriggerLane): void {
    if (this.state !== 'running' || lane.state !== 'idle' || lane.current) return;
    while (lane.queue.length > 0) {
      const request = lane.queue.shift()!;
      if (request.settled || request.signal.aborted) {
        if (!request.settled) this.rejectRequest(request, abortError(request.signal));
        continue;
      }
      lane.current = request;
      lane.state = 'busy';
      try {
        lane.worker!.postMessage({
          type: 'evaluate',
          poolId: this.poolId,
          requestId: request.requestId,
          triggerId: request.evaluation.triggerId,
          input: request.evaluation.input,
          readInput: request.evaluation.readInput,
          ...(request.evaluation.resetReason === undefined
            ? {}
            : { resetReason: request.evaluation.resetReason }),
        });
      } catch (error) {
        this.rejectRequest(request, asError(error, 'trigger_thread_submit_failed'));
        this.restartLane(lane, 'submit_failed');
      }
      return;
    }
  }

  private abortRequest(lane: TriggerLane, request: PendingRequest): void {
    if (request.settled) return;
    const queued = lane.queue.indexOf(request);
    if (queued >= 0) lane.queue.splice(queued, 1);
    const active = lane.current === request;
    this.rejectRequest(request, abortError(request.signal));
    if (active) this.restartLane(lane, 'request_aborted');
  }

  private resolveRequest(request: PendingRequest, value: unknown): void {
    if (request.settled) return;
    request.settled = true;
    request.signal.removeEventListener('abort', request.onAbort);
    request.resolve(structuredClone(value));
  }

  private rejectRequest(request: PendingRequest, error: Error): void {
    if (request.settled) return;
    request.settled = true;
    request.signal.removeEventListener('abort', request.onAbort);
    request.reject(error);
  }

  private failLane(lane: TriggerLane, error: Error): void {
    lane.readyReject?.(error);
    if (lane.current) this.rejectRequest(lane.current, error);
    lane.consecutiveWorkerFailures += 1;
    this.diagnostic('trigger.thread.failed', {
      trigger_id: lane.current?.evaluation.triggerId ?? null,
      trigger_ids: lane.modules.map(({ triggerId }) => triggerId),
      lane_id: lane.id,
      error_code: error.message.split(':', 1)[0],
    }, 'error');
    if (this.state !== 'running') return;
    if (lane.consecutiveWorkerFailures >= 3) {
      const worker = lane.worker;
      lane.worker = undefined;
      lane.current = undefined;
      lane.state = 'stopped';
      void worker?.terminate().catch(() => -1);
      for (const request of [...lane.queue]) {
        this.rejectRequest(request, new Error('trigger_thread_disabled_after_worker_failures'));
      }
      lane.queue.length = 0;
      this.diagnostic('trigger.module.disabled', {
        trigger_ids: lane.modules.map(({ triggerId }) => triggerId),
        reason: 'worker_failures',
        consecutive_failures: lane.consecutiveWorkerFailures,
      }, 'error');
      return;
    }
    this.restartLane(lane, 'worker_error');
  }

  private restartLane(lane: TriggerLane, reason: string): void {
    if (lane.restartTask || this.state !== 'running') return;
    const worker = lane.worker;
    lane.worker = undefined;
    lane.current = undefined;
    lane.state = 'restarting';
    lane.restartTask = Promise.resolve(worker?.terminate()).catch(() => -1).then(async () => {
      if (this.state !== 'running') return;
      await this.spawn(lane);
      this.diagnostic('trigger.thread.restarted', {
        trigger_ids: lane.modules.map(({ triggerId }) => triggerId),
        lane_id: lane.id,
        reason,
      }, 'warn');
    }).catch((error) => {
      lane.state = 'stopped';
      for (const request of [...lane.queue]) {
        this.rejectRequest(request, asError(error, 'trigger_thread_restart_failed'));
      }
      lane.queue.length = 0;
      this.diagnostic('trigger.thread.replacement_failed', {
        trigger_ids: lane.modules.map(({ triggerId }) => triggerId),
        lane_id: lane.id,
        reason,
        error_code: asError(error, 'trigger_thread_restart_failed').message.split(':', 1)[0],
      }, 'error');
    }).finally(() => {
      lane.restartTask = undefined;
      this.pump(lane);
    });
  }

  private currentWorkerHeapUsedBytes(): number {
    return this.lanes.reduce((sum, lane) => sum + (lane.heapUsedBytes ?? 0), 0);
  }

  private diagnostic(
    name: string,
    attributes?: Readonly<Record<string, unknown>>,
    severity?: RuntimeDiagnosticSeverity,
  ): void {
    try {
      this.options.diagnostic?.(name, attributes, severity);
    } catch {
      // Diagnostics never alter Trigger execution.
    }
  }
}
