pretty_irb 0.1.2 → 0.1.4

data/lib/pretty_irb/shell.rb

@@ -1,387 +1,435 @@
-require "irb"
-require "reline"
-
-module PrettyIRB
-  class Shell
-    BANNER = <<~BANNER
-      ╔═══════════════════════════════════════════════════════════╗
-      ║         🎨 Welcome to Pretty IRB 🎨                      ║
-      ║                                                           ║
-      ║  Enhanced Interactive Ruby Shell with Auto-Correct       ║
-      ║  Type 'help' for commands or 'exit' to quit              ║
-      ╚═══════════════════════════════════════════════════════════╝
-    BANNER
-
-    def initialize(config = {})
-      @config = config
-      @history_manager = HistoryManager.new
-      @snippet_manager = SnippetManager.new
-      @variable_inspector = VariableInspector.new(TOPLEVEL_BINDING)
-      @input_buffer = ""
-      @binding = TOPLEVEL_BINDING
-    end
-
-    def run
-      puts BANNER.light_magenta
-      start_repl
-    end
-
-    private
-
-    def start_repl
-      loop do
-        begin
-          input = read_input
-          break if input.nil? || input.downcase == "exit"
-
-          if input.downcase == "help"
-            show_help
-            next
-          end
-
-          # Handle special commands
-          if input.downcase == "clear"
-            system("clear") || system("cls")
-            next
-          end
-
-          # Handle history commands
-          if input.start_with?("history")
-            handle_history_command(input)
-            next
-          end
-
-          # Handle snippet commands
-          if input.start_with?("snippet")
-            handle_snippet_command(input)
-            next
-          end
-
-          # Handle variable commands
-          if input.start_with?("vars")
-            handle_variable_command(input)
-            next
-          end
-
-          # Handle benchmark commands
-          if input.start_with?("benchmark") || input.start_with?("bench")
-            handle_benchmark_command(input)
-            next
-          end
-
-          # Handle cheatsheet commands
-          if input.start_with?("cheat")
-            handle_cheatsheet_command(input)
-            next
-          end
-
-          # Handle AI help commands
-          if input.start_with?("?")
-            handle_ai_command(input)
-            next
-          end
-
-          execute_code(input)
-          @history_manager.add(input)
-
-        rescue Interrupt
-          puts "\n" + Formatter.format_warning("Interrupted. Type 'exit' to quit.")
-        rescue EOFError
-          puts "\n" + Formatter.format_success("Goodbye!")
-          break
-        rescue StandardError => e
-          handle_error(e)
-        end
-      end
-    end
-
-    def read_input
-      prompt = "pretty_irb".light_magenta + " >> ".light_magenta
-      input = Reline.readline(prompt, true)
-      return nil if input.nil?
-
-      # If the line appears to start a block or otherwise is incomplete,
-      # collect additional lines until the code is complete.
-      if incomplete_code?(input)
-        buffer = input + "\n"
-        while (more = Reline.readline("... ".light_magenta, true))
-          break if more.nil?
-          buffer += more + "\n"
-          break unless incomplete_code?(buffer)
-        end
-
-        return buffer
-      end
-
-      input
-    end
-
-    # Heuristic to detect incomplete Ruby code.
-    # Returns true if there are unclosed block tokens, brackets, or quotes.
-    def incomplete_code?(code)
-      return true if code.nil? || code.strip.empty?
-
-      # Count block starters and 'end'
-      opens = 0
-      opens += code.scan(/\b(do|def|class|module|if|case|begin|for|while|until|unless)\b/).size
-      closes = code.scan(/\bend\b/).size
-      return true if opens > closes
-
-      # Brackets / braces / parentheses
-      return true if code.count('(') > code.count(')')
-      return true if code.count('[') > code.count(']')
-      return true if code.count('{') > code.count('}')
-
-      # Unclosed quotes (simple heuristic, ignores escaped quotes)
-      single_quotes = code.scan(/(?<!\\)'/).size
-      double_quotes = code.scan(/(?<!\\)"/).size
-      return true if single_quotes.odd? || double_quotes.odd?
-
-      # Line continuation with trailing backslash
-      return true if code.rstrip.end_with?('\\')
-
-      false
-    end
-
-    def execute_code(code)
-      # Auto-correct the code
-      corrected_code = AutoCorrector.auto_correct_code(code)
-
-      # Show if code was corrected
-      if corrected_code != code
-        puts Formatter.format_warning("↻ Auto-corrected: #{corrected_code}")
-      end
-
-      begin
-        # Evaluate the code
-        result = eval(corrected_code, @binding)
-
-        # Display the result
-        unless result.nil?
-          puts Formatter.format_output(result)
-        end
-
-      rescue NameError => e
-        handle_name_error(e)
-      rescue NoMethodError => e
-        handle_method_error(e)
-      rescue SyntaxError => e
-        handle_syntax_error(e)
-      rescue StandardError => e
-        handle_error(e)
-      end
-    end
-
-    def handle_name_error(error)
-      puts Formatter.format_error("NameError: #{error.message}")
-      suggestions = AutoCorrector.suggest_corrections(error)
-      puts suggestions if suggestions && !suggestions.empty?
-    end
-
-    def handle_method_error(error)
-      puts Formatter.format_error("NoMethodError: #{error.message}")
-      suggestions = AutoCorrector.suggest_corrections(error)
-      puts suggestions if suggestions && !suggestions.empty?
-    end
-
-    def handle_syntax_error(error)
-      puts Formatter.format_error("SyntaxError: #{error.message}")
-      suggestions = AutoCorrector.suggest_corrections(error)
-      puts suggestions if suggestions && !suggestions.empty?
-    end
-
-    def handle_error(error)
-      puts Formatter.format_error("#{error.class}: #{error.message}")
-      if error.backtrace
-        puts Formatter.format_warning("Backtrace:")
-        error.backtrace.first(5).each do |line|
-          puts "  #{line}"
-        end
-      end
-    end
-
-    def show_help
-      help_text = <<~HELP
-        #{'=== Pretty IRB Commands ==='.light_magenta}
-
-        exit, quit          Exit the shell
-        help                Show this help message
-        clear               Clear the screen
-
-        #{'=== History Commands ==='.light_magenta}
-
-        history                       Show all command history
-        history search KEYWORD        Search history
-        history last N                Show last N commands
-        history export FILE           Export to file
-        history clear                 Clear history
-
-        #{'=== Snippet Commands ==='.light_magenta}
-
-        snippet list                  List all snippets
-        snippet save NAME CODE        Save code snippet
-        snippet load NAME             Load and execute snippet
-        snippet show NAME             Show snippet code
-        snippet delete NAME           Delete snippet
-        snippet search KEYWORD        Search snippets
-
-        #{'=== Variable Commands ==='.light_magenta}
-
-        vars                          List all variables
-        vars VARNAME                  Inspect variable
-        vars type:TYPE                Variables of type
-        vars search:KEYWORD           Search variables
-        vars memory                   Memory usage
-
-        #{'=== Benchmark Commands ==='.light_magenta}
-
-        bench CODE                    Benchmark code
-        bench compare CODE1 vs CODE2  Compare two snippets
-        bench memory CODE             Profile memory
-
-        #{'=== Cheatsheet Commands ==='.light_magenta}
-
-        cheat                         Show Ruby cheatsheet
-        cheat TOPIC                   Cheatsheet for topic
-                                      (array, hash, string, enumerable, file, regex, date)
-
-        #{'=== AI Help Commands ==='.light_magenta}
-
-        ?explain(method)    Explain a Ruby method
-        ?example(topic)     Get code examples
-        ?debug(code)        Analyze code for issues
-        ?practices(topic)   Learn best practices
-        ?ref(keyword)       Quick reference
-
-        #{'You can type any valid Ruby code and it will be executed!'.light_green}
-      HELP
-      puts help_text
-    end
-
-    def handle_history_command(input)
-      case input.downcase
-      when "history"
-        puts @history_manager.all
-      when /^history\s+search\s+(.+)$/
-        keyword = $1.strip
-        puts @history_manager.search(keyword)
-      when /^history\s+last\s+(\d+)$/
-        n = $1.to_i
-        puts @history_manager.last(n)
-      when /^history\s+export\s+(.+)$/
-        filename = $1.strip
-        puts @history_manager.export(filename)
-      when "history clear"
-        puts @history_manager.clear
-      else
-        puts Formatter.format_warning("Usage: history | history search KEYWORD | history last N | history export FILE | history clear")
-      end
-    end
-
-    def handle_snippet_command(input)
-      case input.downcase
-      when "snippet list"
-        puts @snippet_manager.list
-      when /^snippet\s+save\s+(\w+)\s+(.+)$/
-        name = $1
-        code = $2
-        puts @snippet_manager.save(name, code)
-      when /^snippet\s+load\s+(\w+)$/
-        name = $1
-        code = @snippet_manager.load(name)
-        if code
-          puts Formatter.format_success("📌 Loaded snippet: #{name}")
-          puts code
-          execute_code(code)
-        else
-          puts code
-        end
-      when /^snippet\s+show\s+(\w+)$/
-        name = $1
-        puts @snippet_manager.show(name)
-      when /^snippet\s+delete\s+(\w+)$/
-        name = $1
-        puts @snippet_manager.delete(name)
-      when /^snippet\s+search\s+(.+)$/
-        keyword = $1.strip
-        puts @snippet_manager.search(keyword)
-      else
-        puts Formatter.format_warning("Usage: snippet list | snippet save NAME CODE | snippet load NAME | snippet show NAME | snippet delete NAME | snippet search KEYWORD")
-      end
-    end
-
-    def handle_variable_command(input)
-      case input.downcase
-      when "vars"
-        puts @variable_inspector.list_variables
-      when /^vars\s+(.+)$/
-        var_name = $1.strip
-        if var_name == "memory"
-          puts @variable_inspector.memory_usage
-        elsif var_name.start_with?("type:")
-          type = var_name.sub("type:", "").strip
-          puts @variable_inspector.by_type(type)
-        elsif var_name.start_with?("search:")
-          keyword = var_name.sub("search:", "").strip
-          puts @variable_inspector.search(keyword)
-        else
-          puts @variable_inspector.inspect_var(var_name)
-        end
-      else
-        puts Formatter.format_warning("Usage: vars | vars VARNAME | vars type:TYPE | vars search:KEYWORD | vars memory")
-      end
-    end
-
-    def handle_benchmark_command(input)
-      case input
-      when /^bench(?:mark)?\s+(.+)$/
-        code = $1
-        puts Benchmarker.benchmark(code)
-      when /^bench(?:mark)?\s+compare\s+(.+?)\s+vs\s+(.+)$/
-        code1 = $1
-        code2 = $2
-        puts Benchmarker.compare(code1, code2)
-      when /^bench(?:mark)?\s+memory\s+(.+)$/
-        code = $1
-        puts Benchmarker.profile_memory(code)
-      else
-        puts Formatter.format_warning("Usage: bench CODE | bench compare CODE1 vs CODE2 | bench memory CODE")
-      end
-    end
-
-    def handle_cheatsheet_command(input)
-      case input.downcase
-      when "cheat"
-        puts CheatSheet.show
-      when /^cheat\s+(\w+)$/
-        topic = $1
-        puts CheatSheet.show(topic)
-      else
-        puts Formatter.format_warning("Usage: cheat | cheat TOPIC")
-        puts "Topics: array, hash, string, enumerable, file, regex, date"
-      end
-    end
-
-    def handle_ai_command(input)
-      case input
-      when /^\?explain\((.*?)\)$/
-        method_name = $1.gsub(/'|"/, "")
-        puts AIHelper.explain_method(method_name)
-      when /^\?example\((.*?)\)$/
-        topic = $1.gsub(/'|"/, "")
-        puts AIHelper.get_example(topic)
-      when /^\?debug\((.*)\)$/
-        code_snippet = $1
-        puts AIHelper.debug_code(code_snippet)
-      when /^\?practices\((.*?)\)$/
-        topic = $1.gsub(/'|"/, "")
-        puts AIHelper.best_practices(topic)
-      when /^\?ref\((.*?)\)$/
-        keyword = $1.gsub(/'|"/, "")
-        puts AIHelper.quick_reference(keyword)
-      else
-        puts Formatter.format_warning("❓ Unknown AI command. Type 'help' for available commands.")
-      end
-    end
-  end
-end
+require "irb"
+
+module PrettyIRB
+  class Shell
+    BANNER = <<~BANNER
+      ╭─────────────────────────────────────────────────────────────╮
+      │   pretty_irb  v#{PrettyIRB::VERSION}  ·  Enhanced Interactive Ruby Shell   │
+      │   type  commands  to list extras  ·  exit / Ctrl+D to quit  │
+      ╰─────────────────────────────────────────────────────────────╯
+    BANNER
+
+    module OutputHook
+      def output_value(omit = false)
+        return super unless PrettyIRB::Shell.active?
+        return if omit
+
+        puts PrettyIRB::Formatter.format_result(@context.last_value)
+      end
+    end
+
+    module ExceptionHook
+      def handle_exception(exc)
+        return super unless PrettyIRB::Shell.active?
+
+        first_line = exc.message.to_s.lines.first&.strip
+        PrettyIRB::Shell.history_manager&.record_error("#{exc.class}: #{first_line}")
+        super
+        PrettyIRB::Shell.print_exception_hint(exc)
+      end
+    end
+
+    module EvalHook
+      def evaluate(code, line_no, *args, **kwargs, &block)
+        raw = nil
+        if PrettyIRB::Shell.active?
+          raw = PrettyIRB::Shell.extract_raw_code(code)
+          if raw && PrettyIRB.configuration.auto_correct
+            corrected = PrettyIRB::AutoCorrector.auto_correct_code(raw)
+            if corrected.strip != raw.strip
+              label = PrettyIRB::UI.paint("  ↳ auto-corrected: ", :light_black)
+              fixed = PrettyIRB::UI.paint(corrected.strip, :yellow)
+              warn "#{label}#{fixed}"
+            end
+          end
+        end
+
+        result = begin
+          if kwargs.empty?
+            super(code, line_no, *args, &block)
+          else
+            super(code, line_no, *args, **kwargs, &block)
+          end
+        rescue ArgumentError
+          super(code, line_no)
+        end
+
+        if PrettyIRB::Shell.active? && raw && !raw.strip.empty?
+          preview = begin
+            str = result.inspect
+            str.length > 80 ? "#{str[0, 77]}..." : str
+          rescue
+            nil
+          end
+          PrettyIRB::Shell.history_manager&.add(raw.strip, preview)
+        end
+
+        result
+      end
+    end
+
+    module RailsIRBInitHook
+      def initialize(*args, **kwargs, &block)
+        if (shell = PrettyIRB::Shell.consume_pending_rails_shell)
+          shell.apply_irb_configuration!
+        end
+
+        if kwargs.empty?
+          super(*args, &block)
+        else
+          super(*args, **kwargs, &block)
+        end
+      end
+    end
+
+    class << self
+      attr_accessor :history_manager, :snippet_manager, :active, :pending_rails_shell
+
+      def active?
+        !!@active
+      end
+
+      def consume_pending_rails_shell
+        shell = @pending_rails_shell
+        @pending_rails_shell = nil
+        shell
+      end
+
+      def extract_raw_code(code)
+        if code.is_a?(String)
+          code
+        elsif code.respond_to?(:code)
+          code.code.to_s
+        end
+      end
+
+      def print_exception_hint(exc)
+        return unless PrettyIRB.configuration.auto_correct
+
+        suggestion = PrettyIRB::AutoCorrector.suggest_corrections(exc).to_s.strip
+        return if suggestion.empty?
+        if exc.message.to_s.include?("Did you mean?") &&
+           !suggestion.include?("Receiver is nil") &&
+           !suggestion.include?("Did you mean `.")
+          return
+        end
+
+        puts "  #{PrettyIRB::UI.paint(suggestion, :yellow)}"
+      rescue
+        nil
+      end
+    end
+
+    def initialize(config = {})
+      @config = config || {}
+    end
+
+    def run
+      STDOUT.sync = true
+      resolved_config = resolve_configuration
+      prev_config = PrettyIRB.configuration
+      PrettyIRB.instance_variable_set(:@configuration, resolved_config)
+      self.class.active = true
+
+      puts render_banner if resolved_config.banner
+
+      self.class.history_manager = HistoryManager.new
+      self.class.snippet_manager = SnippetManager.new
+
+      # Step 1: patch IRB classes idempotently before anything else
+      install_output_hook
+      install_eval_hook
+      install_exception_hook
+
+      @using_rails_irb_console = rails_irb_console_available?
+      if @using_rails_irb_console
+        run_with_rails_irb_console
+      else
+        run_standalone_or_legacy_rails
+      end
+    ensure
+      @using_rails_irb_console = false
+      self.class.consume_pending_rails_shell
+      self.class.active = false
+      PrettyIRB.instance_variable_set(:@configuration, prev_config)
+    end
+
+    # Applies PrettyIRB's Reline tuning, prompt, and IRB_RC hook onto the
+    # current IRB.conf after IRB.setup has initialized it.
+    def apply_irb_configuration!
+      tune_input_mode
+      configure_prompt
+      install_irb_rc_hook
+    end
+
+    private
+
+    def resolve_configuration
+      PrettyIRB.configuration.merge(@config)
+    end
+
+    def rails_app
+      return @config[:app] if @config.is_a?(Hash) && @config[:app]
+      return Rails.application if defined?(Rails) && Rails.respond_to?(:application)
+
+      nil
+    end
+
+    def rails_irb_console_available?
+      return false unless rails_app
+      return true if defined?(Rails::Console::IRBConsole)
+
+      require "rails/commands/console/irb_console"
+      !!defined?(Rails::Console::IRBConsole)
+    rescue LoadError
+      false
+    end
+
+    # Rails 7.2+ / 8.x delegates IRB initialization, helper registration
+    # (`reload!`, `app`, `helper`, `controller`), and backtrace cleaning to
+    # Rails::Console::IRBConsole. Hook into IRB::Irb#initialize so PrettyIRB
+    # layers on top right after IRB.setup(nil) completes.
+    def run_with_rails_irb_console
+      install_rails_irb_hook
+      self.class.pending_rails_shell = self
+      build_rails_irb_console(rails_app).start
+    end
+
+    def build_rails_irb_console(app)
+      init_arity = Rails::Console::IRBConsole.instance_method(:initialize).arity
+      if init_arity == 1
+        Rails::Console::IRBConsole.new(app)
+      else
+        Rails::Console::IRBConsole.new(app, @config[:rails_options] || {})
+      end
+    end
+
+    # Standalone Ruby or Rails 7.0 / 7.1 (before Rails::Console::IRBConsole).
+    def run_standalone_or_legacy_rails
+      app = rails_app
+      load_legacy_rails_console_helpers if app
+
+      # IRB.setup populates IRB.conf (and resets IRB_RC to nil).
+      IRB.setup(nil)
+
+      configure_legacy_rails_irb(app) if app
+      apply_irb_configuration!
+
+      irb = IRB::Irb.new
+      IRB.conf[:MAIN_CONTEXT] = irb.context
+      irb.run(IRB.conf)
+    end
+
+    def load_legacy_rails_console_helpers
+      require "rails/console/app"
+      require "rails/console/helpers"
+    rescue LoadError
+      nil
+    end
+
+    def configure_legacy_rails_irb(app)
+      IRB.conf[:IRB_NAME] = rails_app_name(app)
+      if !rails_local_env? && !ENV.key?("IRB_USE_AUTOCOMPLETE")
+        IRB.conf[:USE_AUTOCOMPLETE] = false
+      end
+
+      if defined?(Rails) && Rails.respond_to?(:backtrace_cleaner) && Rails.backtrace_cleaner
+        prev_filter = IRB.conf[:BACKTRACE_FILTER]
+        IRB.conf[:BACKTRACE_FILTER] = lambda do |backtrace|
+          backtrace = prev_filter.call(backtrace) if prev_filter.respond_to?(:call)
+          Rails.backtrace_cleaner.filter(backtrace)
+        end
+      end
+
+      if defined?(IRB::ExtendCommandBundle) && defined?(Rails::ConsoleMethods)
+        IRB::ExtendCommandBundle.include(Rails::ConsoleMethods)
+      end
+    end
+
+    def rails_app_name(app)
+      if app.respond_to?(:name) && app.name
+        app.name.to_s
+      elsif app.class.respond_to?(:module_parent_name) && app.class.module_parent_name
+        app.class.module_parent_name.downcase
+      else
+        "rails"
+      end
+    end
+
+    def rails_local_env?
+      return true unless defined?(Rails) && Rails.respond_to?(:env) && Rails.env
+
+      if Rails.env.respond_to?(:local?)
+        Rails.env.local?
+      else
+        Rails.env.development? || Rails.env.test?
+      end
+    end
+
+    # ── Terminal / input detection ────────────────────────────────────────────
+    # On Windows, Reline uses a Win32 API console handler.  In "legacy
+    # console" mode (old conhost.exe — PowerShell or cmd.exe without
+    # Windows Terminal) Alt+key sequences are read in a way that blocks the
+    # input loop, freezing the entire REPL.
+
+    def tune_input_mode
+      return unless Gem.win_platform?
+
+      legacy = begin
+        Reline::IOGate.instance_variable_get(:@legacy_console)
+      rescue
+        false
+      end
+
+      if legacy
+        IRB.conf[:USE_MULTILINE]    = false   # avoid Reline Windows hang
+        IRB.conf[:USE_AUTOCOMPLETE] = false
+        msg = "  ℹ Legacy Windows console detected — running in single-line mode.\n" \
+              "    Open Windows Terminal for full multiline & autocomplete support."
+        puts style_text(msg, :light_black)
+      end
+    rescue
+      nil
+    end
+
+    # ── Banner & Prompt ───────────────────────────────────────────────────────
+
+    def render_banner
+      title = "#{style_text('pretty_irb', :white, :bold)} #{style_text("v#{PrettyIRB::VERSION}", :cyan)}"
+      lines = [
+        "#{style_text('Enhanced Interactive Ruby Shell', :white)} #{style_text("· Ruby #{RUBY_VERSION}", :light_black)}"
+      ]
+
+      if (rails_line = build_rails_banner_line)
+        lines << rails_line
+      end
+
+      lines << :divider
+      lines << [
+        "#{style_text('commands', :cyan)} #{style_text('extras', :light_black)}",
+        "#{style_text('vars', :cyan)} #{style_text('scope', :light_black)}",
+        "#{style_text('bench', :cyan)} #{style_text('profile', :light_black)}",
+        "#{style_text('help', :cyan)} #{style_text('IRB', :light_black)}",
+        "#{style_text('Ctrl+D', :white)} #{style_text('exit', :light_black)}"
+      ].join(style_text("  ·  ", :light_black))
+
+      PrettyIRB::UI.box(lines, title: title, min_width: 66)
+    end
+
+    def build_rails_banner_line
+      app = rails_app
+      return nil unless app
+
+      rails_ver = defined?(Rails) && Rails.respond_to?(:version) ? "Rails #{Rails.version}" : "Rails"
+      parts = [style_text("#{rails_app_name(app)} · #{rails_ver}", :light_black)]
+      if (badge = rails_env_badge)
+        parts << badge
+      end
+      parts << PrettyIRB::UI.badge("SANDBOX", color: :yellow, bold: true) if rails_sandbox?(app)
+      parts.join("  ")
+    end
+
+    def rails_sandbox?(app = rails_app)
+      app && app.respond_to?(:sandbox?) && app.sandbox?
+    end
+
+    def configure_prompt
+      prefix       = build_prompt_prefix
+      caret_styles = rails_production? ? %i[red bold] : %i[cyan bold]
+      IRB.conf[:PROMPT] ||= {}
+      IRB.conf[:PROMPT][:PRETTY_IRB] = {
+        PROMPT_I: "#{prefix} #{style_text('❯ ', *caret_styles)}",
+        PROMPT_N: "#{prefix} #{style_text('│ ', :light_black)}",
+        PROMPT_S: "#{prefix} #{style_text('┆ ', :light_black)}",
+        PROMPT_C: "#{prefix} #{style_text('· ', :light_black)}",
+        RETURN:   "=> %s\n"
+      }
+      IRB.conf[:PROMPT_MODE]  = :PRETTY_IRB
+      IRB.conf[:SAVE_HISTORY] ||= 1000
+    end
+
+    def build_prompt_prefix
+      base  = style_text("[%03n]", :light_black)
+      badge = rails_env_badge
+      badge ? "#{base} #{badge}" : base
+    end
+
+    def rails_production?
+      defined?(Rails) && Rails.respond_to?(:env) && Rails.env.to_s == "production"
+    end
+
+    def rails_env_badge
+      return nil unless PrettyIRB.configuration.rails_env_prompt
+      return nil unless defined?(Rails) && Rails.respond_to?(:env) && Rails.env
+
+      env = Rails.env.to_s
+      return nil if env.empty?
+
+      styled_env = case env
+                   when "development" then style_text("dev", :cyan)
+                   when "test"        then style_text("test", :light_blue)
+                   when "production"  then style_text("PROD", :red, :bold)
+                   else                    style_text(env, :yellow)
+                   end
+
+      if rails_sandbox?
+        "(#{styled_env}#{style_text(':', :light_black)}#{style_text('sandbox', :yellow)})"
+      else
+        "(#{styled_env})"
+      end
+    end
+
+    def install_irb_rc_hook
+      prev_rc = IRB.conf[:IRB_RC]
+      IRB.conf[:IRB_RC] = proc do |ctx|
+        prev_rc.call(ctx) if prev_rc.respond_to?(:call)
+        ctx.prompt_mode      = :PRETTY_IRB
+        ctx.auto_indent_mode = true
+        ctx.workspace.main.extend(PrettyIRB::IRBCommands)
+        if !@using_rails_irb_console && defined?(Rails::ConsoleMethods)
+          ctx.workspace.main.extend(Rails::ConsoleMethods)
+        end
+        IRB.conf[:MAIN_CONTEXT] ||= ctx
+      end
+    end
+
+    def style_text(text, *styles)
+      PrettyIRB::UI.paint(text, *styles)
+    end
+
+    # ── Hooks ─────────────────────────────────────────────────────────────────
+
+    def install_output_hook
+      return if IRB::Irb.ancestors.include?(OutputHook)
+
+      IRB::Irb.prepend(OutputHook)
+    rescue => e
+      warn style_text("  pretty_irb: output hook skipped (#{e.message})", :light_black)
+    end
+
+    def install_eval_hook
+      return if IRB::Context.ancestors.include?(EvalHook)
+
+      IRB::Context.prepend(EvalHook)
+    rescue => e
+      warn style_text("  pretty_irb: eval hook skipped (#{e.message})", :light_black)
+    end
+
+    def install_exception_hook
+      return unless IRB::Irb.instance_methods.include?(:handle_exception)
+      return if IRB::Irb.ancestors.include?(ExceptionHook)
+
+      IRB::Irb.prepend(ExceptionHook)
+    rescue => e
+      warn style_text("  pretty_irb: exception hook skipped (#{e.message})", :light_black)
+    end
+
+    def install_rails_irb_hook
+      return if IRB::Irb.ancestors.include?(RailsIRBInitHook)
+
+      IRB::Irb.prepend(RailsIRBInitHook)
+    rescue => e
+      warn style_text("  pretty_irb: rails init hook skipped (#{e.message})", :light_black)
+    end
+  end
+end