/**
 * v5 InferenceTaskService — business logic over the durable task queue.
 *
 * The DB row (managed via IInferenceTaskRepository) is the source of truth
 * for *state*. This service owns the in-process *signaling* layer that
 * lets blocked HTTP handlers wake up promptly when a worker POSTs a
 * result, instead of polling the DB. The signals carry no state — they
 * just say "go re-read row X" — so a missed signal degrades to a slower
 * poll-then-resolve (still correct, just less responsive).
 *
 * Single-instance assumption: EventEmitter only fires on the mcpd that
 * holds the row's blocked handler. With multiple mcpd replicas, the
 * worker's POST might land on instance A while the caller's handler is
 * on instance B; B never sees the local emitter. Solution at scale is
 * pg LISTEN/NOTIFY (v6 swap) — no schema change needed; just replace the
 * EventEmitter wakeup with a NOTIFY emit.
 */
import { EventEmitter } from 'node:events';
import type { InferenceTask, InferenceTaskStatus } from '@prisma/client';
import type { IInferenceTaskRepository, CreateInferenceTaskInput } from '../repositories/inference-task.repository.js';
import { NotFoundError } from './mcp-server.service.js';

/**
 * One streaming chunk pushed back from a worker. Lives in-memory only —
 * we don't persist every SSE delta to the DB (too expensive for
 * high-frequency streams). The final assembled body lands on the row
 * via `complete()` when the worker emits its terminal `done:true`.
 */
export interface InferenceTaskChunk {
  data: string;
  done?: boolean;
}

/**
 * The wait handle returned to a blocked HTTP handler. `done` resolves
 * when the row reaches a terminal status (completed | error | cancelled);
 * `chunks` is an async iterator that yields streaming deltas as they
 * arrive (only meaningful when the underlying request asked for streaming).
 */
export interface InferenceTaskWaiter {
  /** The DB row's terminal state. Throws on `error`/`cancelled` with an explanatory message. */
  done: Promise<InferenceTask>;
  /** Async iterator of streaming chunks. Yields nothing for non-streaming tasks. */
  chunks: AsyncGenerator<InferenceTaskChunk>;
}

export interface IInferenceTaskService {
  /** Create a new pending task. Caller is expected to immediately attempt dispatch. */
  enqueue(input: CreateInferenceTaskInput): Promise<InferenceTask>;
  /** Wait for a task to reach a terminal state, with optional timeout (ms). */
  waitFor(taskId: string, timeoutMs: number): InferenceTaskWaiter;
  /** Conditional CAS claim — returns the row if `pending`, null if already claimed. */
  tryClaim(taskId: string, claimedBy: string): Promise<InferenceTask | null>;
  /** Worker reported first chunk — flips `claimed` → `running`. Idempotent. */
  markRunning(taskId: string): Promise<InferenceTask | null>;
  /** Worker pushed a streaming chunk. Wakes up the in-process waiter; no DB write. */
  pushChunk(taskId: string, chunk: InferenceTaskChunk): boolean;
  /**
   * Subscribe directly to a task's chunk stream as a callback. Returns
   * an unsubscribe function. Used by VirtualLlmService's legacy
   * PendingTaskRef.onChunk bridge — the high-level `waitFor` API is
   * better when you need an async iterator + done promise together.
   */
  subscribeChunksUnsafe(taskId: string, cb: (chunk: InferenceTaskChunk) => void): () => void;
  /** Worker POSTed the final result. Persists the body + flips to `completed`. */
  complete(taskId: string, responseBody: Record<string, unknown> | null): Promise<InferenceTask | null>;
  /** Worker (or mcpd) marked the task as failed. */
  fail(taskId: string, errorMessage: string): Promise<InferenceTask | null>;
  /** Caller (or GC) marked the task as cancelled. */
  cancel(taskId: string): Promise<InferenceTask | null>;
  /** Worker disconnected mid-task — revert claimed/running rows back to pending. */
  revertHeldBy(claimedBy: string): Promise<InferenceTask[]>;
  /** Pending rows for one or more pool keys, ordered FIFO. Used by drainPending on bind. */
  findPendingForPools(poolNames: string[], limit?: number): Promise<InferenceTask[]>;
  findById(id: string): Promise<InferenceTask | null>;
  list(filter: {
    ownerId?: string;
    status?: InferenceTaskStatus | InferenceTaskStatus[];
    poolName?: string;
    agentId?: string;
    limit?: number;
  }): Promise<InferenceTask[]>;
  /** GC sweep: pending too long → error; terminal too long → delete. */
  gcSweep(opts: { pendingTimeoutMs: number; terminalRetentionMs: number }): Promise<{ erroredOut: number; deleted: number }>;
}

const DEFAULT_PENDING_TIMEOUT_MS = 60 * 60 * 1000; // 1 h
const DEFAULT_TERMINAL_RETENTION_MS = 7 * 24 * 60 * 60 * 1000; // 7 d

export class InferenceTaskService implements IInferenceTaskService {
  /**
   * Wakeup channel per task. The result-handler (worker POSTs landing on
   * mcpd) emits a terminal status; the streaming push emits chunks.
   * Cleared once a waiter resolves to avoid leaks. EventEmitter listener
   * cap is bumped to a generous default — typical concurrent waiters per
   * task is 1, but `/stream` SSE consumers and the original HTTP handler
   * can both subscribe.
   */
  private readonly events = new EventEmitter();

  constructor(
    private readonly repo: IInferenceTaskRepository,
    private readonly clock: () => Date = () => new Date(),
  ) {
    this.events.setMaxListeners(50);
  }

  async enqueue(input: CreateInferenceTaskInput): Promise<InferenceTask> {
    return this.repo.create(input);
  }

  waitFor(taskId: string, timeoutMs: number): InferenceTaskWaiter {
    // Set up the chunk channel up-front so a worker pushing chunks before
    // the consumer subscribes doesn't drop them. We use a queue + wake
    // pattern (same shape as the v3 chat.service streaming bridge).
    const chunkQueue: InferenceTaskChunk[] = [];
    let chunkResolve: (() => void) | null = null;
    let finished = false;
    let finishError: Error | null = null;

    const onChunk = (chunk: InferenceTaskChunk): void => {
      chunkQueue.push(chunk);
      if (chunkResolve !== null) {
        const r = chunkResolve;
        chunkResolve = null;
        r();
      }
    };
    const onTerminal = (): void => {
      finished = true;
      if (chunkResolve !== null) {
        const r = chunkResolve;
        chunkResolve = null;
        r();
      }
    };

    this.events.on(`chunk:${taskId}`, onChunk);
    this.events.on(`terminal:${taskId}`, onTerminal);

    const cleanup = (): void => {
      this.events.off(`chunk:${taskId}`, onChunk);
      this.events.off(`terminal:${taskId}`, onTerminal);
    };

    // The terminal promise: poll once at start (in case the row is
    // already terminal — common for already-completed tasks the caller
    // is just fetching the result of), then race the timeout against
    // the wakeup signal.
    const done = (async (): Promise<InferenceTask> => {
      try {
        const initial = await this.repo.findById(taskId);
        if (initial === null) throw new NotFoundError(`InferenceTask not found: ${taskId}`);
        if (isTerminal(initial.status)) return ensureSuccess(initial);

        // Wait for terminal event OR timeout. If the wake happens before
        // the timer fires, we re-fetch and return; otherwise we error.
        const timer = new Promise<never>((_, reject) => {
          setTimeout(() => reject(new Error(`InferenceTask wait timed out after ${String(timeoutMs)}ms (task ${taskId})`)), timeoutMs);
        });
        const wake = new Promise<void>((resolve) => {
          this.events.once(`terminal:${taskId}`, () => resolve());
        });
        await Promise.race([timer, wake]);
        finished = true;
        const final = await this.repo.findById(taskId);
        if (final === null) throw new NotFoundError(`InferenceTask not found: ${taskId}`);
        return ensureSuccess(final);
      } catch (err) {
        finishError = err as Error;
        finished = true;
        throw err;
      } finally {
        cleanup();
      }
    })();

    const chunks = (async function* gen(): AsyncGenerator<InferenceTaskChunk> {
      while (true) {
        while (chunkQueue.length > 0) {
          const c = chunkQueue.shift()!;
          yield c;
          if (c.done === true) return;
        }
        if (finished) {
          if (finishError !== null) throw finishError;
          return;
        }
        await new Promise<void>((r) => { chunkResolve = r; });
      }
    })();

    return { done, chunks };
  }

  async tryClaim(taskId: string, claimedBy: string): Promise<InferenceTask | null> {
    return this.repo.tryClaim(taskId, claimedBy, this.clock());
  }

  async markRunning(taskId: string): Promise<InferenceTask | null> {
    return this.repo.markRunning(taskId, this.clock());
  }

  pushChunk(taskId: string, chunk: InferenceTaskChunk): boolean {
    // Streaming chunks are in-memory only. If the row no longer exists or
    // there are no listeners, the emit is a no-op — return false so the
    // result-route can decide whether to log a warning.
    return this.events.emit(`chunk:${taskId}`, chunk);
  }

  subscribeChunksUnsafe(taskId: string, cb: (chunk: InferenceTaskChunk) => void): () => void {
    this.events.on(`chunk:${taskId}`, cb);
    return (): void => {
      this.events.off(`chunk:${taskId}`, cb);
    };
  }

  async complete(taskId: string, responseBody: Record<string, unknown> | null): Promise<InferenceTask | null> {
    const updated = await this.repo.markCompleted(taskId, responseBody, this.clock());
    if (updated !== null) this.events.emit(`terminal:${taskId}`);
    return updated;
  }

  async fail(taskId: string, errorMessage: string): Promise<InferenceTask | null> {
    const updated = await this.repo.markError(taskId, errorMessage, this.clock());
    if (updated !== null) this.events.emit(`terminal:${taskId}`);
    return updated;
  }

  async cancel(taskId: string): Promise<InferenceTask | null> {
    const updated = await this.repo.markCancelled(taskId, this.clock());
    if (updated !== null) this.events.emit(`terminal:${taskId}`);
    return updated;
  }

  async revertHeldBy(claimedBy: string): Promise<InferenceTask[]> {
    const held = await this.repo.findHeldBy(claimedBy);
    const reverted: InferenceTask[] = [];
    for (const t of held) {
      const r = await this.repo.revertToPending(t.id);
      if (r !== null) reverted.push(r);
    }
    return reverted;
  }

  async findPendingForPools(poolNames: string[], limit?: number): Promise<InferenceTask[]> {
    return this.repo.findPendingForPools(poolNames, limit);
  }

  async findById(id: string): Promise<InferenceTask | null> {
    return this.repo.findById(id);
  }

  async list(filter: {
    ownerId?: string;
    status?: InferenceTaskStatus | InferenceTaskStatus[];
    poolName?: string;
    agentId?: string;
    limit?: number;
  }): Promise<InferenceTask[]> {
    return this.repo.list(filter);
  }

  async gcSweep(opts: { pendingTimeoutMs: number; terminalRetentionMs: number } = {
    pendingTimeoutMs: DEFAULT_PENDING_TIMEOUT_MS,
    terminalRetentionMs: DEFAULT_TERMINAL_RETENTION_MS,
  }): Promise<{ erroredOut: number; deleted: number }> {
    const now = this.clock();
    const pendingCutoff = new Date(now.getTime() - opts.pendingTimeoutMs);
    const terminalCutoff = new Date(now.getTime() - opts.terminalRetentionMs);

    // Pass 1: pending tasks that have aged out → flip to error so the
    // caller (if still waiting) gets a clean failure.
    const stale = await this.repo.findStalePending(pendingCutoff);
    let erroredOut = 0;
    for (const t of stale) {
      const r = await this.repo.markError(t.id, `Task expired in pending state after ${String(opts.pendingTimeoutMs / 1000)}s`, now);
      if (r !== null) {
        this.events.emit(`terminal:${t.id}`);
        erroredOut += 1;
      }
    }

    // Pass 2: terminal tasks past retention → delete.
    const expired = await this.repo.findExpiredTerminal(terminalCutoff);
    const deleted = await this.repo.deleteMany(expired.map((t) => t.id));

    return { erroredOut, deleted };
  }
}

function isTerminal(status: InferenceTaskStatus): boolean {
  return status === 'completed' || status === 'error' || status === 'cancelled';
}

/**
 * Surface a non-success terminal state as a thrown error so the HTTP
 * handler propagates it cleanly. `error` rows carry an `errorMessage`;
 * `cancelled` rows don't, so we synthesize one.
 */
function ensureSuccess(task: InferenceTask): InferenceTask {
  if (task.status === 'completed') return task;
  if (task.status === 'cancelled') {
    throw new Error(`InferenceTask cancelled: ${task.id}`);
  }
  if (task.status === 'error') {
    throw new Error(task.errorMessage ?? `InferenceTask failed: ${task.id}`);
  }
  // Theoretically unreachable — waitFor only resolves on terminal events.
  throw new Error(`InferenceTask in unexpected state: ${String(task.status)}`);
}
