# frozen_string_literal: true

# This tool is used to get the context for a specific chunk.
class ContextualChunkTool
  extend Langchain::ToolDefinition

  # We need to store the collection to query it
  attr_reader :collection

  # Pass in your Chroma collection when you create the tool
  def initialize(client)
    @client = client
  end

  # 1. Define the function signature for the LLM
  define_function :get_surrounding_chunks,
                  description: "Gets the chunk immediately before, the chunk itself, and the chunk immediately after a specific chunk ID." do
    property :chunk_id,
             type: "string",
             description: "The *full* ID of the chunk to get context for, in the exact 'transcriptID-chunkIndex' format (e.g., '31732-0').",
             required: true
  end

  # 2. Define the Ruby method that performs the work
  def get_surrounding_chunks(chunk_id:)
    Rails.logger.info "[ContextTool] get_surrounding_chunks called with chunk_id: #{chunk_id}"
    # --- 1. Parse the incoming ID ---
    parts = chunk_id.rpartition("-")

    return "Error: Invalid chunk_id format. Expected 'transcriptID-chunkIndex', but got '#{chunk_id}'." unless parts[1] == "-" && parts.last.match?(/^\d+$/)

    transcript_id = parts.first
    current_index = parts.last.to_i

    # --- 2. Calculate previous and next indices ---
    prev_index = current_index - 1
    next_index = current_index + 1

    # --- 3. Fetch the three chunks ---
    prev_chunk_id = prev_index >= 0 ? "#{transcript_id}-#{prev_index}" : nil
    prev_chunk = fetch_chunk_text_from_db(prev_chunk_id)

    current_chunk_id = chunk_id
    current_chunk = fetch_chunk_text_from_db(current_chunk_id)

    next_chunk_id = "#{transcript_id}-#{next_index}"
    next_chunk = fetch_chunk_text_from_db(next_chunk_id)

    # --- 4. Assemble and return the context ---
    response = [
      { id: prev_chunk_id, text: prev_chunk },
      { id: current_chunk_id, text: current_chunk },
      { id: next_chunk_id, text: next_chunk }
    ]

    # Format the output clearly for the LLM
    response_text = response
                    .reject { |chunk| chunk[:text].nil? }
                    .map { |chunk| "--- Chunk #{chunk[:id]} ---\nCitation: 〈#{chunk[:id]}〉\n#{chunk[:text]}" }
                    .join("\n\n")

    return "No context found for '#{chunk_id}'." if response_text.empty?

    response_text
  rescue StandardError => e
    puts "Error in ContextualChunkTool: #{e.message}\n#{e.backtrace.join("\n")}"
    "Error processing chunk_id '#{chunk_id}'. Make sure it is in the correct format."
  end

  # --- This is your database logic, NOW IMPLEMENTED ---

  def fetch_chunk_text_from_db(full_chunk_id)
    return nil if full_chunk_id.nil?

    Rails.logger.info "[ContextTool] Retrieving #{full_chunk_id}"

    Chroma.connect_host = "http://localhost:8000"
    Chroma.logger = Logger.new(Rails.logger)
    Chroma.log_level = Chroma::LEVEL_ERROR
    collection = Chroma::Resources::Collection.get(@client.index_name)

    Rails.logger.info "[ContextTool] Fetching from DB: #{full_chunk_id}"

    begin
      results = collection.get(
        where: { "chunk_id" => full_chunk_id }, # This performs an exact metadata match
        limit: 1,
        include: ["documents"] # We only need the document/text
      )

      # The chroma-db gem returns a hash: {"ids"=>[], "documents"=>[], ...}
      document = results["documents"]&.first

      # Log if not found, which is normal for neighbors at the start/end
      Rails.logger.info "[ContextTool] Not found: #{full_chunk_id}" if document.nil?

      document # This will be the text string or nil if not found
    rescue Chroma::Error => e
      # Log the error but return nil so the agent can continue
      puts "ChromaDB error in fetch_chunk_text_from_db: #{e.message}"
      nil
    end
  end
end