# frozen_string_literal: true

# Transcription A/B harness - runs the same recordings through different Whisper settings and scores
# them, so a change can be proven rather than assumed.
#
# Why it exists: the transcripts have real Czech errors ("badíček" for "balíček", "myštěný" for
# "zmeškaný"), and there are several plausible fixes - a different model, different speaker hints, less
# lossy audio from the phone. Without measurement they are all just opinions, and some of them make
# things worse. The existing tests (test_remote_transcriber.rb, test_pipeline.rb) are offline fakes that
# check plumbing, never accuracy.
#
# Usage:
#   ruby components/transcription-worker/bench/ab.rb [corpus_dir]      # default: bench/corpus
#
# Corpus layout - the reference is what you SAID, typed out by hand:
#   corpus/call-01.m4a
#   corpus/call-01.txt        <- optional reference. Without it you still get transcripts to read,
#                                but no score, because there is nothing to be right or wrong against.
#
# Needs the same mTLS env as production (WHISPER_URL, WHISPER_MTLS_P12, WHISPER_MTLS_PASS).
# Results are cached under bench/out/, because the service is CPU-only and a re-run is minutes, not
# seconds. Delete a cached file to force that one job again.
#
# grep anchors: bench-variants, bench-normalise, bench-wer.

require "json"
require "fileutils"
require_relative "../remote_transcriber"

module Helpdesk
  module Bench
    # The settings worth comparing. Everything here is a real parameter the hosted service accepts -
    # it takes exactly five fields (file, model, language, min_speakers, max_speakers) and no others,
    # so there is no prompt or decoding knob to vary. grep anchor: bench-variants.
    VARIANTS = [
      { name: "baseline",   model: "large-v2", min_speakers: 2, max_speakers: 3 }, # what production does today
      { name: "v2-2spk",    model: "large-v2", min_speakers: 2, max_speakers: 2 }, # stop it inventing a third voice
      { name: "v3",         model: "large-v3", min_speakers: 2, max_speakers: 3 }, # the service calls v2 better for Czech - check
      { name: "v3-2spk",    model: "large-v3", min_speakers: 2, max_speakers: 2 },
    ].freeze

    AUDIO_EXT = %w[.m4a .mp3 .wav .ogg .flac .mp4 .opus .webm .aac .audio].freeze

    module_function

    # ---- scoring, all pure so it is testable without touching the network -------------------------

    # Strip the service's "[00:00:01] [Speaker 1]: " prefixes down to the words that were said.
    # grep anchor: bench-normalise.
    def strip_markup(text)
      text.to_s.lines.map { |l| l.sub(/\A\s*\[\d{2}:\d{2}:\d{2}\]\s*/, "").sub(/\A\s*\[[^\]]+\]:\s*/, "") }
          .join(" ")
    end

    # Compare like with like: case and punctuation are not transcription errors. Diacritics ARE kept -
    # in Czech they change the word, so folding them would hide real mistakes.
    # NOTE the en/en-dash characters in the punctuation class below are deliberate and are NOT a style
    # violation of the repo's no-dashes rule: that rule governs prose we write, while this strips
    # punctuation out of Czech text Whisper produces, which does contain both. Do not "fix" them.
    PUNCT = /[.,!?;:"„“”…()\-–—]/.freeze

    def normalise(text)
      strip_markup(text).downcase
                        .gsub(PUNCT, " ")
                        .gsub(/\s+/, " ")
                        .strip
    end

    def words(text) = normalise(text).split(" ")

    # Word error rate: edit distance over WORDS divided by the reference length. 0.0 is perfect, and it
    # can exceed 1.0 when the output is longer than the reference (Whisper padding a mumble with
    # invented words is a real failure mode, so that is not clamped). grep anchor: bench-wer.
    def wer(reference, hypothesis)
      r = words(reference)
      h = words(hypothesis)
      return { wer: nil, ref_words: 0 } if r.empty?
      d = levenshtein(r, h)
      { wer: d.fdiv(r.length), edits: d, ref_words: r.length, hyp_words: h.length }
    end

    # Align two word sequences and label every position. Full matrix with a backtrace, because knowing
    # WHERE they diverge is what makes the output usable - a bare distance tells you a number, an
    # alignment tells you which words to look at. grep anchor: bench-align.
    # -> [[:ok|:sub|:del|:ins, ref_word_or_nil, hyp_word_or_nil], ...]
    def align(r, h)
      m = Array.new(r.length + 1) { Array.new(h.length + 1, 0) }
      (0..r.length).each { |i| m[i][0] = i }
      (0..h.length).each { |j| m[0][j] = j }
      r.each_with_index do |rw, i|
        h.each_with_index do |hw, j|
          m[i + 1][j + 1] = [m[i][j + 1] + 1, m[i + 1][j] + 1, m[i][j] + (rw == hw ? 0 : 1)].min
        end
      end
      ops = []
      i = r.length; j = h.length
      while i.positive? || j.positive?
        if i.positive? && j.positive? && m[i][j] + (r[i - 1] == h[j - 1] ? 0 : 1) == m[i][j]
          ops << [r[i - 1] == h[j - 1] ? :ok : :sub, r[i - 1], h[j - 1]]; i -= 1; j -= 1
        elsif i.positive? && m[i - 1][j] + 1 == m[i][j]
          ops << [:del, r[i - 1], nil]; i -= 1
        elsif j.positive? && m[i][j - 1] + 1 == m[i][j]
          ops << [:ins, nil, h[j - 1]]; j -= 1
        else
          ops << [r[i - 1] == h[j - 1] ? :ok : :sub, r[i - 1], h[j - 1]]; i -= 1; j -= 1
        end
      end
      ops.reverse
    end

    # Group runs of disagreement into chunks, each with the reference words and what replaced them.
    # Adjacent edits belong together: "na shledanou" -> "naschledanou" is ONE disagreement, and scoring
    # it as two makes a formatting quirk look twice as bad as a misheard word.
    def chunks(reference, hypothesis)
      out = []
      cur = nil
      align(words(reference), words(hypothesis)).each do |op, rw, hw|
        if op == :ok
          out << cur if cur
          cur = nil
        else
          cur ||= { ref: [], hyp: [] }
          cur[:ref] << rw if rw
          cur[:hyp] << hw if hw
        end
      end
      out << cur if cur
      out.flat_map { |c| split_chunk(c[:ref], c[:hyp]) }
    end

    # A run of consecutive edits can mix kinds - "děkuju mockrát na shledanou" -> "děkuji moc krát
    # naschledanou" is one run but three separate phenomena, and calling the whole thing a
    # misrecognition overstates it. Peel off any leading piece that has a benign explanation, then
    # recurse on the rest; whatever cannot be explained stays as one :real chunk.
    def split_chunk(ref, hyp)
      whole = [{ ref: ref, hyp: hyp, kind: classify_chunk(ref, hyp) }]
      return whole if ref.length + hyp.length <= 2 || whole[0][:kind] != :real
      (1..ref.length).each do |i|
        (1..hyp.length).each do |j|
          next if i == ref.length && j == hyp.length # no progress
          k = classify_chunk(ref[0, i], hyp[0, j])
          next if k == :real
          rest_r = ref[i..] || []
          rest_h = hyp[j..] || []
          rest = (rest_r.empty? && rest_h.empty?) ? [] : split_chunk(rest_r, rest_h)
          return [{ ref: ref[0, i], hyp: hyp[0, j], kind: k }] + rest
        end
      end
      whole
    end

    # What KIND of disagreement is this? Only :real should count against the audio. Deciding this
    # automatically is the difference between a score you can trust and one you have to re-argue every
    # time. grep anchor: bench-classify.
    NUM_WORDS = %w[nula jedna jedné jeden jedno dva dvě tři čtyři pět šest sedm osm devět deset
                   jedenáct dvanáct třináct čtrnáct patnáct šestnáct sedmnáct osmnáct devatenáct
                   dvacet třicet čtyřicet padesát šedesát sedmdesát osmdesát devadesát sto stě set
                   tisíc lomeno].freeze
    def classify_chunk(ref, hyp)
      r = ref.join; h = hyp.join
      return :boundary if r == h                       # same letters, different spacing
      return :diacritic if strip_diacritics(r) == strip_diacritics(h)
      return :variant if spoken_variant?(r, h)
      return :number if (ref + hyp).any? { |w| NUM_WORDS.include?(w) || w.match?(/\d/) }
      :real
    end

    # Czech first-person verbs have a spoken form in -u and a written one in -i ("děkuju" / "děkuji",
    # "potřebuju" / "potřebuji"). Whisper writes the standard form whichever was said, so counting that
    # as a transcription error would penalise it for being right. Same stem, endings differing only
    # u <-> i, is the whole rule.
    def spoken_variant?(a, b)
      return false if a.length < 3 || b.length < 3
      return false unless a[0..-2] == b[0..-2]
      %w[u i].include?(a[-1]) && %w[u i].include?(b[-1]) && a[-1] != b[-1]
    end

    def strip_diacritics(s)
      s.tr("áčďéěíňóřšťúůýžÁČĎÉĚÍŇÓŘŠŤÚŮÝŽ", "acdeeinorstuuyzACDEEINORSTUUYZ")
    end

    # Word error rate ignoring differences that are not misrecognitions. Boundary and diacritic-only
    # chunks are dropped; everything else counts. This is the number worth comparing between runs.
    def effective_wer(reference, hypothesis)
      r = words(reference)
      return { wer: nil } if r.empty?
      cs = chunks(reference, hypothesis)
      counted = cs.reject { |c| %i[boundary diacritic variant].include?(c[:kind]) }
      edits = counted.sum { |c| [c[:ref].length, c[:hyp].length].max }
      { wer: edits.fdiv(r.length), edits: edits, ref_words: r.length,
        by_kind: cs.group_by { |c| c[:kind] }.transform_values(&:length) }
    end

    # Character error rate with spaces removed - immune to where the model put word boundaries, and a
    # useful sanity check on the word-level numbers.
    def cer(reference, hypothesis)
      r = normalise(reference).delete(" ").chars
      return nil if r.empty?
      levenshtein(r, normalise(hypothesis).delete(" ").chars).fdiv(r.length)
    end

    # Standard Levenshtein over arrays, two rows rather than a full matrix.
    def levenshtein(a, b)
      return b.length if a.empty?
      return a.length if b.empty?
      prev = (0..b.length).to_a
      a.each_with_index do |ax, i|
        cur = [i + 1]
        b.each_with_index do |bx, j|
          cur << [prev[j + 1] + 1, cur[j] + 1, prev[j] + (ax == bx ? 0 : 1)].min
        end
        prev = cur
      end
      prev[b.length]
    end

    # Pull one speaker's words out of a merged dual-channel transcript.
    def speaker_side(transcript, label)
      transcript.to_s.lines.select { |l| l.include?("[#{label}]:") }
                .map { |l| l.sub(/.*\]:/, "") }.join(" ")
    end

    # Score a dual-channel transcript PER SPEAKER. Scoring a merged transcript against a linear script
    # measures the wrong thing: the merge is ordered by timestamp and the script by reading order, so
    # whenever turns group differently the alignment collapses and reports a huge error rate for text
    # that is actually correct. On a real call that inflated 10% to 44%. Comparing each speaker's words
    # against only their own lines removes the ordering entirely.
    # -> { "Operator" => {wer:, cer:, ...}, "Caller" => {...} }
    # grep anchor: bench-per-channel.
    def score_channels(transcript, refs)
      refs.each_with_object({}) do |(label, ref), out|
        hyp = speaker_side(transcript, label)
        out[label] = effective_wer(ref, hyp).merge(cer: cer(ref, hyp), hyp_words: words(hyp).length)
      end
    end

    # How many distinct voices the diariser decided there were. A two-party phone call should be 2;
    # 3 means it split someone, 1 means it merged them.
    def speaker_count(text) = text.to_s.scan(/\[Speaker\s+(\d+)\]/i).flatten.uniq.length

    # The words the reference has that the hypothesis does not, and vice versa - this is what tells you
    # WHICH words are being mangled, which is the whole point of looking at a corpus.
    def word_diff(reference, hypothesis, limit: 12)
      r = words(reference).tally
      h = words(hypothesis).tally
      missing = r.reject { |w, n| h[w].to_i >= n }.keys.first(limit)
      spurious = h.reject { |w, n| r[w].to_i >= n }.keys.first(limit)
      { missing: missing, spurious: spurious }
    end

    # ---- running ----------------------------------------------------------------------------------

    def corpus(dir)
      Dir.glob(File.join(dir, "*")).select { |f| AUDIO_EXT.include?(File.extname(f).downcase) }.sort
    end

    def reference_for(audio)
      ref = audio.sub(/#{Regexp.escape(File.extname(audio))}\z/, ".txt")
      File.exist?(ref) ? File.read(ref) : nil
    end

    def run(dir)
      files = corpus(dir)
      if files.empty?
        warn "no audio in #{dir}. Put recordings there, plus a matching .txt of what was actually said."
        return 1
      end
      out = File.join(__dir__, "out")
      FileUtils.mkdir_p(out)
      transcriber = RemoteTranscriber.new

      rows = []
      files.each do |audio|
        ref = reference_for(audio)
        base = File.basename(audio, ".*")
        puts "\n#{base}#{ref ? '' : '  (no reference: transcripts only, no score)'}"
        VARIANTS.each do |v|
          cached = File.join(out, "#{base}.#{v[:name]}.txt")
          if File.exist?(cached)
            text = File.read(cached)
            note = "cached"
          else
            t0 = Time.now
            res = transcriber.transcribe(audio, model: v[:model],
                                         min_speakers: v[:min_speakers], max_speakers: v[:max_speakers])
            unless res[:state] == "transcribed"
              puts "  %-10s FAILED (%s)" % [v[:name], res[:reason] || res[:state]]
              next
            end
            text = res[:transcript]
            File.write(cached, text)
            note = "#{(Time.now - t0).round}s"
          end
          eff = ref ? effective_wer(ref, text) : nil
          rows << { file: base, variant: v[:name], wer: eff && eff[:wer], speakers: speaker_count(text) }
          puts "  %-10s speakers=%d  %s  %s" % [
            v[:name], speaker_count(text),
            eff && eff[:wer] ? "WER=#{(eff[:wer] * 100).round(1)}%  CER=#{(cer(ref, text) * 100).round(1)}%"
                             : "WER=n/a",
            note
          ]
          if ref && eff[:wer]
            # Only the misrecognitions are worth reading; boundary/diacritic/variant chunks are noise
            # the score already discounts, and printing them buries the ones that matter.
            bad = chunks(ref, text).reject { |c| %i[boundary diacritic variant].include?(c[:kind]) }
            bad.first(8).each do |c|
              puts "             [%-8s] %s -> %s" % [c[:kind], c[:ref].join(" "), c[:hyp].join(" ")]
            end
            puts "             (+#{bad.length - 8} more)" if bad.length > 8
            skipped = eff[:by_kind].select { |k, _| %i[boundary diacritic variant].include?(k) }
            puts "             discounted: #{skipped.map { |k, n| "#{n} #{k}" }.join(', ')}" if skipped.any?
          end
        end
      end

      scored = rows.select { |r| r[:wer] }
      unless scored.empty?
        puts "\n=== mean WER across #{scored.map { |r| r[:file] }.uniq.length} recording(s), lower is better ==="
        scored.group_by { |r| r[:variant] }.each do |name, rs|
          mean = rs.sum { |r| r[:wer] } / rs.length
          puts "  %-10s %.1f%%" % [name, mean * 100]
        end
        puts "\nA difference under a point or two on a handful of calls is noise, not an improvement."
      end
      puts "\ntranscripts cached in #{out} (delete one to re-run that job)"
      0
    end
  end
end

exit(Helpdesk::Bench.run(ARGV[0] || File.join(__dir__, "corpus"))) if $PROGRAM_NAME == __FILE__