module BuilderApm
  module Models
    class Instrumenter
      def self.start
        new.subscribe_to_notifications
      end

      def subscribe_to_notifications
        ActiveSupport::Notifications.subscribe('sql.active_record') do |*args| 
          starttime = Time.now.to_f * 1000
          handle_sql_active_record(*args) 
          duration = (Time.now.to_f * 1000) - starttime
          
          Thread.current[:db_tracing] = (Thread.current[:db_tracing] ||= 0) + duration
        end
        ActiveSupport::Notifications.subscribe('instantiation.active_record') do |*args| 
          starttime = Time.now.to_f * 1000
          handle_instantiation_active_record(*args) 
          duration = (Time.now.to_f * 1000) - starttime
          
          Thread.current[:db_tracing] = (Thread.current[:db_tracing] ||= 0) + duration
        end
      end

      private

      def handle_sql_active_record(*args)
        data = args.extract_options!
        event = ActiveSupport::Notifications::Event.new(*args, data)
        name = event.payload[:name]
        return if name == "SCHEMA" || Thread.current[:request_id].nil?

        triggering_line = determine_triggering_line(caller)
        Thread.current[:db_runtime] += event.duration
        sql_query_data = build_sql_query_data(event, triggering_line)
        store_sql_query_data(sql_query_data)
      end

      def handle_instantiation_active_record(*args)
        data = args.extract_options!
        event = ActiveSupport::Notifications::Event.new(*args, data)

        update_last_sql_query_data_with_instantiation_info(event)
      end

      def determine_triggering_line(call_stack)
        app_stack = call_stack.select { |line| line.include?(Rails.root.to_s) }
        app_stack.first.to_s.gsub(Rails.root.to_s, '')
      end

      def build_sql_query_data(event, triggering_line)
        {
          request_id: Thread.current[:request_id],
          sql_id: SecureRandom.uuid,
          sql: event.payload[:sql],
          params: event.payload[:binds].map { |a| a.value },
          triggering_line: triggering_line,
          name: event.payload[:name],
          cached: event.payload[:cached] || false,
          start_time: event.time, 
          end_time: event.end, 
          duration: event.duration,
          record_count: 0,
          class_name: ''
        }
      end

      def store_sql_query_data(sql_query_data)
      Thread.current[:sql_event_id] = sql_query_data[:sql_id]
      
      # Create the stack if it doesn't exist yet
      stack = (Thread.current[:stack] ||= [])
      
      if stack&.any?
        stack.last[:sql_events].push(sql_query_data)
      else
        stack.push({sql_events: [sql_query_data], children: []})
      end
      # Do the N+1 check if it wasn't done yet
      if BuilderApm.configuration.enable_n_plus_one_profiler && Thread.current[:has_n_plus_one] == false
        start_time = Time.now.to_f * 1000
        perform_n_plus_one_check(stack)
        duration = (Time.now.to_f * 1000) - start_time
        Thread.current[:n_plus_one_duration] ||= 0
        Thread.current[:n_plus_one_duration] += duration
      end
    end

      def update_last_sql_query_data_with_instantiation_info(event)
        stack = Thread.current[:stack]
        request_id = Thread.current[:request_id]
      
        return if stack.nil? || stack.empty? || request_id.nil?
      
        begin
          last_sql = stack.last[:sql_events].pop
          return unless last_sql && last_sql[:sql_id]
      
          if last_sql[:sql_id] == Thread.current[:sql_event_id]
            last_sql[:record_count] = event.payload[:record_count]
            last_sql[:class_name] = event.payload[:class_name]
          end
          Thread.current[:stack].last[:sql_events].push(last_sql)
        ensure
          Thread.current[:sql_event_id] = nil
        end
      end 
      
      # Add this method to perform the N+1 check
      def perform_n_plus_one_check(stack)
        sql_queries = stack.map { |frame| frame[:sql_events] }.flatten

        # Group queries count by table and triggering_line
        queries_count = Hash.new { |h, k| h[k] = {count: 0, indices: []} }
        sql_queries.each_with_index do |query, index|
          match = query[:sql].match(/FROM ['`"]*([^ '`"]+)['`"]*/i)
          if match
            table_and_line = "#{match[1]}|#{query[:triggering_line]}"
            queries_count[table_and_line][:count] += 1
            queries_count[table_and_line][:indices].push(index)
          end
        end

        # If any N+1 issue is found, set 'has_n_plus_one' to true and return
        queries_count.each do |_, value|
          if value[:count] > 2
            Thread.current[:has_n_plus_one] = true
            return
          end
        end
      end
    end
  end
end