require 'apartment/elevators/generic'

begin
  require "apartment-sidekiq"
rescue LoadError
end

require "apartment/adapters/postgresql_adapter"

module Apartment
  SHARD_PREFIXES = ["SHARD_DB", "HEROKU_POSTGRESQL"]

  def self.shard_configurations
    $shard_configurations ||= begin
      shard_to_env = {}

      ENV.keys.each do |k|
        m = /^(#{SHARD_PREFIXES.join("|")})_(\w+)_URL$/.match(k)
        next unless m

        url = ENV[k]
        shard_to_env[m[2].downcase] = ActiveRecord::Base.configurations.resolve(url).configuration_hash
      end

      shard_to_env.freeze unless Rails.env.test?

      shard_to_env
    end
  end

  module Tenant
    self.singleton_class.send(:alias_method, :original_postgresql_adapter, :postgresql_adapter)

    def self.postgresql_adapter(config)
      if Apartment.with_multi_server_setup
        adapter = Adapters::PostgresMultiDBSchemaAdapter
        adapter.new(config)
      else
        original_postgresql_adapter(config)
      end
    end

    def self.split_tenant(tenant)
      bits = tenant.split(":", 2)
      bits.unshift(nil) if bits.length == 1
      bits[0] = "default" if bits[0].nil? || bits[0].empty?
      bits
    end
  end

  module Adapters
    class PostgresMultiDBSchemaAdapter < Apartment::Adapters::PostgresqlSchemaAdapter
      def initialize(*args, **kwargs)
        super
        @excluded_model_set = Set.new(Apartment.excluded_models)
      end

      def process_excluded_model(excluded_model)
        @excluded_model_set << excluded_model
        super
      end

      def is_excluded_model?(model)
        @excluded_model_set.include?(model.to_s)
      end

      def db_connection_config(tenant)
        shard, schema = Tenant.split_tenant(tenant)
        if shard == "default"
          @config
        else
          Apartment.shard_configurations[shard.downcase]
        end
      end

      def drop_command(conn, tenant)
        shard, schema = Tenant.split_tenant(tenant)
        conn.execute(%(DROP SCHEMA "#{schema}" CASCADE))
      end

      def tenant_exists?(tenant)
        return true unless Apartment.tenant_presence_check
        shard, schema = Tenant.split_tenant(tenant)

        Apartment.connection.schema_exists?(schema)
      end

      def create_tenant_command(conn, tenant)
        shard, schema = Tenant.split_tenant(tenant)
        # NOTE: This was causing some tests to fail because of the database strategy for rspec
        if conn.open_transactions.positive?
          conn.execute(%(CREATE SCHEMA "#{schema}")).inspect
        else
          schema = %(BEGIN;
          CREATE SCHEMA "#{schema}";
          COMMIT;)

          conn.execute(schema)
        end
      rescue *rescuable_exceptions => e
        rollback_transaction(conn)
        raise e
      end

      def connect_to_new(tenant = nil)
        return reset if tenant.nil?

        current_tenant = @current
        tenants_array = tenant.is_a?(Array) ? tenant.map(&:to_s) : tenant.to_s
        tenant_schemas = map_to_schemas(Array(tenants_array))
        query_cache_enabled = ActiveRecord::Base.connection.query_cache_enabled

        @current = tenants_array

        raise ActiveRecord::StatementInvalid, "PandaPal/Apartment Mutli-DB support does not currently support DB roles" if ActiveRecord::Base.current_role != ActiveRecord::Base.default_role

        unless ActiveRecord::Base.connected?
          Apartment.establish_connection multi_tenantify(tenant, false)
          Apartment.connection.verify!
        end

        Apartment.connection.enable_query_cache! if query_cache_enabled

        raise ActiveRecord::StatementInvalid, "Could not find schema for tenant #{tenant} (#{tenant_schemas.inspect})" unless schema_exists?(tenant_schemas)

        Apartment.connection.schema_search_path = full_search_path
      rescue *rescuable_exceptions => e
        @current = current_tenant
        raise_schema_connect_to_new(tenant, e)
      end

      protected

      def persistent_schemas
        map_to_schemas(super)
      end

      def map_to_schemas(tenants)
        all_shard = nil
        tenants.map do |schema|
          shard, schema = Tenant.split_tenant(schema)
          all_shard ||= shard
          raise "Cannot mix shards in persistent_schemas" if shard != all_shard
          schema
        end
      end
    end
  end

  module CorePatches
    module ActiveRecord
      module FutureResult
        # called in the original thread that knows the tenant
        def initialize(_pool, *_args, **_kwargs)
          @tenant = Apartment::Tenant.current
          super
        end

        # called in the new thread with a connection that needs switching
        def exec_query(_conn, *_args, **_kwargs)
          Apartment::Tenant.switch!(@tenant) unless Apartment::Tenant.current == @tenant
          super
        end
      end

      module Base
        extend ActiveSupport::Concern

        included do
          self.singleton_class.send(:alias_method, :pre_apartment_current_shard, :current_shard)

          def self.current_shard
            # This implementation is definitely a hack, but it should be fairly compatible. If you need to leverage
            #   Rails' sharding natively - you should just need to add models to the excluded_models list in Apartment
            #   to effectively disable this patch for that model.
            if (adapter = Thread.current[:apartment_adapter]) && adapter.is_a?(Apartment::Adapters::PostgresMultiDBSchemaAdapter) && !adapter.is_excluded_model?(self)
              shard, schema = Apartment::Tenant.split_tenant(adapter.current)
              return "apt:#{shard}" unless shard == "default"
            end

            pre_apartment_current_shard
          end
        end
      end
    end
  end

  # Fix Apartment's lack of support for load_async
  ::ActiveRecord::FutureResult.prepend CorePatches::ActiveRecord::FutureResult

  # Hack ActiveRecord shards/connection_pooling to support our multi-DB approach
  ::ActiveRecord::Base.include CorePatches::ActiveRecord::Base
end

Apartment.configure do |config|
  config.excluded_models ||= []
  config.excluded_models |= ['PandaPal::Organization', 'PandaPal::Session']

  config.with_multi_server_setup = true unless Rails.env.test?

  config.tenant_names = lambda do
    if PandaPal::Organization < PandaPal::OrganizationConcerns::MultiDatabaseSharding
      base_config = Apartment.connection_config
      shard_configurations = Apartment.shard_configurations

      PandaPal::Organization.all.to_a.each_with_object({}) do |org, hash|
        shard = org.shard || "default"
        hash[org.tenant_name] = shard == "default" ? base_config : shard_configurations[shard.downcase]
      end
    else
      PandaPal::Organization.pluck(:name)
    end
  end
end

Rails.application.config.middleware.use Apartment::Elevators::Generic, lambda { |request|
  if match = request.path.match(/\/(?:orgs?|organizations?)\/(\d+)/)
    PandaPal::Organization.find_by(id: match[1]).try(:name)
  elsif match = request.path.match(/\/(?:orgs?|organizations?|o)\/(\w+)/)
    PandaPal::Organization.find_by(name: match[1]).try(:name)
  elsif request.path.start_with?('/rails/active_storage/blobs/')
    PandaPal::Organization.find_by(id: request.params['organization_id']).try(:name)
  end
}

module PandaPal::Plugins::ApartmentCache
  private

  if Rails.version >= '5.0'
    def normalize_key(key, options)
      "tenant:#{Apartment::Tenant.current}/#{super}"
    end
  else
    def namespaced_key(*args, **kwargs)
      "tenant:#{Apartment::Tenant.current}/#{super}"
    end
  end
end
ActiveSupport::Cache::Store.send(:prepend, PandaPal::Plugins::ApartmentCache)

if defined?(ActionCable)
  module ActionCable
    module Channel
      class Base
        def self.broadcasting_for(model)
          # Rails 5 #stream_for passes #channel_name as part of model. Rails 6 doesn't and includes it via #broadcasting_for.
          model = [channel_name, model] unless model.is_a?(Array)
          serialize_broadcasting([ Apartment::Tenant.current, model ])
        end

        def self.serialize_broadcasting(object)
          case
          when object.is_a?(Array)
            object.map { |m| serialize_broadcasting(m) }.join(":")
          when object.respond_to?(:to_gid_param)
            object.to_gid_param
          else
            object.to_param
          end
        end
      end
    end
  end

  module PandaPal::Plugins::ActionCableApartment
    module Connection
      def tenant=(name)
        @tenant = name
      end

      def tenant
        @tenant || 'public'
      end

      def dispatch_websocket_message(*args, **kwargs)
        Apartment::Tenant.switch(tenant) do
          super
        end
      end
    end
  end

  ActionCable::Connection::Base.prepend(PandaPal::Plugins::ActionCableApartment::Connection)
end

if defined?(Delayed)
  module PandaPal::Plugins
    class ApartmentDelayedJobsPlugin < ::Delayed::Plugin
      callbacks do |lifecycle|
        lifecycle.around(:enqueue) do |job, *args, &block|
          current_tenant = Apartment::Tenant.current

          #make sure enqueue on public tenant unless we are testing since delayed job is set to run immediately
          Apartment::Tenant.switch!('public') unless Rails.env.test?
          job.tenant = current_tenant
          begin
            block.call(job, *args)
          rescue Exception => e
            Rails.logger.error("Error enqueing job #{job.to_s} - #{e.backtrace}")
          ensure
            #switch back to prev tenant
            Apartment::Tenant.switch!(current_tenant)
          end
        end

        lifecycle.before(:perform) do |worker, *args, &block|
          tenant = args.first.tenant
          Apartment::Tenant.switch!(tenant) if tenant.present?
          Rails.logger.debug("Running job with tenant #{Apartment::Tenant.current}")
        end

        lifecycle.around(:invoke_job) do |job, *args, &block|
          begin
            block.call(job, *args)
          ensure
            Apartment::Tenant.switch!('public')
            Rails.logger.debug("Resetting Tenant back to: #{Apartment::Tenant.current}")
          end
        end

        lifecycle.after(:failure) do |job, *args|
          Rails.logger.error("Job failed on tenant: #{Apartment::Tenant.current}")
        end
      end
    end
  end

  Delayed::Worker.plugins << PandaPal::Plugins::ApartmentDelayedJobsPlugin
end