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.
  class ZipTask < Rake::FileTask

    # :nodoc:
    module IncludeFiles

      # Include the specified files or directories.
      def include(*files)
        if Hash === files.last
          options = files.pop
        else
          options = {}
        end

        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, 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
        else
          raise "Unrecognized option #{options.keys.join(", ")}"
        end
        self
      end

      # Exclude the specified file or directories.
      def exclude(*files)
        (@files ||= FileList[]).exclude *files
        self
      end
      alias :add :include

      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|
            @expand_sources << proc { artifacts(file).map(&:to_s) }
            expander = ZipExpander.new(file)
            @add_files << proc { |zip| expander.expand(zip, @path) }
            expander
          end.first
        else
          raise "Unrecognized option #{options.keys.join(", ")}"
        end
      end

    protected

      def setup_path(path = nil)
        @path = "#{path}/" if path
        expand_src = proc { artifacts(*@files || []).map(&:to_s) }
        @expand_sources = [ expand_src ]
        @add_files = [] << 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
              end
            else
              puts "Adding #{@path}#{File.basename(file)}" if Rake.application.options.trace
              zip.add "#{@path}#{File.basename(file)}", file
            end
          end
        end
      end

      def include_as(source, as)
        @expand_sources << proc { artifacts(source).map(&:to_s) }
        @add_files << proc do |zip|
          file = artifacts(source).first.to_s
          if File.directory?(file)
            in_directory(file) do |file, rel_path|
              puts "Adding #{@path}#{as}/#{rel_path}" if Rake.application.options.trace
              zip.add file, "#{@path}#{as}/#{rel_path}"
            end
          else
            puts "Adding #{@path}#{as}" if Rake.application.options.trace
            zip.add "#{@path}#{as}", file
          end
        end
      end

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

      def expand_sources()
        @expand_sources.map(&:call).flatten
      end

      def add_file(zip)
        @add_files.each { |action| action.call zip }
      end

    end


    # Which files go where.
    class Path

      include IncludeFiles
      
      def initialize(path)
        setup_path path
      end
  
    end

    include IncludeFiles

    def initialize(*args)
      super
      @paths = { nil=>self }
      setup_path
      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) { |zip| create zip }
        rescue
          rm task.name, :verbose=>false rescue nil
          raise
        end
      end
    end


    class ZipExpander

      def initialize(zip_file)
        @zip_file = zip_file
      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(artifacts(@zip_file).map(&:to_s).first) 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) }
              # TODO: read and write file
            end
          end
        end
      end

    end


    # Returns a path to which you can include/exclude files.
    #
    #   zip(..).include("foo", :path=>"bar")
    # is equivalen to:
    #   zip(..).path("bar").include("foo")
    def path(path)
      path.blank? ? @paths[nil] : (@paths[path] ||= Path.new(path))
    end

    # Pass options to the task. Returns self.
    def with(options)
      options.each do |key, value|
        self[key] = value
      end
      self
    end

    def []=(key, value)
      fail "#{self.class} does not support the attribute #{key}"
    end
    
    def invoke_prerequisites()
      super
      @paths.collect { |name, path| path.expand_sources }.flatten.each { |src| file(src).invoke }
    end

    def needed?()
      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.expand_sources }.flatten.
        each { |src| File.directory?(src) ? FileList[File.join(src, "**", "*")] | [src] : src }.flatten.
        select { |file| File.exist?(file) }.collect { |file| File.stat(file).mtime }.max
      File.stat(name).mtime < (most_recent || Rake::EARLY)
    end

  protected

    def create(zip)
      @paths.each { |name, obj| obj.add_file zip }
    end

    def zip_map()
      unless @zip_map
        args = @paths.collect do |path, files|
          if path
            dest_for = proc { |f| File.join(path, f) }
          else
            dest_for = proc { |f| f }
          end
          artifacts(*files).map(&:to_s).collect do |file|
            if File.directory?(file)
              # Include all files inside the directory, with path starting
              # from the directory itself.
              prefix = File.dirname(file) + File::SEPARATOR
              Dir[File.join(file, "**", "*")].sort.
                reject { |file| File.directory?(file) || files.exclude?(file) }.
                collect { |file| [ dest_for[file.sub(prefix, "")], file ] }
            else
              # Include just that one file, sans its diectory path.
              [ dest_for[File.basename(file)], file ]
            end
          end
        end.flatten
        @zip_map = Hash[ *args ]
      end
      @zip_map
    end

  end

  # 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:
  #   returning(zip("test.zip")) { |task|
  #     task.include "srcs"
  #     task.include "README", "LICENSE"
  #   end
  def zip(file)
    ZipTask.define_task(file)
  end


  # The UnzipTask expands the contents of a ZIP file into a target directory.
  # You can include any number of files and directories, use exclusion patterns,
  # and expand files from relative paths.
  #
  # The file(s) to unzip is the first prerequisite.
  class UnzipTask < Rake::Task

    # The target directory.
    attr_accessor :target

    def initialize(*args)
      super
      @paths = {}
      enhance do |task|
        fail "Where do you want the file unzipped" unless target

        # 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, :verbose=>false
        prerequisites.each do |file|
          Zip::ZipFile.open(file) 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)
                puts "Extracting #{dest}" if Rake.application.options.trace
                mkpath File.dirname(dest), :verbose=>false rescue nil
                entry.extract(dest) { true }
              end
            end
          end
        end
        # Let other tasks know we updated the target directory.
        touch target, :verbose=>false
      end
    end

    # Specifies directory to unzip to and return self.
    def into(target)
      self.target = target
      self
    end

    # 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

    # 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

    # 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(path)
      @paths[path] ||= FromPath.new(path)
    end

    def needed?()
      #return true unless target && File.exist?(target)
      #return true if prerequisites.any? { |prereq| File.stat(prereq).mtime > File.stat(target).mtime }
      #false
      true
    end

    # :nodoc:
    class FromPath

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

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

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

      # :nodoc:
      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

  # Defines a task that will unzip the specified file, into the directory
  # specified by calling #into. It is the second call to into that creates
  # and returns the task.
  #
  # You can unzip only some files by specifying an inclusion or exclusion
  # pattern, and unzip files from a path in the ZIP file. See UnzipTask
  # for more information.
  #
  # For example:
  #   unzip("test.zip").into("test")
  #   unzip("test.zip").into("etc").include("README", "LICENSE") 
  #   unzip("test.zip").into("src").from_path("srcs")
  def unzip(file)
    task = nil
    namespace { task = UnzipTask.define_task("unzip"=>file) }
    task
  end

end