module Mongrel
  module Const
    POST         = 'POST'.freeze         unless const_defined?(:POST)
    QUERY_STRING = 'QUERY_STRING'.freeze unless const_defined?(:QUERY_STRING)
    UPLOAD_ID    = 'upload_id'.freeze
  end
  
  HttpHandler.class_eval do
    def request_aborted(params)
    end
  end
  
  HttpRequest.class_eval do
    def initialize_with_abort(params, socket, dispatchers)
      initialize_without_abort(params, socket, dispatchers)
      dispatchers.each {|d| d.request_aborted(params) if @body.nil? && d }
    end
    
    alias_method :initialize_without_abort, :initialize
    alias_method :initialize, :initialize_with_abort
  end
end


class MerbUploadHandler < Mongrel::HttpHandler

  def initialize(options = {})
    @path_info      = Array(options[:path_info])
    @frequency      = options[:frequency] || 3
    @request_notify = true
    if options[:drb]
      require 'drb'
      DRb.start_service
      Mongrel.const_set :Uploads, DRbObject.new(nil, "druby://#{Merb::Server.host}:#{Merb::Server.drb_server_port}").upload_progress
    else
      require File.dirname(__FILE__)+'/merb_upload_progress'
      Mongrel.const_set :Uploads, Merb::UploadProgress.new
    end
    Mongrel::Uploads.debug = true if options[:debug]
  end

  def request_begins(params)
    upload_notify(:add, params, params[Mongrel::Const::CONTENT_LENGTH].to_i)
  end

  def request_progress(params, clen, total)
    upload_notify(:mark, params, clen)
  end

  def process(request, response)
    upload_notify(:finish, request.params)
  end
  
  def request_aborted(params)
    return unless upload_id = valid_upload?(params)
    Mongrel::Uploads.finish(upload_id)
    puts "request aborted!"
  end

  private
    def upload_notify(action, params, *args)
      return unless upload_id = valid_upload?(params)
      if action == :mark
        last_checked_time = Mongrel::Uploads.last_checked(upload_id)
        return unless last_checked_time && Time.now - last_checked_time > @frequency
      end
      Mongrel::Uploads.send(action, upload_id, *args) 
      Mongrel::Uploads.update_checked_time(upload_id) unless action == :finish
    end
    
    def valid_upload?(params)
      @path_info.any? { |p| params[Mongrel::Const::PATH_INFO].include?(p) } &&
        params[Mongrel::Const::REQUEST_METHOD] == Mongrel::Const::POST &&
        Mongrel::HttpRequest.query_parse(params[Mongrel::Const::QUERY_STRING])[Mongrel::Const::UPLOAD_ID]
    end
end