# app/tools/vector_rag_tool.rb
class VectorRagTool
  # IGNORE All the langchain stuff
  extend Langchain::ToolDefinition

  define_function :retrieve,
                  description: "Semantic transcript search with optional metadata filter" do
    property :query,  type: "string",
                      description: "User question or search terms",
                      required: true

    property :k,      type: "integer",
                      description: "How many chunks to return"

    property :filter, type: "object",
                      description: "The filters for the query as-is",
                      required: false do
      property :tags,
               description: "Exact match on the pipe-delimited `tags` string",
               type: "string",
               required: false

      property :respondent_ids,
               type: "array",
               description: "Return chunks whose respondent_id is in this list",
               required: false do
        item type: "string"
      end
    end
  end

  def initialize(client)
    @client = client # Langchain::Vectorsearch::Chroma
  end

  def retrieve(query:, k: 20, filter: nil)
    Rails.logger.info "[VectorRagTool] Retrieving #{query}"
    # {} -> 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><citation>〈#{r.id}〉</citation><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)
    Rails.logger.info "[VectorRagTool] building where filters..."
    return {} if filter.blank?

    conditions = []

    # Handle section filter
    conditions << { "master_section_id" => { "$in" => [filter["section"].to_i] } } if filter["section"].present?

    # Handle respondent IDs
    if filter["respondent_ids"].is_a?(Array) && filter["respondent_ids"].any?
      conditions << { "respondent_id" => { "$in" => filter["respondent_ids"].compact.map(&:to_i) } }
    elsif filter["respondent_id"].present?
      conditions << { "respondent_id" => filter["respondent_id"].to_i }
    end

    # Combine all conditions with AND
    return {} if conditions.empty?
    return conditions.first if conditions.length == 1

    { "$and" => conditions }
  end
end