(require 'bibtex)
(require 'ox)

(defvar ox-bib-find-entries-in-org-file nil
  "Whether to look for bibliogrpahy entries in an org mode file
or a bibtex file.")

(defun ox-bib-get-bibliography ()
  (mapcar (lambda (s)
            (org-no-properties s))
	  ;; Don't explode if before the first headline
          (or
           (condition-case nil
               (plist-get (org-export-get-environment nil t) :bibliography)
             (error nil))
           (save-excursion
             (widen)
             (goto-char (point-min))
             (plist-get (org-export-get-environment) :bibliography)))))

(defun ox-bib-get-key (path &optional desc)
  "Extract cite key from a link path and desc.

Currently expects links in the format:
\[[type:key;pre;post;][desc]]"
  (nth 0 (split-string path ";")))

(defun ox-bib-get-pre (path &optional desc)
  "Extract pre-citation arg from a link path and desc.

Currently expects links in the format:
\[[type:key;pre;post;][desc]]"
  (or (nth 1 (split-string path ";")) ""))

(defun ox-bib-get-post (path &optional desc)
  "Extract post-citation arg from a link path and desc.

Currently expects links in the format:
\[[type:key;pre;post;][desc]]"
  (or (nth 2 (split-string path ";")) ""))

;; Adapted from org-find-entry-with-custom-id; there might be a better
;; way to do this
(defun org-find-entry-with-custom-id (ident)
  (let ((id ident)
        (case-fold-search nil))
    (save-excursion
      (save-restriction
        (widen)
        (goto-char (point-min))
        (when (re-search-forward
               (concat "^[ \t]*:CUSTOM_ID:[ \t]+" (regexp-quote id) "[ \t]*$")
               nil t)
          (org-back-to-heading t)
          (point))))))

(defun ox-bib-cite-follow (path)
  (let* ((files (ox-bib-get-bibliography))
	 (key (ox-bib-get-key path))
	 file done)
    (when ox-bib-find-entries-in-org-file
      (setq files
	    (mapcar (lambda (s)
		      (replace-regexp-in-string "\\.bib\\'" ".org" s))
		    files)))
    (while (and (car-safe files) (not done))
      (setq file (car-safe files))
      (when (file-exists-p file)
        (let* ((buf (find-file-noselect file))

	       (pos (with-current-buffer buf
		      (cond
		       (ox-bib-find-entries-in-org-file
			(org-find-entry-with-custom-id key))
		       (t (bibtex-search-entry key))))))
	  (when pos
	    (setq done (cons buf pos)))))
      (setq files (cdr files)))
    (if done
	(progn
          (pop-to-buffer (car done))
	  (goto-char (cdr done))
          (when ox-bib-find-entries-in-org-file
            (org-show-context)))
      (message "Citation not found: %s" key))))

(defun ox-bib-textcite-export (path desc format)
  (cond
   ((org-export-derived-backend-p format 'latex)
    (let ((key (ox-bib-get-key path desc))
	  (pre (ox-bib-get-pre path desc))
	  (post (ox-bib-get-post path desc)))
      (format "\\textcite[%s][%s]{%s}" pre post key)))
   (t
    (format "<textcite to: %s>" path))))

(defun ox-bib-parencite-export (path desc format)
  (cond
   ((org-export-derived-backend-p format 'latex)
    (let ((key (ox-bib-get-key path desc))
          (pre (ox-bib-get-pre path desc))
          (post (ox-bib-get-post path desc)))
      (format "\\parencite[%s][%s]{%s}" pre post key)))
   (t
    (format "<parencite to: %s>" path))))

(org-add-link-type "textcite" #'ox-bib-cite-follow #'ox-bib-textcite-export)
(org-add-link-type "parencite" #'ox-bib-cite-follow #'ox-bib-parencite-export)