;;;; copyright (c) 2004 David Lichteblau <david@lichteblau.com>, GPL

(defpackage :zip
  (:use :cl)
  (:export #:zipfile
	   #:open-zipfile
	   #:close-zipfile
	   #:with-zipfile
	   #:zipfile-entries
	   #:get-zipfile-entry
	   #:zipfile-entry-name
	   #:do-zipfile-entries
	   #:zipfile-entry-contents
	   #:unzip))

(in-package :zip)

(defun make-byte-array (n)
  (make-array n :element-type '(unsigned-byte 8)))

(defun fetch-short (array offset)
  (logior (elt array offset)
	  (ash (elt array (+ offset 1)) 8)))

(defun fetch-int (array offset)
  (logior (elt array offset)
	  (ash (elt array (+ offset 1)) 8)
	  (ash (elt array (+ offset 2)) 16)
	  (ash (elt array (+ offset 3)) 24)))

(defmacro define-record (reader (&key (length (gensym))) &rest fields)
  `(progn
     (defconstant ,length
	 ,(loop
	      for (nil type) in fields
	      sum (ecase type (:int 4) (:short 2))))
     (defun ,reader (s)
       (let ((bytes (make-byte-array ,length)))
	 (read-sequence bytes s)
	 bytes))
     ,@(loop
	   for (name type) in fields
	   for offset = 0 then (+ offset length)
	   for length = (ecase type (:int 4) (:short 2))
	   for reader = (ecase type (:int 'fetch-int) (:short 'fetch-short))
	   unless (eq name :dummy)
	   collect `(defun ,name (r)
		      (,reader r ,offset)))))

(define-record read-end-header (:length +end-header-length+)
  (end/signature :int)
  (end/this-disc :short)
  (end/central-directory-disc :short)
  (end/disc-files :short)
  (end/total-files :short)
  (end/central-directory-size :int)
  (end/central-directory-offset :int)
  (end/comment-length :short))

(define-record read-directory-entry ()
  (cd/signature :int)
  (cd/version-made-by :short)
  (cd/version :short)
  (cd/flags :short)
  (cd/method :short)
  (cd/time :short)
  (cd/date :short)
  (cd/crc :int)
  (cd/compressed-size :int)
  (cd/size :int)
  (cd/name-length :short)
  (cd/extra-length :short)
  (cd/comment-length :short)
  (cd/disc-number :short)
  (cd/internal-attributes :short)
  (cd/external-attributes :int)
  (cd/offset :int))

(define-record read-local-header ()
  (file/signature :int)
  (file/version :short)
  (file/flags :short)
  (file/method :short)
  (file/time :short)
  (file/date :short)
  (file/crc :int)
  (file/compressed-size :int)
  (file/size :int)
  (file/name-length :short)
  (file/extra-length :short))

(defun seek-to-end-header (s)
  (let* ((len (+ 65536 +end-header-length+))
	 (guess (max 0 (- (file-length s) len))))
    (file-position s guess)
    (let ((v (make-byte-array (min (file-length s) len))))
      (read-sequence v s)
      (let ((n (search #(80 75 5 6) v :from-end t)))
	(unless n
	  (error "end of central directory header not found"))
	(file-position s (+ guess n))))))

(defstruct zipfile
  stream
  entries)

(defstruct zipfile-entry
  name
  stream
  offset)

(defun read-entry-object (s)
  (let* ((header (read-directory-entry s))
	 (name (make-string (cd/name-length header))))
    (assert (= (cd/signature header) #x02014b50))
    (read-sequence name s)
    (prog1
	(make-zipfile-entry :name name :stream s :offset (cd/offset header))
      (file-position s (+ (file-position s)
			  (cd/extra-length header)
			  (cd/comment-length header))))))

(defun open-zipfile (pathname)
  (let ((s (open pathname)))
    (unwind-protect
	(progn
	  (seek-to-end-header s)
	  (let* ((end (read-end-header s))
		 (n (end/total-files end))
		 (entries (make-hash-table :test #'equal))
		 (zipfile (make-zipfile :stream s :entries entries)))
	    (file-position s (end/central-directory-offset end))
	    (dotimes (x n)
	      (let ((entry (read-entry-object s)))
		(setf (gethash (zipfile-entry-name entry) entries) entry)))
	    #+sbcl (let ((s s)) (sb-ext:finalize zipfile (lambda ()(close s))))
	    (setf s nil)
	    zipfile))
      (when s
	(close s)))))

(defgeneric close-zipfile (zipfile))
(defgeneric get-zipfile-entry (name zipfile))
(defgeneric zipfile-entry-contents (entry))

(defmethod close-zipfile ((zipfile zipfile))
  (close (zipfile-stream zipfile)))

(defmethod get-zipfile-entry (name (zipfile zipfile))
  (gethash name (zipfile-entries zipfile)))

(defmethod zipfile-entry-contents ((entry zipfile-entry))
  (let ((s (zipfile-entry-stream entry))
	header)
    (file-position s (zipfile-entry-offset entry))
    (setf header (read-local-header s))
    (assert (= (file/signature header) #x04034b50))
    (file-position s (+ (file-position s)
			(file/name-length header)
			(file/extra-length header)))
    (let ((data (make-byte-array (file/compressed-size header))))
      (read-sequence data s)
      (deflate:rfc1951-uncompress-octets data))))

(defmacro with-zipfile ((file pathname) &body body)
  `(let ((,file (open-zipfile ,pathname)))
     (unwind-protect
	 (progn ,@body)
       (close-zipfile ,file))))

(defmacro do-zipfile-entries ((name entry zipfile) &body body)
  (setf name (or name (gensym)))
  (setf entry (or entry (gensym)))
  `(block nil
     (maphash (lambda (,name ,entry)
		(declare (ignorable ,name ,entry))
		,@body)
	      (zipfile-entries ,zipfile))))

(defun unzip (pathname target-directory &key (if-exists :error) verbose)
  (when (or (pathname-name target-directory)
	    (pathname-type target-directory))
    (error "pathname not a directory, lacks trailing slash?"))
  (with-zipfile (zip pathname)
    (do-zipfile-entries (name entry zip)
      (let ((filename (merge-pathnames name target-directory))
	    (data (zipfile-entry-contents entry)))
	(ecase verbose
	  ((nil))
	  ((t) (write-string name) (terpri))
	  (:dots (write-char #\.)))
	(force-output)
	(ensure-directories-exist filename)
	(with-open-file (s filename :direction :output :if-exists if-exists
			 :element-type '(unsigned-byte 8))
	  (write-sequence data s))))))
