# typed: strict

require "arel/visitors/to_sql"
class Arel::Visitors::ToSql
  extend T::Sig
  prepend(
    Module.new do
      extend T::Sig
      sig { params(o: T.untyped, collector: T.untyped).returns(T.untyped) }
      def visit_Arel_Attributes_Attribute(o, collector)
        join_name = o.relation.table_alias || o.relation.name
        ar_table = o.relation.instance_variable_get("@klass")
        if ar_table &&
             (
               attribute_def =
                 AttrJsonRecordAliases::ImplHelper.get_json_attr_def(
                   ar_table,
                   o.name,
                 )
             )
          attr_type =
            T.cast(
              attribute_def.type.type,
              T.any(Symbol, ActiveModel::Type::Value),
            )
          attr_type_cast = ar_table.json_attribute_type_cast(attr_type)
          column = "#{quote_column_name("json_attributes")}->>'#{o.name}'"
          str = "((#{quote_table_name(join_name)}.#{column})#{attr_type_cast})"
          return collector << str
        end

        collector << quote_table_name(join_name) << "." <<
          quote_column_name(o.name)
      end
    end,
  )
end

require "active_record/relation/query_methods"
require "active_record/relation/calculations"
module ActiveRecord::QueryMethods
  extend T::Sig
  prepend(
    Module.new do
      extend T::Sig

      sig { params(field: T.untyped).returns(T.untyped) }
      def arel_column(field)
        klass = T.cast(T.unsafe(self).klass, T.class_of(ActiveRecord::Base))
        is_order_field =
          field.instance_variable_get("@_is_order_field") || false
        if attribute_def =
             AttrJsonRecordAliases::ImplHelper.get_json_attr_def(klass, field)
          attr_type_cast =
            T.unsafe(klass).json_attribute_type_cast(attribute_def.type.type)
          adapter_class = T.unsafe(self).adapter_class
          column =
            "#{adapter_class.quote_column_name("json_attributes")}->>'#{field}'"
          column_expr =
            "(#{adapter_class.quote_table_name(klass.table_name)}.#{column})#{attr_type_cast}"
          if is_order_field
            Arel::Nodes::SqlLiteral.new(column_expr)
          else
            Arel::Nodes::SqlLiteral.new(
              "#{column_expr} AS #{adapter_class.quote_column_name(field)}",
            )
          end
        else
          super
        end
      end

      sig { params(field: T.untyped).returns(T.untyped) }
      def order_column(field)
        field = field.dup
        field.instance_variable_set(:"@_is_order_field", true)
        ret = super(field)
        field.instance_variable_set(:"@_is_order_field", nil)
        ret
      end
    end,
  )
end

module ActiveRecord::Calculations
  extend T::Sig
  prepend(
    Module.new do
      extend T::Sig
      sig { params(field: T.untyped).returns(T.untyped) }
      def aggregate_column(field)
        return super(field) if field == :all
        field = field.to_s if field.is_a?(Symbol)
        field = field.dup if field.frozen?
        field.instance_variable_set(:"@_is_order_field", true)
        ret = super(field)
        field.instance_variable_set(:"@_is_order_field", nil)
        ret
      end
    end,
  )
end

module AttrJsonRecordAliases
  extend T::Sig
  extend T::Helpers
  extend ActiveSupport::Concern
  abstract!
  requires_ancestor { ActiveRecord::Base }

  included do
    include AttrJson::Record
    original_write_attribute = instance_method(:write_attribute)
    define_method(:write_attribute) do |name, value|
      T.bind(self, AttrJsonRecordAliases)
      klass = T.cast(self.class, T.class_of(ActiveRecord::Base))
      ret = original_write_attribute.bind(self).call(name, value)
      if attribute_def = ImplHelper.get_json_attr_def(klass, name)
        public_send(attribute_def.container_attribute)[
          attribute_def.store_key
        ] = read_attribute(name)
      end

      ret
    end
  end

  sig { params(attributes: T::Hash[Symbol, T.untyped]).returns(T.untyped) }
  def update_json_columns(attributes)
    klass = T.cast(self.class, T.class_of(ActiveRecord::Base))
    updated_json = klass.connection.quote(JSON.dump(attributes))
    klass.connection.execute(<<~SQL)
        UPDATE #{klass.quoted_table_name}
        SET json_attributes = json_attributes || #{updated_json}
        WHERE #{_query_constraints_hash.map { |k, v| "#{k} = #{v}" }.join(" AND ")}
      SQL
  end

  module ImplHelper
    extend T::Sig

    sig do
      params(klass: T.class_of(ActiveRecord::Base), field: T.untyped).returns(
        T.nilable(AttrJson::AttributeDefinition),
      )
    end
    def self.get_json_attr_def(klass, field)
      if klass < AttrJsonRecordAliases
        registry =
          T.cast(
            T.unsafe(klass).attr_json_registry,
            AttrJson::AttributeDefinition::Registry,
          )
        registry[field.to_sym]
      else
        nil
      end
    end
  end

  module ClassMethods
    extend T::Sig
    extend T::Helpers
    requires_ancestor { T.class_of(ActiveRecord::Base) }

    sig do
      params(type: T.any(Symbol, ActiveModel::Type::Value)).returns(String)
    end
    def json_attribute_type_cast(type)
      case type
      when :integer
        "::integer"
      when :string
        "::text"
      when :time_no_zone
        "::timestamp"
      else
        ""
      end
    end
  end

  mixes_in_class_methods(ClassMethods)
end
