# app/services/rag_assistant.rb
# frozen_string_literal: true

require "tiktoken_ruby"
require "securerandom"

# Write docs
class RagAssistant
  PROMPT = <<~RAGPROMPT
    YOUR MANDATORY WORKFLOW

    Your default tool is ChunkSearch. You MUST use it to find specific, citable information.

    For Specific Questions (e.g., "What was the Q3 budget?"):

    Use ChunkSearch to find the answer.

    Formulate your response based only on the "Content" provided.

    You MUST end each fact with its ID citation as shown.

    Example Answer: "The Q3 budget was set at $500,000. 〈5〉"

    For Broad Questions (e.g., "Summarize the 'Q3 Planning' section"):

    You may first use ChunkSearch to locate the section (e.g., ChunkSearch(query: "Q3 Planning section")).

    From the "Source" string in the result, extract the respondent ID (e.g., '789').

    Use DocumentBrowser(id: 789) to get the complete text of that document.

    Write your summary based on this full text.

    You MUST still cite the overall section you are summarizing.

    Example Answer: "The 'Q3 Planning' section outlines the new market expansion strategy and the supporting marketing campaign (Source: Section 'Q3 Planning' in 'Q3 Kickoff')."

    CRITICAL RULES

    Always Cite: Never state a fact from a transcript without its corresponding Source: citation. For summaries of full sections, cite the section itself.

    Use Chunks for Sourcing: Your primary goal is to use ChunkSearch results for sourcing. Only use FullTextReader when you need broader context that a single chunk cannot provide (like writing a summary).

    Stick to the Context: Do NOT make up information or answer from your general knowledge. If the tools do not provide an answer, say so.

    Clarity: If the user's query is ambiguous (e.g., "what did he say?"), ask for clarification before using any tools.

    MOST IMPORTANT: When sourcing, always cite your source using this format: 〈id〉.  Use the ID that is returned to you in the citation. Do not copy the citation. Include only the Chunk ID within the special brackets and nothing else. Always base your answers on the most recent documents shared in the context, as the earlier documents may have been filtered out in subsequent requests.

    Execute your tools sequentially. Do not call multiple tools at the exact same time. Wait for the result of one ChunkSearch before deciding to search again.
  RAGPROMPT

  MAX_TOKENS = 2_000
  OVERLAP = 200
  Chunk = Struct.new(:text, :id, :metadata)

  def get_client
    @client
  end

  def initialize(assistant_conversation, task_id: nil, project_id: nil, respondent_category_id: nil)
    @conversation = assistant_conversation
    @project_id = project_id
    @respondent_category_id = respondent_category_id
    raggable       = assistant_conversation.raggable
    respondents    = raggable.respondents

    @raggable_type = raggable.instance_of?(Project) ? "Project" : "RespondentCategory"

    Rails.logger.debug("Raggable Type -> #{@raggable_type}")

    language =
      case @raggable_type
      when "RespondentCategory"
        if raggable.has_homogenous_language_respondents
          raggable.respondents.first&.transcript_language || "en"
        else
          "en"
        end
      when "Project"
        if raggable.has_homogenous_language_subprojects
          raggable.respondent_categories.first&.respondents&.first&.transcript_language || "en"
        else
          "en"
        end
      else
        "en"
      end

    @conversation.update!(language:)
    language_prompt = "All user questions and your responses are in #{Respondent::LANGUAGES.transform_values(&:titleize)[language.to_sym]}. Always think, reason, and respond only in this language. Never use any other language."

    # ------------------- LLM -------------------------------------------------
    #    @llm = Langchain::LLM::OpenAI.new(
    #      api_key: Rails.application.credentials.openai_secret,
    #      default_options: { model: "gpt-5", temperature: 0.1, stream: true }
    #    )
    model_map = FeatureModelMap.for_feature!("RAG_ASSISTANT")
    @llm = LangchainLlmFactory.build(model_map:)

    # ------------------- Vector DB ------------------------------------------
    index_name = raggable.index_name.presence || SecureRandom.hex(24)
    Rails.logger.debug("Generated Index Name -> #{index_name}")
    @client = Langchain::Vectorsearch::Chroma.new(
      url: "http://localhost:8000",
      index_name:,
      llm: @llm
    )

    if raggable.index_name.blank?
      @client.create_default_schema
      build_index(respondents, task_id:) # expensive only once
      raggable.update!(index_name:) # ← persist for next time
    end

    # ------------------- Streaming hook -------------------------------------
    @current_message = nil
    @stream_callback = proc do |delta|
      # Rails.logger.info("[LLM] delta: #{delta}")
      token = begin
        delta.dig("delta", "content")
      rescue StandardError
        nil
      end
      next if token.blank? || @current_message.nil?

      @current_message.broadcast_append_chunk(token)
    end

    vector_tool = VectorRagTool.new(@client) # VectorRagTool instance
    contextual_chunk_tool = ContextualChunkTool.new(@client)

    instructions = <<~INSTRUCTIONS
      You are a helpful qualitative research assistant with several years of experience and expertise. You will carry out a variety of tasks for the user using the tools available to you. You will answer questions about the qualitative research that has been conducted based on the content of numerous transcripts available to you. These transcripts are searchable and you can search for chunks from multiple transcripts. Each document is a transcript of a conversation between a moderator and some respondents.

      Each document with a different ID is a different transcript. It's likely that respondents will have the same name across documents like Respondent 0,1,2,3,4 or Speaker 0,1,2,3,4 etc. in multiple documents.

      These refer to different people if they are in different documents. You should refer to the document name when you are evaluating or discussing context.

      Very important: You need to always cite sources for your answers. Cite all relevant sources for all your claims, not just one. There is no situation in which you are allowed to answer without citing a source. Copy the exact value from each document's <citation> tag, such as 〈123-0〉. Use only these exact brackets: 〈 and 〉. Do not use similar brackets like 〈 and 〉. Always base your answers on the most recent documents shared in the context, as the earlier documents may have been filtered out in subsequent requests.

      The user has the option to exclude documents that are sent to you with each message.

      Tool Usage:
      1. To search, use the 'VectorSearchTool' to find the most relevant chunks.
      2. Read the results. If a chunk seems incomplete or mentions something you don't know, search for it to learn more if you need to.
      3. You may need to do this multiple times to get the most relevant results.
      4. Once you are confident you have the full picture, synthesize all the retrieved text into a single, comprehensive answer.
      5. The transcripts are quite long and detailed, so it's very possible that one search may not be enough. A typical transcript will be at least 1-2 hours worth of spoken conversation.
      6. Consider requesting a reasonable number of chunks to ensure the most relevant results.
      7. When mentioning chunks, copy the document's <citation> value exactly to reference the chunk.

      Rules:

      1. Never use lazy references like "the first" or "the last" in your answer. Always name what you are referring to.
      2. If the name of what you are referring to is ambiguous, look for it in nearby chunks till you find it.
      3. If the conversation refers to something unnamed, you can provide a reference that mentions the name of the document and then use the reference as it is being used in the conversation.
      4. When referring to something in one of the documents, ensure you mention the document name when you are referring to it.
      5. It's possible that whatever you are referring to ALSO exists in other documents. You should be precise to ensure there is minimal confusion.

      Help the user by looking through the documents extensively and ensure that you provide compehensive and accurate answer. You can use the contextual chunk tool to check a few chunks before and after the search result to clarify that you are reading the extracted portions in the correct context. Make a plan and execute it. Work methodically. Do not output private reasoning, hidden thoughts, scratchpad notes, planning steps, or tool-selection analysis. Only output the final user-facing answer with citations.

      #{language_prompt}
    INSTRUCTIONS
    # ------------------- Assistant ------------------------------------------

    # Tool execution callback to broadcast status updates
    tool_callback = proc do |tool_call_id, tool_name, method_name, tool_arguments|
      Rails.logger.info "[RagAssistant] Tool executing: #{tool_name}.#{method_name} with args: #{tool_arguments}"
      broadcast_status("Using tools...")
    end

    @assistant = Langchain::Assistant.new(
      llm: @llm,
      tools: [vector_tool, contextual_chunk_tool],
      instructions:,
      add_message_callback: @stream_callback,
      tool_execution_callback: tool_callback,
      parallel_tool_calls: gemini_llm? ? false : true
    ) { |tok| @stream_callback.call(tok) }
  end

  # --------------------------------------------------------------------------
  # PUBLIC API
  # --------------------------------------------------------------------------

  # Broadcasts a status update to the typing indicator
  def broadcast_status(status_text, tool_name: nil, tool_args: nil)
    # Generate a descriptive message based on what's happening
    message = case tool_name
              when "VectorRagTool"
                query = tool_args&.dig(:query)
                truncated_query = query&.truncate(45, omission: "...")
                "Searching transcripts for: #{truncated_query}"
              when "ContextualChunkTool"
                "Reading context"
              else
                status_text
              end

    Turbo::StreamsChannel.broadcast_replace_to(
      [@conversation, "messages"],
      target: "assistant_status_indicator",
      partial: "assistant_conversations/status_indicator",
      locals: { status_message: message }
    )
  end

  def assistant_message_role(role)
    return role unless gemini_llm?

    case role
    when "assistant"
      "model"
    when "tool"
      "function"
    else
      role
    end
  end

  def gemini_llm?
    @llm.is_a?(Langchain::LLM::GoogleGemini) || @llm.is_a?(Langchain::LLM::GoogleVertexAI)
  end

  def ask(prompt, filters)
    Rails.logger.info("ASKING:::: #{prompt}")
    Langchain.logger = Rails.logger
    Langchain.logger.level = Logger::INFO
    Rails.logger.info("[FILTERS] FILTERS ARE: #{filters}")
    # 1. Persist user question
    @current_message = @conversation.assistant_messages.order(:created_at).last

    # Broadcast initial "thinking" status
    broadcast_status("Analyzing your question...")

    message_history = @conversation.assistant_messages.order(:created_at).where.not(id: @current_message.id)

    # Collect only clean conversational pairs for history
    # Skip any assistant message that has tool calls, AND skip the tool responses that follow
    # However, this means multi-turn tool context is lost on replay.
    # The cleaner long-term fix is to store and restore the full serialised message thread properly, but this stops the crash.
    skip_next = false
    message_history.each do |msg|
      next if skip_next && %w[tool function].include?(msg.role)

      skip_next = false

      # Skip assistant messages that are just tool invocations
      if %w[assistant model].include?(msg.role) && msg.message.blank?
        skip_next = true # also skip the following tool response(s)
        next
      end

      next if %w[tool function].include?(msg.role)

      @assistant.add_message(
        content: msg.message,
        role: assistant_message_role(msg.role)
      )
    end

    user_prompt = <<~PROMPT
      <question>
        #{prompt}
      </question>
      <filters>
        #{filters}
      </filters>
    PROMPT

    # Rails.logger.info("******************* ||||||||||||||||| LLM Prompt:  #{message_history} #{user_prompt}")
    @assistant.add_message(content: user_prompt)

    # Broadcast that we're now processing/synthesizing
    broadcast_status("Thinking through the answer...")

    response = @assistant.run(auto_tool_execution: true)

    Rails.logger.info("[RagAssistantResponse] Raw Response - #{response}")

    full_answer = response.reverse.find { |message| message.llm? && message.content.present? }&.content.to_s

    if full_answer.blank?
      Rails.logger.error("[RagAssistant] Assistant completed without a text answer. state=#{@assistant.state} tail=#{response.last(5).map do |message|
        { role: message.role, standard_role: message.standard_role, has_tool_calls: message.tool_calls.any?, content_length: message.content.to_s.length }
      end}")
      full_answer = "I could not generate a final answer for that request. Please try again."
    end

    # Rails.logger.info("******************* ||||||||||||||||| LLM Answer: #{full_answer}")

    @current_message.update!(message: full_answer)

    # Broadcast the final, polished HTML to replace the streaming version
    Turbo::StreamsChannel.broadcast_replace_to(
      [@conversation, "messages"],
      target: ActionView::RecordIdentifier.dom_id(@current_message),
      partial: "assistant_conversations/message",
      locals: { message: @current_message }
    )

    full_answer
  rescue StandardError => e
    Rails.logger.error("Error in assistant ask job - #{e.class}: #{e.message}")
    Rails.logger.error(e.backtrace.take(15).join("\n")) # Logs the top 15 lines of the trace

    @current_message&.update!(message: "There was an error processing your request. Please contact the team.")
    Turbo::StreamsChannel.broadcast_replace_to(
      [@conversation, "messages"],
      target: ActionView::RecordIdentifier.dom_id(@current_message),
      partial: "assistant_conversations/message",
      locals: { message: @current_message }
    )
    "There was an error processing your request. Please contact the team."
  ensure
    # Always clear so the next ask starts fresh
    @current_message = nil
  end

  def build_index(respondents, task_id:)
    build_gemini_index(respondents, task_id:)
  end

  # --------------------------------------------------------------------------
  # INTERNALS
  # --------------------------------------------------------------------------
  def build_openai_index(respondents, task_id:)
    total = respondents.count
    task = SystemTask.find(task_id) if task_id
    tokenizer = Tiktoken.encoding_for_model("text-embedding-003") # Check model name
    idx = 0
    progress = 0
    respondents.find_each do |resp|
      chunks =
        if resp.transcript_sections.exists?
          # ===== normal per-section path =====================================
          resp.transcript_sections.flat_map do |sec|
            result = tokenize_section(sec, tokenizer)
            sleep(1.second)
            result
          end
        else
          # ===== fallback: one big transcript string =========================
          tokenize_full_transcript(resp, tokenizer)
        end
      @client.add_texts(
        texts: chunks.map(&:text),
        ids: chunks.map(&:id),
        metadatas: chunks.map(&:metadata)
      )
      idx += 1
      progress = ((idx / total.to_f) * 100).round
      task&.set_progress!(org_id: task.user.organization_id, progress:, metadata: { project_id: @project_id,
                                                                                    respondent_category_id: @respondent_category_id,
                                                                                    raggable_type: @raggable_type })
    end
    task&.set_progress!(org_id: task.user.organization_id, progress: 100, metadata: { project_id: @project_id,
                                                                                      respondent_category_id: @respondent_category_id,
                                                                                      raggable_type: @raggable_type })
  end

  def tokenize_section(section, tok)
    toks = tok.encode(section.content)
    start = 0
    idx = 0
    out = []

    while start < toks.length
      # Determine the ideal end point
      finish = [start + MAX_TOKENS, toks.length].min

      # This is the new, crucial part:
      # We will shrink the slice from the right until it's a valid decodable chunk.
      slice = nil
      decoded_text = nil

      # Work backwards from the ideal 'finish' until we find a valid UTF-8 sequence
      current_finish = finish
      while current_finish > start
        begin
          slice = toks[start...current_finish]
          decoded_text = tok.decode(slice)
          # If we get here, decoding was successful! Break the inner loop.
          break
        rescue Tiktoken::UnicodeError # Catches UTF-8 decoding errors in tiktoken-ruby
          # Decoding failed. Let's try a shorter slice.
          current_finish -= 1
        end
      end

      # If decoded_text is still nil, it means no valid slice was found (edge case).
      # You might want to handle this, but usually it will succeed.
      break unless decoded_text

      out << Chunk.new(
        decoded_text, # Use the successfully decoded text
        "#{section.id}-#{idx}",
        {
          section_id: section.id,
          respondent_id: section.respondent_id,
          master_section_id: section.master_section_id,
          chunk_idx: idx,
          tags: "|--|#{section.respondent.tags.pluck(:id).join('|--|')}|--|",
          start_token: start
        }
      )

      idx += 1

      # The new start must be based on the actual finish point of the valid slice
      new_start = current_finish - OVERLAP

      # Prevent an infinite loop if overlap is too large or chunk is too small
      break if new_start <= start

      start = new_start
    end

    out
  end

  def build_gemini_index(respondents, task_id:)
    total = respondents.count
    task = SystemTask.find(task_id) if task_id
    idx = 0
    progress = 0

    text_splitter = Baran::RecursiveCharacterTextSplitter.new(
      chunk_size: MAX_TOKENS,
      chunk_overlap: OVERLAP,
      separators: ["\n\n", "\n", " ", ""] # Default separators
    )

    splitter = text_splitter

    respondents.find_each do |resp|
      chunks = if resp.transcript_sections.exists?
                 resp.transcript_sections.flat_map do |sec|
                   section_metadata = {
                     section_id: sec.id,
                     respondent_id: resp.id,
                     tags: resp.tags.pluck(:name).collect { |x| "|--|#{x}|--|" }.join(" ")
                   }

                   create_chunks(
                     text: sec.content.to_s,
                     base_id: sec.id,
                     metadata: section_metadata,
                     splitter:
                   )
                 end
               else
                 fallback_metadata = {
                   section_id: "",
                   respondent_id: resp.id,
                   tags: resp.tags.pluck(:name).collect { |x| "|--|#{x}|--|" }.join(" ")
                 }

                 create_chunks(
                   text: resp.full_text.to_s,
                   base_id: resp.id,
                   metadata: fallback_metadata,
                   splitter:
                 )
               end

      Rails.logger.info("[RagAssistant] Adding #{chunks.size} chunks to Chroma for respondent #{resp.id}")
      chunks.each_slice(50).with_index do |batch, batch_idx|
        Rails.logger.info("[RagAssistant] Adding Chroma batch #{batch_idx + 1}, size=#{batch.size}, respondent_id=#{resp.id}")
        add_chunks_to_chroma(batch)
      end

      idx += 1
      Rails.logger.info "On #{idx} / #{total}"
      progress = ((idx / total.to_f) * 100).round
      task&.set_progress!(org_id: task.user.organization_id, progress:, metadata: { project_id: @project_id,
                                                                                    respondent_category_id: @respondent_category_id,
                                                                                    raggable_type: @raggable_type })
    end
    task&.set_progress!(org_id: task.user.organization_id, progress: 100, metadata: { project_id: @project_id,
                                                                                      respondent_category_id: @respondent_category_id,
                                                                                      raggable_type: @raggable_type })
  end

  def create_chunks(text:, base_id:, metadata:, splitter:)
    # This returns an array of hashes: [{ cursor: 0, text: "...", metadata: {...} }, ...]
    chunk_hashes = splitter.chunks(text.to_s, metadata:)

    # Map these hashes back to your `Chunk` struct
    chunk_hashes.map.with_index do |chunk_hash, idx|
      Chunk.new(
        chunk_hash[:text],                # The chunked text
        "#{base_id}-#{idx}",              # Your custom ID
        chunk_hash[:metadata].merge(chunk_idx: idx) # Merged metadata
      )
    end
  end

  def add_chunks_to_chroma(chunks)
    texts = chunks.map(&:text)
    embeddings = @llm.embed_texts(texts:)

    chroma_embeddings = chunks.map.with_index do |chunk, idx|
      Chroma::Resources::Embedding.new(
        id: chunk.id.to_s,
        embedding: embeddings[idx],
        metadata: chunk.metadata,
        document: chunk.text
      )
    end

    Chroma::Resources::Collection.get(@client.index_name).add(chroma_embeddings)
  end

  # ---- chunk an entire respondent transcript --------------------------------
  def tokenize_full_transcript(resp, tok)
    tokenize_string(
      text: resp.full_text.to_s,
      base_id: resp.id,
      metadata: {
        section_id: "",
        respondent_id: resp.id,
        tags: "|--|#{resp.tags.pluck(:id).join('|--|')}|--|"
      },
      tokenizer: tok
    )
  end

  def tokenize_string(text:, base_id:, metadata:, tokenizer:)
    tokens = tokenizer.encode(text)
    chunks = []
    idx    = 0
    start  = 0

    while start < tokens.length
      finish = [start + MAX_TOKENS, tokens.length].min
      slice = nil
      decoded_text = nil
      current_finish = finish

      # Work backwards from the ideal 'finish' until a valid UTF-8 sequence is found
      while current_finish > start
        begin
          slice = tokens[start...current_finish]
          decoded_text = tokenizer.decode(slice)
          break # Decoding was successful, exit inner loop
        rescue Tiktoken::UnicodeError # Catches UTF-8 decoding errors in tiktoken-ruby
          # Decoding failed, try a shorter slice
          current_finish -= 1
        end
      end

      # If decoding failed for all slices (unlikely edge case), break
      break unless decoded_text

      chunks << Chunk.new(
        decoded_text,
        "#{base_id}-#{idx}",
        metadata.merge(chunk_idx: idx, start_token: start)
      )

      idx += 1
      new_start = current_finish - OVERLAP
      break if new_start <= start

      start = new_start
    end

    chunks
  end
end