import { stateFile } from "./paths";
import {
  readJson,
  readState,
  readSalvaged,
  writeJsonAtomic,
  withStateLock,
  SALVAGED,
  type Reason,
  type SongState,
} from "./state";
import { keyOf, sourceWav, wordsFor, defaultWindowFor } from "./clips";
import { readWavWindow } from "./wav";
import type { Candidate, Mode, VerdictBody, VerdictResult } from "./types";

import { yinFrame, regionPitch } from "../song/pitch.mjs";
import { REASONS, OTHER_CODE } from "../song/reasons.mjs";

// ---------------------------------------------------------------------------
// Applying one verdict.
//
// These are the SAME files decode-um.mjs writes, and the static page stays
// usable, so both writers have to agree about what a verdict means. The
// transitions below mirror decode-um.mjs branch for branch:
//
//   sort     -> accepted / unclean / rejected, plus a re-trimmed window
//   salvage  -> the salvage machinery (cleaned, deliberately dirty, dropped)
//   verify   -> identical to salvage; the static page exports both as UMSALV/2
//   keeps    -> the salvage machinery, plus which side the stray word is on
//
// A moved window makes a NEW key (video@start), which is why every branch that
// can move one has to reject the old key as well as accept the new one.
// ---------------------------------------------------------------------------

const pct = (a: number[], q: number) => {
  const s = [...a].sort((u, v) => u - v);
  return s[Math.max(0, Math.min(s.length - 1, Math.floor(q * (s.length - 1))))];
};

function noteReason(
  reasons: Record<string, Reason>,
  key: string,
  verdict: number,
  code: number,
  text?: string,
) {
  if (!code) return;
  if (code === OTHER_CODE) {
    const t = text ?? "";
    reasons[key] = { verdict, code, why: t ? `other: ${t}` : "other (no text)", text: t };
    return;
  }
  const list = verdict === 2 ? REASONS.dirty : REASONS.bad;
  reasons[key] = { verdict, code, why: list[code - 1] ?? `#${code}` };
}

/** Pitch, spread and level for a window, exactly as decode-um.mjs measures them. */
async function measure(video: string, start: number, end: number) {
  const region = await readWavWindow(sourceWav(video), start - 0.01, end + 0.01);
  const sr = region.sampleRate;
  const a = Math.max(0, Math.round((start - region.from) * sr));
  const b = Math.min(region.x.length, Math.round((end - region.from) * sr));
  const seg = region.x.slice(a, b);
  if (seg.length < 0.05 * sr) return null;
  const r = regionPitch(seg, Math.round(0.005 * sr), seg.length - 1, sr, 55, 300);
  if (!r.f0) return null;
  const W = 2048;
  const hop = Math.round(sr * 0.01);
  const fs: number[] = [];
  for (let i = 0; i + W < seg.length; i += hop) {
    const q = yinFrame(seg, i, W, sr, 55, 300);
    if (q.conf >= 0.5 && q.f0 > 0) fs.push(q.f0);
  }
  const spread = fs.length ? pct(fs.map((v) => Math.abs(1200 * Math.log2(v / r.f0))), 0.8) : 0;
  let s2 = 0;
  for (let i = 0; i < seg.length; i += 1) s2 += seg[i] * seg[i];
  return {
    f0: +r.f0.toFixed(2),
    spreadCents: +spread.toFixed(1),
    rms: +(20 * Math.log10(Math.sqrt(s2 / seg.length) + 1e-9)).toFixed(1),
  };
}

// ---------------------------------------------------------------------------
// The journal, and undo.
//
// A verdict is not reversible by re-judging: a re-trimmed window creates a NEW
// key and REJECTS the old one, so "put it back" means restoring membership in
// several piles at once plus the map entries that came with it. Guessing at
// that would corrupt exactly the files that hold thousands of human decisions.
//
// So each verdict records the previous state of every key it touches, and undo
// replays that record rather than trying to invert the logic. Only the keys
// involved are captured -- the piles themselves are thousands of entries and
// snapshotting them per verdict would turn a keystroke into megabytes.
// ---------------------------------------------------------------------------

const JOURNAL = () => stateFile("journal.json");
/** Enough to cover a sitting; this is an undo stack, not an audit log. */
const JOURNAL_MAX = 200;

export type JournalEntry = {
  at: number;
  key: string;
  finalKey: string;
  mode: Mode;
  verdict: number;
  /** Every key whose membership was captured -- the old key and any new one. */
  keys: string[];
  before: {
    accepted: string[];
    rejected: string[];
    unclean: string[];
    provisional: string[];
    impure: string[];
    keepside: Record<string, number | null>;
    decided: Record<string, string | null>;
    reasons: Record<string, Reason | null>;
    corepitch: Record<string, { f0: number; dur: number } | null>;
    words: Record<string, { word: string; custom?: boolean } | null>;
  };
  /** Keys this verdict appended to _salvaged.json, to be removed on undo. */
  salvagedAdded: string[];
  needsContextAdded: string | null;
};

function snapshot(st: SongState, keys: string[]): JournalEntry["before"] {
  const has = (list: string[]) => keys.filter((k) => list.includes(k));
  const pick = <T,>(map: Record<string, T>) =>
    Object.fromEntries(keys.map((k) => [k, map[k] ?? null])) as Record<string, T | null>;
  return {
    accepted: has(st.accepted),
    rejected: has(st.rejected),
    unclean: has(st.unclean),
    provisional: has(st.provisional),
    impure: has(st.impure),
    keepside: pick(st.keepside),
    decided: pick(st.decided),
    reasons: pick(st.reasons),
    corepitch: pick(st.corepitch),
    words: pick(st.words),
  };
}

const salvagedKeys = async () =>
  new Set((await readSalvaged()).candidates.map((c) => keyOf(c as unknown as Candidate)));

export async function readJournal(): Promise<JournalEntry[]> {
  return readJson<JournalEntry[]>(JOURNAL(), []);
}

/**
 * Put the last verdict back exactly as it was.
 *
 * Runs under the same lock as applyVerdict, so it cannot interleave with a
 * verdict landing from another tab or from decode-um.mjs.
 */
export async function undoLastVerdict() {
  return withStateLock(async () => {
    const journal = await readJournal();
    const entry = journal[journal.length - 1];
    if (!entry) return { ok: false as const, error: "nothing to undo" };

    const st = await readState();
    const A = new Set(st.accepted);
    const R = new Set(st.rejected);
    const unclean = new Set(st.unclean);
    const prov = new Set(st.provisional);
    const impure = new Set(st.impure);
    const keepside = { ...st.keepside };
    const decided = { ...st.decided };
    const reasons = { ...st.reasons };
    const core = { ...st.corepitch };
    const words = { ...st.words };

    const setMembership = (set: Set<string>, wasIn: string[]) => {
      for (const k of entry.keys) {
        if (wasIn.includes(k)) set.add(k);
        else set.delete(k);
      }
    };
    setMembership(A, entry.before.accepted);
    setMembership(R, entry.before.rejected);
    setMembership(unclean, entry.before.unclean);
    setMembership(prov, entry.before.provisional);
    setMembership(impure, entry.before.impure);

    const restore = <T,>(map: Record<string, T>, was: Record<string, T | null>) => {
      for (const k of entry.keys) {
        const v = was[k];
        if (v == null) delete map[k];
        else map[k] = v;
      }
    };
    restore(keepside, entry.before.keepside);
    restore(decided, entry.before.decided);
    restore(reasons, entry.before.reasons);
    restore(core, entry.before.corepitch);
    restore(words, entry.before.words);

    // The hand-set window this verdict created has to go too, or the clip keeps
    // sounding at the edited boundaries even though the verdict is gone.
    if (entry.salvagedAdded.length) {
      const salvaged = await readSalvaged();
      const list = (salvaged.candidates as unknown as Candidate[]).filter(
        (c) => !entry.salvagedAdded.includes(keyOf(c)),
      );
      await writeJsonAtomic(SALVAGED(), { video: "_salvaged", candidates: list });
    }
    if (entry.needsContextAdded) {
      const nc = await readJson<Record<string, { edge: string; at: number }>>(
        stateFile("needs-context.json"),
        {},
      );
      delete nc[entry.needsContextAdded];
      await writeJsonAtomic(stateFile("needs-context.json"), nc);
    }

    const accepted = [...A];
    const rejected = [...R].filter((k) => !A.has(k));
    await writeJsonAtomic(stateFile("accepted.json"), { accepted, rejected });
    await writeJsonAtomic(stateFile("unclean.json"), [...unclean]);
    await writeJsonAtomic(stateFile("provisional.json"), [...prov]);
    await writeJsonAtomic(stateFile("impure.json"), [...impure]);
    await writeJsonAtomic(stateFile("keepside.json"), keepside);
    await writeJsonAtomic(stateFile("decided.json"), decided);
    await writeJsonAtomic(stateFile("corepitch.json"), core);
    await writeJsonAtomic(stateFile("reasons.json"), reasons);
    await writeJsonAtomic(stateFile("words.json"), words);
    await writeJsonAtomic(JOURNAL(), journal.slice(0, -1));

    return {
      ok: true as const,
      undone: { key: entry.key, finalKey: entry.finalKey, mode: entry.mode, verdict: entry.verdict },
      remaining: journal.length - 1,
      counts: {
        accepted: accepted.length,
        rejected: rejected.length,
        unclean: unclean.size,
        impure: impure.size,
      },
    };
  });
}

export async function applyVerdict(body: VerdictBody): Promise<VerdictResult> {
  return withStateLock(async () => {
    const { key, mode } = body;
    const [video, startStr] = [key.slice(0, key.lastIndexOf("@")), key.slice(key.lastIndexOf("@") + 1)];
    const origStart = Number(startStr);

    // Re-read everything. Nothing is held across requests: decode-um.mjs may
    // have rewritten any of these between one verdict and the next.
    const st = await readState();
    const A = new Set(st.accepted);
    const R = new Set(st.rejected);
    const unclean = new Set(st.unclean);
    const prov = new Set(st.provisional);
    const impure = new Set(st.impure);
    const keepside = { ...st.keepside };
    const decided = { ...st.decided };
    const reasons = { ...st.reasons };
    const core = { ...st.corepitch };
    const words = { ...st.words };

    const win = body.window;
    let finalKey = key;
    let moved = false;

    // Snapshot BEFORE anything moves. Both the key being judged and the key a
    // re-trim could create are captured, because a moved window rejects one and
    // accepts the other -- undoing only the key that was clicked would leave
    // the new one accepted forever.
    const maybeNewKey =
      win && win.to - win.from >= 0.12 ? `${video}@${(+win.from).toFixed(2)}` : null;
    const touchedKeys = [...new Set([key, maybeNewKey].filter(Boolean) as string[])];
    const journalBefore = snapshot(st, touchedKeys);
    const salvagedBefore = await salvagedKeys();

    // -- the label -----------------------------------------------------------
    // Written under BOTH keys when a re-trim renames the clip: the arranger
    // looks up the key it will actually place (the new one), and re-opening the
    // clip in this tool still finds the label under the key it was judged as.
    const labelWrites: string[] = [];

    if (mode === "sort") {
      noteReason(reasons, key, body.verdict, body.code ?? 0, body.text);
      if (body.verdict === 1) {
        // Judged clean by ear -- outranks any machine verdict, so drop the
        // provisional flag that would otherwise penalise it in the arranger.
        A.add(key);
        R.delete(key);
        unclean.delete(key);
        prov.delete(key);
        // A window adjusted during the fast pass is still a hand-set window and
        // must be honoured, not thrown away because this mode is meant to be quick.
        if (win) {
          // "Did it move?" is decided against the window the card was GIVEN,
          // which the client sends back untouched. Recomputing the default here
          // would re-read audio to answer a question the client already knows,
          // and would disagree with it the moment clipWindow's tuning changed.
          const d = body.defaultWindow;
          const changed =
            !d || Math.abs(win.from - d.from) >= 0.005 || Math.abs(win.to - d.to) >= 0.005;
          if (changed && win.to - win.from >= 0.12) {
            const nk = `${video}@${win.from.toFixed(2)}`;
            const salvaged = await readSalvaged();
            const list = salvaged.candidates as unknown as Candidate[];
            if (!list.some((c) => keyOf(c) === nk)) {
              list.push({
                video,
                token: body.word ?? "",
                start: +win.from.toFixed(3),
                end: +win.to.toFixed(3),
                dur: +(win.to - win.from).toFixed(3),
                f0: 0,
                spreadCents: 0,
                rms: -20,
                peak: 0,
                marginBefore: 0.05,
                marginAfter: 0.05,
                gapPrev: "",
                gapNext: "",
                noTrim: true,
                sortedFrom: key,
              });
              await writeJsonAtomic(SALVAGED(), { video: "_salvaged", candidates: list });
            }
            if (nk !== key) {
              A.delete(key);
              R.add(key);
              moved = true;
            }
            A.add(nk);
            R.delete(nk);
            prov.delete(nk);
            finalKey = nk;
          }
        }
      } else if (body.verdict === 2) {
        unclean.add(key);
        A.delete(key);
        R.delete(key);
        prov.delete(key);
      } else {
        R.add(key);
        A.delete(key);
        unclean.delete(key);
        prov.delete(key);
      }
    } else {
      // salvage / verify / keeps -- one machinery, three meanings.
      const isKeeps = mode === "keeps";
      const sideRaw = body.verdict;
      const keep = isKeeps ? (sideRaw === 0 ? 0 : 2) : body.verdict;
      decided[key] = isKeeps ? "keeps" : "salvage";
      if (!isKeeps) noteReason(reasons, key, keep, body.code ?? 0, body.text);

      unclean.delete(key); // the holding pile is ephemeral either way

      if (keep !== 1 && keep !== 2) {
        R.add(key);
        A.delete(key);
        impure.delete(key);
        delete keepside[key];
      } else {
        const start = +(win?.from ?? origStart).toFixed(3);
        const end = +(win?.to ?? origStart + 0.2).toFixed(3);
        const m = end - start < 0.12 ? null : await measure(video, start, end);
        if (!m) {
          // Too short, or no pitch at all -- decode-um.mjs drops these rather
          // than admitting a clip the renderer cannot tune.
          R.add(key);
          A.delete(key);
          impure.delete(key);
        } else {
          const nk = `${video}@${start.toFixed(2)}`;
          const salvaged = await readSalvaged();
          const list = salvaged.candidates as unknown as Candidate[];
          if (!list.some((c) => keyOf(c) === nk)) {
            list.push({
              video,
              token: body.word ?? "",
              start,
              end,
              dur: +(end - start).toFixed(3),
              f0: m.f0,
              spreadCents: m.spreadCents,
              rms: m.rms,
              peak: 0,
              marginBefore: 0.05,
              marginAfter: 0.05,
              gapPrev: "",
              gapNext: "",
              noTrim: true,
              salvagedFrom: key,
            });
            await writeJsonAtomic(SALVAGED(), { video: "_salvaged", candidates: list });
          }
          core[nk] = { f0: m.f0, dur: +(end - start).toFixed(3) };
          if (nk !== key) {
            A.delete(key);
            R.add(key);
            impure.delete(key);
            moved = true;
          }
          A.add(nk);
          R.delete(nk);
          prov.delete(nk);
          if (keep === 2) impure.add(nk);
          else impure.delete(nk);
          if (isKeeps) {
            if (nk !== key) delete keepside[key];
            keepside[nk] = sideRaw; // 1 head | 2 tail | 3 both
          }
          finalKey = nk;
        }
      }
    }

    labelWrites.push(finalKey);
    if (finalKey !== key) labelWrites.push(key);
    if (body.word) {
      for (const k of labelWrites) {
        words[k] = body.custom ? { word: body.word, custom: true } : { word: body.word };
      }
    }

    // Only ever a genuinely absent source file -- every source is local today,
    // so this records an edge to fetch later rather than silently truncating.
    let needsContextAdded: string | null = null;
    if (body.needsContext) {
      const nc = await readJson<Record<string, { edge: string; at: number }>>(
        stateFile("needs-context.json"),
        {},
      );
      if (!nc[finalKey]) needsContextAdded = finalKey;
      nc[finalKey] = { edge: body.needsContext, at: Math.floor(Date.now() / 1000) };
      await writeJsonAtomic(stateFile("needs-context.json"), nc);
    }

    const accepted = [...A];
    const rejected = [...R].filter((k) => !A.has(k));
    await writeJsonAtomic(stateFile("accepted.json"), { accepted, rejected });
    await writeJsonAtomic(stateFile("unclean.json"), [...unclean]);
    await writeJsonAtomic(stateFile("provisional.json"), [...prov]);
    await writeJsonAtomic(stateFile("impure.json"), [...impure]);
    await writeJsonAtomic(stateFile("keepside.json"), keepside);
    await writeJsonAtomic(stateFile("decided.json"), decided);
    await writeJsonAtomic(stateFile("corepitch.json"), core);
    await writeJsonAtomic(stateFile("reasons.json"), reasons);
    if (body.word) await writeJsonAtomic(stateFile("words.json"), words);

    // The journal is written LAST, so a verdict that failed part-way through
    // leaves nothing claiming to be undoable.
    const salvagedAfter = await salvagedKeys();
    const journal = await readJournal();
    journal.push({
      at: Date.now(),
      key,
      finalKey,
      mode,
      verdict: body.verdict,
      keys: [...new Set([...touchedKeys, finalKey])],
      before: touchedKeys.includes(finalKey)
        ? journalBefore
        : // A finalKey nobody predicted (the salvage branch derives it from a
          // measured window) still needs its previous state, and "absent from
          // everything" is exactly what it was.
          {
            ...journalBefore,
            keepside: { ...journalBefore.keepside, [finalKey]: st.keepside[finalKey] ?? null },
            decided: { ...journalBefore.decided, [finalKey]: st.decided[finalKey] ?? null },
            reasons: { ...journalBefore.reasons, [finalKey]: st.reasons[finalKey] ?? null },
            corepitch: { ...journalBefore.corepitch, [finalKey]: st.corepitch[finalKey] ?? null },
            words: { ...journalBefore.words, [finalKey]: st.words[finalKey] ?? null },
            accepted: st.accepted.includes(finalKey)
              ? [...journalBefore.accepted, finalKey]
              : journalBefore.accepted,
            rejected: st.rejected.includes(finalKey)
              ? [...journalBefore.rejected, finalKey]
              : journalBefore.rejected,
            unclean: st.unclean.includes(finalKey)
              ? [...journalBefore.unclean, finalKey]
              : journalBefore.unclean,
            provisional: st.provisional.includes(finalKey)
              ? [...journalBefore.provisional, finalKey]
              : journalBefore.provisional,
            impure: st.impure.includes(finalKey)
              ? [...journalBefore.impure, finalKey]
              : journalBefore.impure,
          },
      salvagedAdded: [...salvagedAfter].filter((k) => !salvagedBefore.has(k)),
      needsContextAdded,
    });
    await writeJsonAtomic(JOURNAL(), journal.slice(-JOURNAL_MAX));

    return {
      ok: true as const,
      key,
      finalKey,
      moved,
      counts: {
        accepted: accepted.length,
        rejected: rejected.length,
        unclean: unclean.size,
        impure: impure.size,
      },
    };
  });
}

export function reasonsForMode(mode: Mode) {
  // Level 2 of the verdict buttons. The lists are APPEND-ONLY: reasons.json
  // stores the number, so renumbering rewrites history.
  return {
    bad: REASONS.bad.map((why, i) => ({ code: i + 1, why })),
    dirty: REASONS.dirty.map((why, i) => ({ code: i + 1, why })),
    other: OTHER_CODE,
    mode,
  };
}