class LmVectorRagTool < RubyLLM::Tool
  description "Semantic transcript search with optional metadata filter"

  params do
    string :query, description: "User question or search terms"
  end

  any_of do
    integer :k, description: "How many chunks to return"
    object :filter, description: "The filters for the query as-is" do
      string :tags, description: "Exact match on the pipe-delimited `tags` string"
      array :respondent_ids, of: :string, description: "Return chunks whose respondent_id is in this list"
    end
  end

  def execute(query:, k: 20, filter: nil)
    # {} -> Chroma syntax
    # translate the high-level filter into the Chroma `where` hash
    where = build_where(filter&.with_indifferent_access)

    Rails.logger.info("[FILTERS] WHERE FILTERS ARE: #{where.presence}")

    # First embed the query
    embedding = @client.llm.embed(text: query).embedding

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

    Rails.logger.info("Embedding size is: #{embedding.size}")

    Rails.logger.info("Results requested: #{k}")

    Rails.logger.info("Where is: #{where.presence}")

    # Query Chroma with the embedding + filters
    results = collection.query(
      query_embeddings: [embedding],
      results: k,
      where: where.presence
    )

    Rails.logger.info("[FILTERS] RESULTS ARE: #{results}")

    # If there are no results, instruct LLM accordingly
    return nil if results.empty?

    results.map do |r|
      "<document><id>#{r.id}</id><name>#{Respondent.find(r.metadata['respondent_id']).name}</name><content>#{r.document}</content></document>"
    end.join("\n\n")
  end

  private

  # turn the function-call args into Chroma syntax
  def build_where(filter)
    return {} if filter.blank?

    where = {}

    # --- Handle multiple tags with a logical OR ---
    if filter["tags"].present? && filter["tags"].is_a?(Array)
      if filter["tags"].compact.count == 1
        first_tag = filter["tags"].compact.first
        where["tags"] = { "$contains" => "|--|#{first_tag}|--|" }
      else
        # Build an array of conditions for the $or operator
        tag_conditions = filter["tags"].compact.map do |tag|
          { "$contains" => "|--|#{tag}|--|" }
        end

        # Apply the $or operator to the tags
        where["$or"] = tag_conditions
      end
    elsif filter["tags"].present?
      # Handle single tag case as before
      where["tags"] = { "$contains" => "|--|#{filter['tags']}|--|" }
    end

    # ---- plural respondent IDs ----------
    if filter["respondent_ids"].is_a?(Array) && filter["respondent_ids"].any?
      where["respondent_id"] = { "$in" => filter["respondent_ids"].compact.map(&:to_i) }
      # ---- single respondent ID -----------
    elsif filter["respondent_id"].present?
      where["respondent_id"] = filter["respondent_id"].to_i
    end

    where
  end
end