require "zip/zip"
require "zip/zipfilesystem"


module Buildr

  # The ZipTask creates a new ZIP file. You can include any number of files and
  # and directories, use exclusion patterns, and include files into specific
  # directories.
  #
  # For example:
  #   zip("test.zip").tap do |task|
  #     task.include "srcs"
  #     task.include "README", "LICENSE"
  #   end
  #
  # See Buildr#zip.
  class ZipTask < Rake::FileTask

    # Which files go where. All the rules for including, excluding and merging files
    # are handled by this object. A zip has at least one path.
    class Path #:nodoc:

      attr_reader :actions
      
      def initialize(zip, path)
        @zip = zip
        @path = "#{path}/" if path
        expand_src = proc { (@files || []).map(&:to_s).uniq }
        @sources = [ expand_src ]
        @actions = [] << proc do |zip|
          expand_src.call.each do |file|
            if File.directory?(file)
              in_directory(file, @files) do |file, rel_path|
                puts "Adding #{@path}#{rel_path}" if Rake.application.options.trace
                zip.add("#{@path}#{rel_path}", file) { true }
              end
            else
              puts "Adding #{@path}#{File.basename(file)}" if Rake.application.options.trace
              zip.add("#{@path}#{File.basename(file)}", file) { true }
            end
          end
        end
      end

      # Documented in ZipTask.
      def include(*files)
        if Hash === files.last
          options = files.pop
        else
          options = {}
        end
        files = files.flatten

        if options[:path]
          path(options[:path]).include *files +[ options.reject { |k,v| k == :path } ]
        elsif options[:as]
          raise "You can only use the :as option in combination with the :path option" unless options.keys.size == 1
          raise "You can only use one file with the :as option" unless files.size == 1
          include_as(files.first.to_s, options[:as])
        elsif options[:merge]
          raise "You can only use the :merge option in combination with the :path option" unless options.keys.size == 1
          files.each { |file| merge file }
        elsif options.keys.empty?
          (@files ||= FileList[]).include files.map(&:to_s)
        else
          raise "Unrecognized option #{options.keys.join(", ")}"
        end
        self
      end
      alias :add :include

      # Documented in ZipTask.
      def exclude(*files)
        (@files ||= FileList[]).exclude *files
        self
      end

      # Documented in ZipTask.
      def merge(*files)
        if Hash === files.last
          options = files.pop
        else
          options = {}
        end

        if options[:path]
          path(options[:path]).merge *files +[ options.reject { |k,v| k == :path } ]
        elsif options.keys.empty?
          files.collect do |file|
            @sources << proc { file.to_s }
            expander = ZipExpander.new(file)
            @actions << proc { |zip| expander.expand(zip, @path) }
            expander
          end.first
        else
          raise "Unrecognized option #{options.keys.join(", ")}"
        end
      end

      # Documented in ZipTask.
      def path(path)
        path.blank? ? self : @zip.path("#{@path}#{path}")
      end

      # Documented in ZipTask.
      def root()
        @zip
      end

      # Returns all the source files.
      def sources()
        @sources.map(&:call).flatten
      end

    protected

      def include_as(source, as)
        @sources << proc { source }
        @actions << proc do |zip|
          file = source.to_s
          if File.directory?(file)
            in_directory(file) do |file, rel_path|
              if as == "."
                dest = (@path || "") + rel_path.split("/")[1..-1].join("/")
              else
                dest = "#{@path}#{as}#{rel_path}"
              end
              puts "Adding #{dest}" if Rake.application.options.trace
              zip.add(dest, file) { true }
            end
          else
            puts "Adding #{@path}#{as}" if Rake.application.options.trace
            zip.add("#{@path}#{as}", file) { true }
          end
        end
      end

      def in_directory(dir, excludes = nil)
        prefix = Regexp.new("^" + Regexp.escape(File.dirname(dir) + File::SEPARATOR))
        Dir["#{dir}/**/*"].
          reject { |file| File.directory?(file) || (excludes && excludes.exclude?(file)) }.
          each { |file| yield file, file.sub(prefix, "") }
      end

    end

    # Extend one Zip file into another.
    class ZipExpander #:nodoc:

      def initialize(zip_file)
        @zip_file = zip_file.to_s
      end

      def include(*files)
        (@includes ||= [])
        @includes |= files
        self
      end

      def exclude(*files)
        (@excludes ||= []) 
        @excludes |= files
        self
      end

      def expand(zip, path)
        @includes ||= ["*"]
        @excludes ||= []
        Zip::ZipFile.open(@zip_file) do |source|
          source.entries.reject { |entry| entry.directory? }.each do |entry|
            if @includes.any? { |pattern| File.fnmatch(pattern, entry.name) } &&
               !@excludes.any? { |pattern| File.fnmatch(pattern, entry.name) }
              puts "Adding #{path}#{entry.name}" if Rake.application.options.trace
              zip.get_output_stream("#{path}#{entry.name}") { |output| output.write source.read(entry) }
            end
          end
        end
      end

    end

    def initialize(*args) #:nodoc:
      super
      @paths = { nil=>Path.new(self, nil) }
      enhance do |task|
        puts "Creating #{task.name}" if verbose
        # We're here because the Zip file does not exist, or one of the files is
        # newer than the Zip contents; in the later case, opening the Zip file
        # will add to its contents instead of replacing it, so we want the Zip
        # gone before we change it. We also don't want to see any partial updates.
        rm task.name, :verbose=>false rescue nil
        mkpath File.dirname(task.name), :verbose=>false
        begin
          Zip::ZipFile.open(task.name, Zip::ZipFile::CREATE) do |zip|
            zip.restore_permissions = true
            create zip
          end
        rescue
          rm task.name, :verbose=>false rescue nil
          raise
        end
      end
    end

    # :call-seq:
    #   include(*files) => self
    #   include(*files, :path=>path) => self
    #   include(file, :as=>name) => self
    #   include(*zips, :merge=>true) => self
    #
    # Include files in the ZIP (or current path) and returns self.
    #
    # This method accepts three options. You can use :path to include files under
    # a specific path, for example:
    #   zip(..).include("foo", :path=>"bar")
    # includes the file bar as foo/bar. See also #path.
    #
    # You can use :as to include a file under a different name, for example:
    #   zip(..).include("foo", :as=>"bar")
    # You can use the :as option in combination with the :path option, but only with
    # a single file at a time.
    #
    # As a special case, you can include the entire contents of a directory by including
    # the directory and using :as=>".". This:
    #   zip(..).include("srcs", :as=>".")
    # will include all the source files, using the directory as a prerequisite. This:
    #   zip(..).include("srcs/*")
    # includes all the same source files, using the source files as a prerequisite.
    #
    # You can use :merge option to include the contents of another ZIP file, for example:
    #   zip(..).include("foo.zip", :merge=>true)
    # You can use the :merge option in combination with the :path option. See also #merge.
    def include(*files)
      @paths[nil].include *files
      self
    end 
    alias :add :include
   
    # :call-seq:
    #   exclude(*files) => self
    # 
    # Excludes files and returns self. Can be used in combination with include to
    # prevent some files from being included.
    def exclude(*files)
      @paths[nil].exclude *files
      self
    end 

    # :call-seq:
    #   merge(*files) => Merge
    #   merge(*files, :path=>name) => Merge
    #
    # Merges ZIP files and returns a merge object. The contents of the merged ZIP file is
    # extracted into this ZIP file (or current path).
    #
    # The returned object supports two methods: include and exclude. You can use these to
    # merge only specific files from the ZIP. For example:
    #   zip(..).merge("src.zip").include("module1/*")
    #
    # This differs from include with the :merge option, which returns self.
    def merge(*files)
      @paths[nil].merge *files
    end 

    # :call-seq:
    #   path(name) => Path
    #
    # Returns a path object. You can use the path object to include files in a given
    # path inside the ZIP file. The path object implements the include, exclude, merge,
    # path and root methods.
    #
    # For example:
    #   zip(..).path("bar").include("foo")
    # Will add the file foo under the name bar/foo.
    #
    # As a shorthand, you can also use the :path option:
    #   zip(..).include("foo", :path=>"bar")
    def path(name)
      name.blank? ? @paths[nil] : (@paths[name] ||= Path.new(self, name))
    end

    # :call-seq:
    #   root() => ZipTask
    #
    # Returns the root path, essentially the ZipTask object itself. In case you are wondering
    # down paths and want to go back.
    def root()
      self
    end

    # :call-seq:
    #   with(options) => self
    #
    # Pass options to the task. Returns self. ZipTask itself does not support any options,
    # but other tasks (e.g. JarTask, WarTask) do.
    #
    # For example:
    #   package(:jar).with(:manifest=>"MANIFEST_MF")
    def with(options)
      options.each do |key, value|
        self[key] = value
      end
      self
    end

    # :call-seq:
    #   [name] = value
    #
    # Used by with method to set specific options. For example:
    #   package(:jar).with(:manifest=>"MANIFEST_MF")
    # Or:
    #   package(:jar)[:manifest] = "MANIFEST_MF"
    def []=(key, value)
      raise ArgumentError, "#{self.class} does not support the option #{key}"
    end

    def prerequisites() #:nodoc:
      super + @paths.collect { |name, path| path.sources }.flatten.each { |src| file(src) }
    end
    
    def needed?() #:nodoc:
      return true unless File.exist?(name)
      # You can do something like:
      #   include("foo", :path=>"foo").exclude("foo/bar", path=>"foo").
      #     include("foo/bar", :path=>"foo/bar")
      # This will play havoc if we handled all the prerequisites together
      # under the task, so instead we handle them individually for each path.
      #
      # We need to check that any file we include is not newer than the
      # contents of the ZIP. The file itself but also the directory it's
      # coming from, since some tasks touch the directory, e.g. when the
      # content of target/classes is included into a WAR.
      most_recent = @paths.collect { |name, path| path.sources }.flatten.
        each { |src| File.directory?(src) ? FileList["#{src}/**/*"] | [src] : src }.flatten.
        select { |file| File.exist?(file) }.collect { |file| File.stat(file).mtime }.max
      File.stat(name).mtime < (most_recent || Rake::EARLY) || super
    end

  protected

    # Sub-classes override this method to perform additional creation tasks,
    # e.g. creating a manifest file in a JAR.
    def create(zip)
      @paths.each { |name, obj| obj.actions.each { |action| action[zip] } }
    end

  end

  # :call-seq:
  #    zip(file) => ZipTask
  #
  # The ZipTask creates a new ZIP file. You can include any number of files and
  # and directories, use exclusion patterns, and include files into specific
  # directories.
  #
  # For example:
  #   zip("test.zip").tap do |task|
  #     task.include "srcs"
  #     task.include "README", "LICENSE"
  #   end
  def zip(file)
    ZipTask.define_task(file)
  end


  # An object for unzipping a file into a target directory. You can tell it to include
  # or exclude only specific files and directories, and also to map files from particular
  # paths inside the zip file into the target directory. Once ready, call #extract.
  #
  # Usually it is more convenient to create a file task for extracting the zip file
  # (see #unzip) and pass this object as a prerequisite to other tasks.
  #
  # See Buildr#unzip.
  class Unzip

    # The zip file to extract.
    attr_accessor :zip_file
    # The target directory to extract to.
    attr_accessor :target

    # Initialize with hash argument of the form target=>zip_file.
    def initialize(args)
      @target, @zip_file = Rake.application.resolve_args(args)
      @paths = {}
    end

    # :call-seq:
    #   extract()
    #
    # Extract the zip file into the target directory.
    #
    # You can call this method directly. However, if you are using the #unzip method,
    # it creates a file task for the target directory: use that task instead as a
    # prerequisite. For example:
    #   build unzip(dir=>zip_file)
    # Or:
    #   unzip(dir=>zip_file).target.invoke
    def extract()
      # If no paths specified, then no include/exclude patterns
      # specified. Nothing will happen unless we include all files.
      if @paths.empty?
        @paths[nil] = FromPath.new(nil)
        @paths[nil].include "*"
      end

      # Otherwise, empty unzip creates target as a file when touching.
      mkpath target.to_s, :verbose=>false
      Zip::ZipFile.open(zip_file.to_s) do |zip|
        entries = zip.collect
        @paths.each do |path, patterns|
          patterns.map(entries).each do |dest, entry|
            next if entry.directory?
            dest = File.expand_path(dest, target.to_s)
            puts "Extracting #{dest}" if Rake.application.options.trace
            mkpath File.dirname(dest), :verbose=>false rescue nil
            entry.extract(dest) { true }
          end
        end
      end
      # Let other tasks know we updated the target directory.
      touch target.to_s, :verbose=>false
    end

    # :call-seq:
    #   include(*files) => self
    #   include(*files, :path=>name) => self
    #
    # Include all files that match the patterns and returns self.
    #
    # Use include if you only want to unzip some of the files, by specifying
    # them instead of using exclusion. You can use #include in combination
    # with #exclude.
    def include(*files)
      if Hash === files.last
        from_path(files.pop[:path]).include *files
      else
        from_path(nil).include *files
      end
      self
    end
    alias :add :include

    # :call-seq:
    #   exclude(*files) => self
    #
    # Exclude all files that match the patterns and return self.
    #
    # Use exclude to unzip all files except those that match the pattern.
    # You can use #exclude in combination with #include.
    def exclude(*files)
      if Hash === files.last
        from_path(files.pop[:path]).exclude *files
      else
        from_path(nil).exclude *files
      end
      self
    end

    # :call-seq:
    #   from_path(name) => Path
    #
    # Allows you to unzip from a path. Returns an object you can use to
    # specify which files to include/exclude relative to that path.
    # Expands the file relative to that path.
    #
    # For example:
    #   unzip("test.jar").into(Dir.pwd).from_path("etc").include("LICENSE")
    # will unzip etc/LICENSE into ./LICENSE.
    #
    # This is different from:
    #  unzip("test.jar").into(Dir.pwd).include("etc/LICENSE")
    # which unzips etc/LICENSE into ./etc/LICENSE.
    def from_path(name)
      @paths[name] ||= FromPath.new(name)
    end

    # Returns the path to the target directory.
    def to_s()
      target.to_s
    end

    class FromPath #:nodoc:

      def initialize(path)
        if path
          @path = path[-1] == ?/ ? path : path + "/"
        else
          @path = ""
        end
      end

      # See UnzipTask#include
      def include(*files) #:doc:
        @include ||= []
        @include |= files
        self
      end

      # See UnzipTask#exclude
      def exclude(*files) #:doc:
        @exclude ||= []
        @exclude |= files
        self
      end

      def map(entries)
        includes = @include || ["*"]
        excludes = @exclude || []
        entries.inject({}) do |map, entry|
          short = entry.name.sub(@path, "")
          if includes.any? { |pat| File.fnmatch(pat, short) } &&
             !excludes.any? { |pat| File.fnmatch(pat, short) }
            map[short] = entry
          end
          map
        end
      end

    end

  end

  # :call-seq:
  #    unzip(to_dir=>zip_file) => Zip
  #
  # Creates a task that will unzip a file into the target directory. The task name
  # is the target directory, the prerequisite is the file to unzip.
  #
  # This method creates a file task to expand the zip file. It returns an Unzip object
  # that specifies how the file will be extracted. You can include or exclude specific
  # files from within the zip, and map to different paths.
  #
  # The Unzip object's to_s method return the path to the target directory, so you can
  # use it as a prerequisite. By keeping the Unzip object separate from the file task,
  # you overlay additional work on top of the file task.
  #
  # For example:
  #   unzip("all"=>"test.zip")
  #   unzip("src"=>"test.zip").include("README", "LICENSE") 
  #   unzip("libs"=>"test.zip").from_path("libs")
  def unzip(args)
    target, zip_file = Rake.application.resolve_args(args)
    task = file(File.expand_path(target.to_s)=>zip_file)
    Unzip.new(task=>zip_file).tap do |setup|
      task.enhance { setup.extract }
    end
  end

end