# frozen_string_literal: true

# WikiIndex — the RAG core. Chunks wiki pages, embeds each chunk (via an injected Embedder), stores
# {text, title, path, vector}, and retrieves the top-K chunks for a query by cosine similarity. Plain
# in-memory cosine (no vector DB — the corpus is a few thousand chunks). Standard library only; JSON persistence.

require "json"

module Helpdesk
  class WikiIndex
    Chunk = Struct.new(:text, :title, :path, :vector)

    def initialize(embedder:, chunk_words: 150)
      @embedder = embedder
      @chunk_words = chunk_words
      @chunks = []
    end

    attr_reader :chunks

    # pages: [{ title:, path:, text: }]  — embeds in batches to limit request count.
    def build(pages, batch: 64)
      units = []
      pages.each do |pg|
        chunk_text(pg[:text]).each { |c| units << { text: c, title: pg[:title], path: pg[:path] } }
      end
      units.each_slice(batch) do |slice|
        vecs = @embedder.embed(slice.map { |u| u[:text] })
        # Fail at index-build time (offline) if the embedder returned fewer/nil vectors. Otherwise the
        # tail chunks get nil vectors, get saved as "vector": null, and crash cosine() on-device — RAG
        # then goes dark for EVERY call with a single swallowed warning. Catch it here, not on hardware.
        raise "embedder returned #{vecs.length} vectors for #{slice.length} chunks" unless vecs.length == slice.length
        slice.each_with_index do |u, i|
          v = vecs[i]
          raise "nil/empty embedding for chunk #{u[:title].inspect}##{i}" unless v.is_a?(Array) && !v.empty?
          @chunks << Chunk.new(u[:text], u[:title], u[:path], v)
        end
      end
      self
    end

    # → [{ text:, title:, path:, score: }] best-first
    def retrieve(query, k: 5)
      return [] if @chunks.empty?
      qv = @embedder.embed_one(query)
      dim = qv.length
      # Defensive: a loaded index built against a different embedding model (or a partial batch) can
      # carry nil/wrong-dim vectors. Skip them (and say so) rather than crash the whole query.
      usable = @chunks.select { |c| c.vector.is_a?(Array) && c.vector.length == dim }
      if usable.size < @chunks.size
        warn "[helpdesk] wiki retrieve: skipped #{@chunks.size - usable.size}/#{@chunks.size} chunk(s) with nil/mismatched vectors (index/model drift?)"
      end
      usable
        .map { |c| [cosine(qv, c.vector), c] }
        .sort_by { |score, _| -score }
        .first(k)
        .map { |score, c| { text: c.text, title: c.title, path: c.path, score: score.round(4) } }
    end

    def save(path)
      File.write(path, JSON.generate(@chunks.map { |c| { text: c.text, title: c.title, path: c.path, vector: c.vector } }))
      self
    end

    def load(path)
      @chunks = JSON.parse(File.read(path)).map { |h| Chunk.new(h["text"], h["title"], h["path"], h["vector"]) }
      self
    end

    def size = @chunks.size

    private

    def chunk_text(text)
      words = text.to_s.split(/\s+/)
      return [] if words.empty?
      words.each_slice(@chunk_words).map { |w| w.join(" ") }
    end

    def cosine(a, b)
      # Require equal dimensions — truncating to the shorter vector masks a model/index mismatch and
      # would score a 1-dim corrupt vector at 1.0, floating garbage to the top of retrieve().
      raise ArgumentError, "cosine dim mismatch: #{a&.length.inspect} vs #{b&.length.inspect}" unless a.is_a?(Array) && b.is_a?(Array) && a.length == b.length
      dot = 0.0
      na = 0.0
      nb = 0.0
      i = 0
      n = a.length
      while i < n
        dot += a[i] * b[i]
        na += a[i] * a[i]
        nb += b[i] * b[i]
        i += 1
      end
      denom = Math.sqrt(na) * Math.sqrt(nb)
      denom.zero? ? 0.0 : dot / denom
    end
  end
end