# encoding: utf-8

require 'rack'
require 'ipaddr'
require 'one_apm/rack/middleware_base'
require 'one_apm/rack/middleware_wrapper'

module OneApm::Rack
  class BrowserMonitoring < MiddlewareBase

    SCAN_LIMIT = 50_000
    ALREADY_INSTRUMENTED_KEY = "oneapm.browser_monitoring_already_instrumented"
    REMOTE_IP = "action_dispatch.remote_ip"

    CHARSET_RE         = /<\s*meta[^>]+charset\s*=[^>]*>/im.freeze
    X_UA_COMPATIBLE_RE = /<\s*meta[^>]+http-equiv\s*=\s*['"]x-ua-compatible['"][^>]*>/im.freeze

    def traced_call(env)
      result = @app.call(env)   # [status, headers, response]

      js_to_inject = OneApm::Manager.browser_timing_header
      if (js_to_inject != "") && should_instrument?(env, result[0], result[1])
        response_string = autoinstrument_source(result[2], result[1], js_to_inject)

        env[ALREADY_INSTRUMENTED_KEY] = true
        if response_string
          response = Rack::Response.new(response_string, result[0], result[1])
          response.finish
        else
          result
        end
      else
        result
      end
    end

    def should_instrument?(env, status, headers)
      OneApm::Manager.config[:'browser_monitoring.auto_instrument'] &&
        status == 200 &&
        !env[ALREADY_INSTRUMENTED_KEY] &&
        is_html?(headers) &&
        !is_attachment?(headers) &&
        !is_streaming?(env) && 
        ip_valid?(env)
    end

    def is_html?(headers)
      headers["Content-Type"] && headers["Content-Type"].include?("text/html")
    end

    def is_attachment?(headers)
      headers['Content-Disposition'].to_s.include?('attachment')
    end

    def is_streaming?(env)
      return false unless defined?(ActionController::Live)
      env['action_controller.instance'].class.included_modules.include?(ActionController::Live)
    end

    def ip_valid?(env)
      whitelist_ips = OneApm::Manager.config[:'browser_monitoring.whitelist_ips']
      return true if whitelist_ips.empty? || env[REMOTE_IP].to_s.empty?
      OneApm::Manager.logger.debug "Remote IP: #{env[REMOTE_IP]}"
      begin
        client = IPAddr.new(env[REMOTE_IP].to_s).to_i 
        whitelist_ips.split(",").any? do |ip|
          low_ip, high_ip = ip.split("-").map(&:strip)
          low = IPAddr.new(low_ip).to_i
          if high_ip
            high = IPAddr.new(high_ip).to_i
            (low..high) === client
          else
            low == client  
          end
        end
      rescue => e
        OneApm::Manager.logger.error "Configuration for browser_monitoring.whitelist_ips has problems:#{whitelist_ips}"
        true
      end
    end 

    def autoinstrument_source(response, headers, js_to_inject)
      begin
        source = gather_source(response)
        close_old_response(response)
        return nil unless source
        inject_js(source, headers, js_to_inject)
      rescue => e
        OneApm::Manager.logger.debug "Skipping RUM instrumentation on exception.", e
        nil
      end
    end

    def inject_js source, headers, js_to_inject
      position = OneApm::Manager.config[:'browser_monitoring.position']
      insertion_index = (position.empty? || position.to_sym != :footer) ? header_index(source) : footer_index(source)
      if insertion_index
        source = source[0...insertion_index] \
          << js_to_inject \
          << source[insertion_index..-1]
        if headers['Content-Length']
          headers['Content-Length'] = calculate_content_length(source).to_s
        end
      else
        OneApm::Manager.logger.debug "Skipping RUM instrumentation. Could not properly determine location to inject script."
      end
      source
    end

    def header_index source
      beginning_of_source = source[0..SCAN_LIMIT]
      if body_start = find_body_start(beginning_of_source)
        meta_tag_positions = [
          find_x_ua_compatible_position(beginning_of_source),
          find_charset_position(beginning_of_source)
        ].compact

        if !meta_tag_positions.empty?
          insertion_index = meta_tag_positions.max
        else
          insertion_index = find_end_of_head_open(beginning_of_source) || body_start
        end
        insertion_index
      else
        msg = "Skipping RUM instrumentation. Unable to find <body> tag in first #{SCAN_LIMIT} bytes of document."
        OneApm::Manager.logger.log_once(:warn, :rum_insertion_failure, msg)
        OneApm::Manager.logger.debug(msg)
        nil
      end
    end

    def footer_index source
      if body_end = find_body_end(source)
        body_end
      else
        msg = "Skipping RUM instrumentation. Unable to find </body> tag from document."
        OneApm::Manager.logger.log_once(:warn, :rum_insertion_failure, msg)
        OneApm::Manager.logger.debug(msg)
        nil
      end
    end

    def gather_source(response)
      source = nil
      response.each {|fragment| source ? (source << fragment.to_s) : (source = fragment.to_s)}
      source
    end

    # Per "The Response > The Body" section of Rack spec, we should close
    # if our response is able. http://rack.rubyforge.org/doc/SPEC.html
    def close_old_response(response)
      if response.respond_to?(:close)
        response.close
      end
    end

    def find_body_start(beginning_of_source)
      beginning_of_source.index("<body")
    end

    def find_body_end source
      source.rindex("</body>")
    end

    def find_x_ua_compatible_position(beginning_of_source)
      match = X_UA_COMPATIBLE_RE.match(beginning_of_source)
      match.end(0) if match
    end

    def find_charset_position(beginning_of_source)
      match = CHARSET_RE.match(beginning_of_source)
      match.end(0) if match
    end

    def find_end_of_head_open(beginning_of_source)
      head_open = beginning_of_source.index("<head")
      beginning_of_source.index(">", head_open) + 1 if head_open
    end

    # String does not respond to 'bytesize' in 1.8.6. Fortunately String#length
    # returns bytes rather than characters in 1.8.6 so we can use that instead.
    def calculate_content_length(source)
      if source.respond_to?(:bytesize)
        source.bytesize
      else
        source.length
      end
    end
  end
end