# Extend Jekyll 'highlight' tag to allow captions.
# You can use following options:
# 1. {% highlight <lang> [linenos] [caption] %}
# 2. {% highlight <lang> [linenos] [caption=some-text_without.duble/quotes] %}

module JekyllTagsExtensions
  module HighlightCaption
    def self.prepended(mod)
      syntax = mod.send(:remove_const, :SYNTAX).source.gsub('\\w', '[^[:^graph:]"=]')
      mod.const_set(:SYNTAX, Regexp.new(syntax))
      options = mod.send(:remove_const, :OPTIONS_REGEX).source.gsub('\\w', '[^[:^graph:]"=]')
      mod.const_set(:OPTIONS_REGEX, Regexp.new(options))
    end

    def render(*args)
      caption = "<figcaption><span class=\"lang\">"
      caption << "#{@lang.upcase}#{@highlight_options[:caption] ? ':' : ''}</span>"
      if @highlight_options[:caption].is_a? String
        caption << "<em>#{@highlight_options[:caption]}</em>" 
      end
      caption << "</figcaption>"

      match = super.match('<figure[^>]*>')
      match.pre_match + match.to_s + caption + match.post_match
    end
  end
end

Jekyll::Hooks.register :site, :pre_render do |site|
  Jekyll::Tags::HighlightBlock.prepend JekyllTagsExtensions::HighlightCaption
end