# encoding: utf-8

require 'singleton'
require 'one_apm/agent/database/obfuscation_helpers'
require 'one_apm/agent/database/obfuscator'
require 'one_apm/agent/database/postgres_explain_obfuscator'

module OneApm
  MYSQL_EXPLAIN_COLUMNS = [
                           "Id",
                           "Select Type",
                           "Table",
                           "Type",
                           "Possible Keys",
                           "Key",
                           "Key Length",
                           "Ref",
                           "Rows",
                           "Extra"
                          ].freeze

  module Agent
    module Database
      MAX_QUERY_LENGTH = 16384

      extend self

      def capture_query(query)
        Helper.correctly_encoded(truncate_query(query))
      end

      def truncate_query(query)
        if query.length > (MAX_QUERY_LENGTH - 4)
          query[0..MAX_QUERY_LENGTH - 4] + '...'
        else
          query
        end
      end

      def obfuscate_sql(sql)
        Obfuscator.instance.obfuscator.call(sql)
      end

      def set_sql_obfuscator(type, &block)
        Obfuscator.instance.set_sql_obfuscator(type, &block)
      end

      def record_sql_method(config_section=:transaction_tracer)
        case OneApm::Manager.config["#{config_section}.record_sql".to_sym].to_s
        when 'off'
          :off
        when 'none'
          :off
        when 'false'
          :off
        when 'raw'
          :raw
        else
          :obfuscated
        end
      end

      RECORD_FOR = [:raw, :obfuscated].freeze

      def should_record_sql?(config_section=:transaction_tracer)
        RECORD_FOR.include?(record_sql_method(config_section))
      end

      def should_collect_explain_plans?(config_section=:transaction_tracer)
        should_record_sql?(config_section) &&
          OneApm::Manager.config["#{config_section}.explain_enabled".to_sym]
      end

      def get_connection(config, &connector)
        ConnectionManager.instance.get_connection(config, &connector)
      end

      def close_connections
        ConnectionManager.instance.close_connections
      end

      def adapter_from_config(config)
        if config[:adapter]
          return config[:adapter].to_s
        elsif config[:uri] && config[:uri].to_s =~ /^jdbc:([^:]+):/
          return $1
        end
      end

      def explain_sql(sql, connection_config, &explainer)
        return nil unless sql && connection_config
        statement = sql.split(";\n")[0] # only explain the first
        explain_plan = explain_statement(statement, connection_config, &explainer)
        return explain_plan || []
      end

      SUPPORTED_ADAPTERS_FOR_EXPLAIN = %w[postgres postgresql mysql2 mysql sqlite].freeze

      def explain_statement(statement, config, &explainer)
        return unless is_select?(statement)

        if statement[-3,3] == '...'
          OneApm::Manager.logger.debug('Unable to collect explain plan for truncated query.')
          return
        end

        if parameterized?(statement)
          OneApm::Manager.logger.debug('Unable to collect explain plan for parameterized query.')
          return
        end

        adapter = adapter_from_config(config)
        if !SUPPORTED_ADAPTERS_FOR_EXPLAIN.include?(adapter)
          OneApm::Manager.logger.debug("Not collecting explain plan because an unknown connection adapter ('#{adapter}') was used.")
          return
        end

        handle_exception_in_explain do
          start = Time.now
          plan = explainer.call(config, statement)
          OneApm::Manager.record_metric("Supportability/Database/execute_explain_plan", Time.now - start)
          return process_resultset(plan, adapter) if plan
        end
      end

      def process_resultset(results, adapter)
        case adapter.to_s
        when 'postgres', 'postgresql'
          process_explain_results_postgres(results)
        when 'mysql2'
          process_explain_results_mysql2(results)
        when 'mysql'
          process_explain_results_mysql(results)
        when 'sqlite'
          process_explain_results_sqlite(results)
        end
      end

      QUERY_PLAN = 'QUERY PLAN'.freeze

      def process_explain_results_postgres(results)
        if results.is_a?(String)
          query_plan_string = results
        else
          lines = []
          results.each { |row| lines << row[QUERY_PLAN] }
          query_plan_string = lines.join("\n")
        end

        unless record_sql_method == :raw
          query_plan_string = OneApm::Agent::Database::PostgresExplainObfuscator.obfuscate(query_plan_string)
        end
        values = query_plan_string.split("\n").map { |line| [line] }

        [[QUERY_PLAN], values]
      end

      def string_explain_plan_results(results)
        [nil, [results]]
      end

      def process_explain_results_mysql(results)
        return string_explain_plan_results(results) if results.is_a?(String)
        headers = []
        values  = []
        if results.is_a?(Array)
          headers = results.first.keys
          results.each do |row|
            values << headers.map { |h| row[h] }
          end
        else
          results.each_hash do |row|
            headers = row.keys
            values << headers.map { |h| row[h] }
          end
        end
        [headers, values]
      end

      def process_explain_results_mysql2(results)
        return string_explain_plan_results(results) if results.is_a?(String)
        headers = results.fields
        values  = []
        results.each { |row| values << row }
        [headers, values]
      end

      SQLITE_EXPLAIN_COLUMNS = %w[addr opcode p1 p2 p3 p4 p5 comment]

      def process_explain_results_sqlite(results)
        return string_explain_plan_results(results) if results.is_a?(String)
        headers = SQLITE_EXPLAIN_COLUMNS
        values  = []
        results.each do |row|
          values << headers.map { |h| row[h] }
        end
        [headers, values]
      end

      def handle_exception_in_explain
        yield
      rescue => e
        begin
          # guarantees no throw from explain_sql
          OneApm::Manager.logger.error("Error getting query plan:", e)
          nil
        rescue
          # double exception. throw up your hands
          nil
        end
      end

      KNOWN_OPERATIONS = [
        'alter',
        'select',
        'update',
        'delete',
        'insert',
        'create',
        'show',
        'set',
        'exec',
        'execute',
        'call'
      ]

      SQL_COMMENT_REGEX = Regexp.new('/\*.*?\*/', Regexp::MULTILINE).freeze

      def parse_operation_from_query(sql)
        sql = sql.gsub(SQL_COMMENT_REGEX, '')
        if sql =~ /(\w+)/
          op = $1.downcase
          return op if KNOWN_OPERATIONS.include?(op)
        end
      end

      def is_select?(statement)
        parse_operation_from_query(statement) == 'select'
      end

      def parameterized?(statement)
        Obfuscator.instance.obfuscate_single_quote_literals(statement) =~ /\$\d+/
      end

      class ConnectionManager
        include Singleton

        def get_connection(config, &connector)
          @connections ||= {}

          connection = @connections[config]

          return connection if connection

          begin
            @connections[config] = connector.call(config)
          rescue => e
            OneApm::Manager.logger.error("Caught exception trying to get connection to DB for explain.", e)
            nil
          end
        end

        def close_connections
          @connections ||= {}
          @connections.values.each do |connection|
            begin
              connection.disconnect!
            rescue
            end
          end

          @connections = {}
        end
      end

      class Statement < String
        attr_accessor :adapter, :config, :explainer
      end
    end
  end
end