#!/usr/bin/env ruby
#
# webshot.rb
#
# Copyright (c) 2007 Naoki Hiroshima
# You can redistribute it and/or modify it under the same terms as GPL2.
#
# Author:: Naoki Hiroshima <n at h7a dot org>

require 'thread'
require 'osx/cocoa'

OSX.require_framework "WebKit"
OSX::NSApplication.sharedApplication

class Webshot < OSX::NSObject
  ImageType = {
    :png  => OSX::NSPNGFileType,
    :jpg  => OSX::NSJPEGFileType,
    :gif  => OSX::NSGIFFileType,
    :tiff => OSX::NSTIFFFileType,
    :bmp  => OSX::NSBMPFileType,
    :pdf  => :dataWithPDFInsideRect
  }

  def self.run
    ws = Webshot.alloc.init
    ws.run
    yield ws

  rescue Exception => e
    puts e.to_s

  ensure
    OSX::NSApplication.sharedApplication.terminate(nil)
  end

  def initialize
    @width = 800
    @height = 600
    @sync = Queue.new

    rect = OSX::NSRect.new(0, 0, 1, 1)
    window = OSX::NSWindow.alloc.
      initWithContentRect_styleMask_backing_defer(rect, OSX::NSBorderlessWindowMask, 2, 0)

    @webView = OSX::WebView.alloc.initWithFrame(rect)
    @webView.mainFrame.frameView.setAllowsScrolling(false)

    OSX::NSNotificationCenter.defaultCenter.
      addObserver_selector_name_object(self,
                                       'webViewProgressEstimateChanged:',
                                       OSX::WebViewProgressEstimateChangedNotification,
                                       nil)

    OSX::NSNotificationCenter.defaultCenter.
      addObserver_selector_name_object(self,
                                       'webViewProgressFinished:',
                                       OSX::WebViewProgressFinishedNotification,
                                       nil)

    OSX::NSNotificationCenter.defaultCenter.
      addObserver_selector_name_object(self,
                                       'webViewProgressStarted:',
                                       OSX::WebViewProgressStartedNotification,
                                       nil)

    window.setContentView(@webView)
    @webView.setFrameLoadDelegate(self)
  end

  def run
    OSX::NSApplication.sharedApplication.setDelegate(self)

    Signal.trap(:INT) {
      OSX::NSApplication.sharedApplication.terminate(nil)
    } if Signal.list.keys.any? {|e| e == 'INT'}

    Thread.new { OSX::NSApplication.sharedApplication.run }
    waitUntilDone
  end

  def width=(width)
    performSelectorOnMainThread_withObject_waitUntilDone('setWidth:', [width], false)
    waitUntilDone
  end

  def save(url, filename)
    format = File.extname(filename).delete('.').to_sym
    unless ImageType[format]
      puts "not supported: #{format}"
      return
    end

    performSelectorOnMainThread_withObject_waitUntilDone('doSave:', [url, filename], false)
    waitUntilDone
  end

  #
  # protected methods
  #

  protected

  def setWidth(param)
    @width = param[0]
    done
  end

  def doSave(param)
    @url = param[0]
    @filename = param[1]

    @url = 'http://' + @url unless @url.match(/^[a-z]+:\/\//)

    @webView.window.setContentSize(OSX::NSSize.new(@width, @height))
    @webView.mainFrame.frameView.setAllowsScrolling false

    nsurl = OSX::NSURL.URLWithString(@url)
    request = OSX::NSURLRequest.requestWithURL(nsurl)
    @webView.mainFrame.loadRequest(request)
  end

  def waitUntilDone
    @sync.pop # synchronize
  end

  def done
    @sync << :done
  end

  #
  # event handlers
  #

  def webViewProgressEstimateChanged(notification)
    print "%3d%%" % (@webView.estimatedProgress * 100).to_i
    STDOUT.flush
    print "\033[4D"
  end

  def webViewProgressStarted(notification)
    print "#{@url} ... "
    STDOUT.flush
  end

  def webViewProgressFinished(notification)
    v = @webView.mainFrame.frameView.documentView
    return if v.nil?

    v.window.setContentSize(v.bounds.size)

    data = nil
    format = File.extname(@filename).delete('.').to_sym
    if format == :pdf
      data = v.dataWithPDFInsideRect v.frame;
    else
      v.lockFocus
      bitmap = OSX::NSBitmapImageRep.alloc.initWithFocusedViewRect(v.bounds)
      v.unlockFocus

      data = bitmap.representationUsingType_properties(ImageType[format], nil)
    end

    data.writeToFile_atomically(@filename, true)

    print "saved\n"
  ensure
    done
  end

  def applicationDidFinishLaunching(notification)
    done
  end

  def webView_didFailLoadWithError_forFrame(sender, error, frame)
    puts error.localizedDescription
    done
  end

  def webView_didFailProvisionalLoadWithError_forFrame(sender, error, frame)
    puts error.localizedDescription
    done
  end
end


##################################################

if __FILE__ == $0

  def help
    puts <<EOT
Usage: #{$0} arguments
  -q        process quietly
  -h        show this
  -[num]    set browser\'s width
  -[fmt]    set output format (#{Webshot::ImageType.keys.join '|'})
   url      URL to save
EOT
  end

  help && exit if ARGV.empty?

  Webshot.run do |webshot|
    format = :png
    ARGV.each do |arg|
      if arg.match(/^-([a-z0-9]*)/)
        if $1 == 'q'
          def print(*args); end
        elsif $1 == 'h'
          help
          break
        elsif $1.to_i > 0
          webshot.width = $1.to_i
        else
          format = $1.to_sym
          if Webshot::ImageType[format].nil?
            help
            break
          end
        end
      else
        filename = arg.sub(/^[a-z]*:\/\//, '').gsub(/[:?\/]/, '_') + '.' + format.to_s
        webshot.save arg, filename
      end
    end
  end
end