# frozen_string_literal: true

require 'active_support/concern'

module RubyLLM
  module ActiveRecord
    # Methods mixed into chat models.
    module ChatMethods
      extend ActiveSupport::Concern

      included do
        before_save :resolve_model_from_strings
      end

      attr_accessor :assume_model_exists, :context

      def model=(value)
        @model_string = value if value.is_a?(String)
        return if value.is_a?(String)

        if self.class.model_association_name == :model
          super
        else
          self.model_association = value
        end
      end

      def model_id=(value)
        @model_string = value
      end

      def model_id
        model_association&.model_id
      end

      def provider=(value)
        @provider_string = value
      end

      def provider
        model_association&.provider
      end

      private

      def resolve_model_from_strings # rubocop:disable Metrics/PerceivedComplexity
        config = context&.config || RubyLLM.config
        @model_string ||= config.default_model unless model_association
        return unless @model_string

        model_info, _provider = Models.resolve(
          @model_string,
          provider: @provider_string,
          assume_exists: assume_model_exists || false,
          config: config
        )

        model_class = self.class.model_class.constantize
        model_record = model_class.find_or_create_by!(
          model_id: model_info.id,
          provider: model_info.provider
        ) do |m|
          m.name = model_info.name || model_info.id
          m.family = model_info.family
          m.context_window = model_info.context_window
          m.max_output_tokens = model_info.max_output_tokens
          m.capabilities = model_info.capabilities || []
          m.modalities = model_info.modalities.to_h
          m.pricing = model_info.pricing.to_h
          m.metadata = model_info.metadata || {}
        end

        self.model_association = model_record
        @model_string = nil
        @provider_string = nil
      end

      public

      def to_llm
        model_record = model_association
        @chat ||= (context || RubyLLM).chat(
          model: model_record.model_id,
          provider: model_record.provider.to_sym,
          assume_model_exists: assume_model_exists || false
        )
        @chat.reset_messages!

        ordered_messages = order_messages_for_llm(messages_association.to_a)
        ordered_messages.each do |msg|
          @chat.add_message(msg.to_llm)
        end
        reapply_runtime_instructions(@chat)

        setup_persistence_callbacks
      end

      def with_instructions(instructions, append: false, replace: nil)
        append = append_instructions?(append:, replace:)
        persist_system_instruction(instructions, append:)

        to_llm.with_instructions(instructions, append:, replace:)
        self
      end

      def with_runtime_instructions(instructions, append: false, replace: nil)
        append = append_instructions?(append:, replace:)
        store_runtime_instruction(instructions, append:)

        to_llm.with_instructions(instructions, append:, replace:)
        self
      end

      def with_tool(...)
        to_llm.with_tool(...)
        self
      end

      def with_tools(...)
        to_llm.with_tools(...)
        self
      end

      def with_model(model_name, provider: nil, assume_exists: false)
        self.model = model_name
        self.provider = provider if provider
        self.assume_model_exists = assume_exists
        resolve_model_from_strings
        save!
        to_llm.with_model(model_association.model_id, provider: model_association.provider.to_sym, assume_exists:)
        self
      end

      def with_temperature(...)
        to_llm.with_temperature(...)
        self
      end

      def with_thinking(...)
        to_llm.with_thinking(...)
        self
      end

      def with_params(...)
        to_llm.with_params(...)
        self
      end

      def with_headers(...)
        to_llm.with_headers(...)
        self
      end

      def with_schema(...)
        to_llm.with_schema(...)
        self
      end

      def on_new_message(&)
        to_llm.on_new_message(&)
        self
      end

      def on_end_message(&)
        to_llm.on_end_message(&)
        self
      end

      def before_message(...)
        to_llm.before_message(...)
        self
      end

      def after_message(...)
        to_llm.after_message(...)
        self
      end

      def before_tool_call(...)
        to_llm.before_tool_call(...)
        self
      end

      def after_tool_result(...)
        to_llm.after_tool_result(...)
        self
      end

      def on_tool_call(...)
        to_llm.on_tool_call(...)
        self
      end

      def on_tool_result(...)
        to_llm.on_tool_result(...)
        self
      end

      def add_message(message_or_attributes)
        llm_message = message_or_attributes.is_a?(RubyLLM::Message) ? message_or_attributes : RubyLLM::Message.new(message_or_attributes)
        content_text, attachments, content_raw = prepare_content_for_storage(llm_message.content)

        attrs = { role: llm_message.role, content: content_text }
        parent_tool_call_assoc = messages_association.klass.reflect_on_association(:parent_tool_call)
        if parent_tool_call_assoc && llm_message.tool_call_id
          tool_call_id = find_tool_call_id(llm_message.tool_call_id)
          attrs[parent_tool_call_assoc.foreign_key] = tool_call_id if tool_call_id
        end

        message_record = messages_association.create!(attrs)
        message_record.update!(content_raw:) if message_record.respond_to?(:content_raw=)

        persist_content(message_record, attachments) if attachments.present?
        persist_tool_calls(llm_message.tool_calls, message_record:) if llm_message.tool_calls.present?

        message_record
      end

      def cost
        RubyLLM::Cost.aggregate(messages_association.map(&:cost))
      end

      def create_user_message(content, with: nil)
        add_message(role: :user, content: build_content(content, with))
      end

      def ask(message = nil, with: nil, &)
        add_message(role: :user, content: build_content(message, with))
        complete(&)
      end

      alias say ask

      def complete(...)
        to_llm.complete(...)
      rescue RubyLLM::Error => e
        cleanup_failed_messages if @message&.persisted? && @message.content.blank?
        cleanup_orphaned_tool_results
        raise e
      end

      private

      def cleanup_failed_messages
        RubyLLM.logger.warn "RubyLLM: API call failed, destroying message: #{@message.id}"
        @message.destroy
      end

      def cleanup_orphaned_tool_results # rubocop:disable Metrics/PerceivedComplexity
        messages_association.reload
        last = messages_association.order(:id).last

        return unless last&.tool_call? || last&.tool_result?

        if last.tool_call?
          last.destroy
        elsif last.tool_result?
          tool_call_message = last.parent_tool_call.message_association
          expected_results = tool_call_message.tool_calls_association.pluck(:id)
          fk_column = tool_call_message.class.reflections['tool_results'].foreign_key
          actual_results = tool_call_message.tool_results.pluck(fk_column)

          if expected_results.sort != actual_results.sort
            tool_call_message.tool_results.each(&:destroy)
            tool_call_message.destroy
          end
        end
      end

      def setup_persistence_callbacks
        return @chat if @chat.instance_variable_get(:@_persistence_callbacks_setup)

        @chat.before_message { persist_new_message }
        @chat.after_message { |msg| persist_message_completion(msg) }

        @chat.instance_variable_set(:@_persistence_callbacks_setup, true)
        @chat
      end

      def replace_persisted_system_instructions(instructions)
        system_messages = messages_association.where(role: :system).order(:id).to_a

        if system_messages.empty?
          messages_association.create!(role: :system, content: instructions)
          return
        end

        primary_message = system_messages.shift
        primary_message.update!(content: instructions) if primary_message.content != instructions
        system_messages.each(&:destroy!)
      end

      def append_instructions?(append:, replace:)
        return append if replace.nil?

        append || (replace == false)
      end

      def persist_system_instruction(instructions, append:)
        transaction do
          if append
            messages_association.create!(role: :system, content: instructions)
          else
            replace_persisted_system_instructions(instructions)
          end
        end
      end

      def order_messages_for_llm(messages)
        system_messages, non_system_messages = messages.partition { |msg| msg.role.to_s == 'system' }
        system_messages + non_system_messages
      end

      def runtime_instructions
        @runtime_instructions ||= []
      end

      def store_runtime_instruction(instructions, append:)
        if append
          runtime_instructions << instructions
        else
          @runtime_instructions = [instructions]
        end
      end

      def reapply_runtime_instructions(chat)
        return if runtime_instructions.empty?

        first, *rest = runtime_instructions
        chat.with_instructions(first)
        rest.each { |instruction| chat.with_instructions(instruction, append: true) }
      end

      def persist_new_message
        @message = messages_association.create!(role: :assistant, content: '')
      end

      # rubocop:disable Metrics/PerceivedComplexity
      def persist_message_completion(message)
        return unless message

        tool_call_id = find_tool_call_id(message.tool_call_id) if message.tool_call_id

        transaction do
          content_text, attachments_to_persist, content_raw = prepare_content_for_storage(message.content)

          attrs = {
            role: message.role,
            content: content_text,
            input_tokens: message.input_tokens,
            output_tokens: message.output_tokens
          }
          attrs[:cached_tokens] = message.cached_tokens if @message.has_attribute?(:cached_tokens)
          if @message.has_attribute?(:cache_creation_tokens)
            attrs[:cache_creation_tokens] = message.cache_creation_tokens
          end
          attrs[:thinking_text] = message.thinking&.text if @message.has_attribute?(:thinking_text)
          attrs[:thinking_signature] = message.thinking&.signature if @message.has_attribute?(:thinking_signature)
          attrs[:thinking_tokens] = message.thinking_tokens if @message.has_attribute?(:thinking_tokens)

          # Add model association dynamically
          attrs[self.class.model_association_name] = model_association

          if tool_call_id
            parent_tool_call_assoc = @message.class.reflect_on_association(:parent_tool_call)
            attrs[parent_tool_call_assoc.foreign_key] = tool_call_id
          end

          @message.assign_attributes(attrs)
          @message.content_raw = content_raw if @message.respond_to?(:content_raw=)
          @message.save!

          persist_content(@message, attachments_to_persist) if attachments_to_persist
          persist_tool_calls(message.tool_calls) if message.tool_calls.present?
        end
      end
      # rubocop:enable Metrics/PerceivedComplexity

      def persist_tool_calls(tool_calls, message_record: @message)
        tool_call_klass = message_record.tool_calls_association.klass
        supports_thought_signature = tool_call_klass.column_names.include?('thought_signature')

        tool_calls.each_value do |tool_call|
          attributes = tool_call.to_h
          attributes.delete(:thought_signature) unless supports_thought_signature
          attributes[:tool_call_id] = attributes.delete(:id)
          message_record.tool_calls_association.create!(**attributes)
        end
      end

      def find_tool_call_id(tool_call_id)
        messages = messages_association
        message_class = messages.klass
        tool_calls_assoc = message_class.tool_calls_association_name
        tool_call_table_name = message_class.reflect_on_association(tool_calls_assoc).table_name

        message_with_tool_call = messages.joins(tool_calls_assoc)
                                         .find_by(tool_call_table_name => { tool_call_id: tool_call_id })
        return nil unless message_with_tool_call

        tool_call = message_with_tool_call.tool_calls_association.find_by(tool_call_id: tool_call_id)
        tool_call&.id
      end

      def persist_content(message_record, attachments)
        return unless message_record.respond_to?(:attachments)

        attachables = prepare_for_active_storage(attachments)
        message_record.attachments.attach(attachables) if attachables.any?
      end

      def prepare_for_active_storage(attachments)
        Utils.to_safe_array(attachments).filter_map do |attachment|
          case attachment
          when ActionDispatch::Http::UploadedFile, ActiveStorage::Blob
            attachment
          when ActiveStorage::Attachment, ActiveStorage::Attached::One, ActiveStorage::Attached::Many
            active_storage_blobs(attachment)
          when Hash
            attachment.values.map { |v| prepare_for_active_storage(v) }
          else
            convert_to_active_storage_format(attachment)
          end
        end.flatten.compact
      end

      def convert_to_active_storage_format(source)
        return if source.blank?

        attachment = source.is_a?(RubyLLM::Attachment) ? source : RubyLLM::Attachment.new(source)

        if attachment.active_storage?
          active_storage_blobs(attachment.source)
        else
          {
            io: StringIO.new(attachment.content),
            filename: attachment.filename,
            content_type: attachment.mime_type
          }
        end
      rescue StandardError => e
        RubyLLM.logger.warn "Failed to process attachment #{source}: #{e.message}"
        nil
      end

      def active_storage_blobs(attachment)
        case attachment
        when ActiveStorage::Blob then attachment
        when ActiveStorage::Attachment, ActiveStorage::Attached::One then attachment.blob
        when ActiveStorage::Attached::Many then attachment.blobs
        end
      end

      def build_content(message, attachments)
        return message if content_like?(message)

        RubyLLM::Content.new(message, attachments)
      end

      def content_like?(object)
        object.is_a?(RubyLLM::Content) || object.is_a?(RubyLLM::Content::Raw)
      end

      def prepare_content_for_storage(content)
        attachments = nil
        content_raw = nil
        content_text = content

        case content
        when RubyLLM::Content::Raw
          content_raw = content.value
          content_text = nil
        when RubyLLM::Content
          attachments = content.attachments if content.attachments.any?
          content_text = content.text
        when Hash, Array
          content_raw = content
          content_text = nil
        end

        [content_text, attachments, content_raw]
      end
    end
  end
end
