// Machine translation of a language: take untranslated lines page by page, send each page to the
// local translator, save what came back. Saving after every page means a stopped or crashed run
// loses at most one page, and the next run continues where this one ended.

import type { Job, TranslatorClient } from "./client";

export interface Source {
  key: string;
  source: string;
}

export interface RunOptions {
  client: Pick<TranslatorClient, "createJob" | "job" | "cancel">;
  projectId: string;
  batchSize: number;
  /** lines per translator job; each job is saved on its own */
  pageSize: number;
  signal: AbortSignal;
  /** the next `pageSize` items still to translate, skipping the first `skip` (the ones that failed) */
  fetchPage: (skip: number, limit: number) => Promise<{ total: number; items: Source[] }>;
  save: (results: { key: string; text: string }[]) => Promise<void>;
  /** a page is about to be sent to the model: these are the lines being translated now */
  onPage?: (items: Source[]) => void;
  onProgress: (p: Progress) => void;
  pollMs?: number;
  /** where an interrupted run stood: the counters and how many lines to skip */
  initial?: { total: number; done: number; failed: number; skip: number };
  /** the job that was being translated when the run was interrupted, and the lines it carries */
  resume?: { jobId: string; keys: string[] };
  /** a page was handed to the translator (`jobId`) or fully handled (`null`): enough to continue from here */
  onCheckpoint?: (cp: Checkpoint) => void;
}

export interface Checkpoint {
  jobId: string | null;
  keys: string[];
  skip: number;
  total: number;
  done: number;
  failed: number;
}

export interface Progress {
  total: number;
  done: number;
  failed: number;
  /** the latest reason a line was not translated */
  lastError: string;
  model: string | null;
}

export type Outcome = "finished" | "stopped";

const sleep = (ms: number, signal: AbortSignal) =>
  new Promise<void>((resolve) => {
    const t = setTimeout(resolve, ms);
    signal.addEventListener("abort", () => {
      clearTimeout(t);
      resolve();
    }, { once: true });
  });

/** Polls a job until it is over; aborting cancels it and returns what was finished by then. */
export async function waitForJob(client: Pick<TranslatorClient, "job" | "cancel">, first: Job,
                                 signal: AbortSignal, pollMs = 700): Promise<Job> {
  const opts = { client, signal, pollMs };
  let job = first;
  let errors = 0;
  // cancelling returns the job as it stands, with whatever was finished before the stop
  const stop = () => opts.client.cancel(job.id).catch(() => job);
  while (job.status === "queued" || job.status === "running") {
    if (opts.signal.aborted) return stop();
    await sleep(opts.pollMs ?? 700, opts.signal);
    if (opts.signal.aborted) return stop();
    try {
      job = await opts.client.job(job.id, opts.signal);
      errors = 0;
    } catch (e) {
      if (opts.signal.aborted) return stop();
      if (++errors >= 4) throw e;
    }
  }
  return job;
}

/** Throws when the translator becomes unreachable or a save fails. */
export async function runTranslation(opts: RunOptions): Promise<Outcome> {
  const start = opts.initial;
  const progress: Progress = { total: start?.total ?? 0, done: start?.done ?? 0, failed: start?.failed ?? 0, lastError: "", model: null };
  const emit = () => opts.onProgress({ ...progress });
  let skip = start?.skip ?? 0;
  let resume = opts.resume;
  const checkpoint = (jobId: string | null, keys: string[]) =>
    opts.onCheckpoint?.({ jobId, keys, skip, total: progress.total, done: progress.done, failed: progress.failed });

  while (!opts.signal.aborted) {
    let items: Source[];
    let job: Job;
    if (resume) {
      // the page was already sent before the page was reloaded: pick its result up
      const taken = resume;
      resume = undefined;
      const known = await opts.client.job(taken.jobId, opts.signal).catch(() => null);
      // the translator has forgotten it: the lines are still untranslated and come with the next page
      if (!known) continue;
      job = known;
      items = taken.keys.map((key, i) => ({ key, source: job.items[i]?.source ?? "" }));
      opts.onPage?.(items);
    } else {
      const page = await opts.fetchPage(skip, opts.pageSize);
      if (opts.signal.aborted || !page.items.length) break;
      // `page.total` still counts the lines that failed earlier, so it is the whole remaining work
      progress.total = progress.done + page.total;
      items = page.items.filter((it) => it.source.trim());
      // lines that stay untranslated keep matching the filter, so the next page starts after them
      skip += page.items.length - items.length;
      if (!items.length) continue;

      opts.onPage?.(items);
      job = await opts.client.createJob(
        opts.projectId, items.map((it, i) => ({ id: String(i), text: it.source })), opts.batchSize);
      checkpoint(job.id, items.map((it) => it.key));
    }
    job = await waitForJob(opts.client, job, opts.signal, opts.pollMs);
    progress.model = job.model ?? progress.model;

    const ready: { key: string; text: string }[] = [];
    for (const it of job.items) {
      const src = items[Number(it.id)];
      if (it.status === "completed" && it.translation?.trim() && src) ready.push({ key: src.key, text: it.translation });
      else if (job.status !== "cancelled") progress.lastError = it.error ?? "Модель вернула пустой перевод";
    }
    if (ready.length) await opts.save(ready);
    progress.done += ready.length;
    if (job.status !== "cancelled") progress.failed += items.length - ready.length;
    skip += items.length - ready.length;
    emit();
    checkpoint(null, []);
    if (job.status === "failed") throw new Error(job.error || "Переводчик не смог выполнить задание");
    if (job.status === "cancelled") break;
  }
  emit();
  return opts.signal.aborted ? "stopped" : "finished";
}