#!/usr/bin/ruby

# Copyright (c) 2010 Geza Kovacs
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

require 'optparse'
require 'tmpdir'

def shell_escape(s)
    "'" + s.gsub("'", "'\\''") + "'"
end

def sh(c, *args)
  outl = []
  if args.length > 0
    c = shell_escape(c) + ' '
    c << args.map {|w| shell_escape(w)}.join(' ')
  end
  IO.popen(c) { |f|
    while not f.eof?
      tval = f.gets
      puts tval
      outl.push(tval)
    end
  }
  return outl.join("")
end

def writef(fn, c)
  File.open(fn, "w") { |f|
    f.puts(c)
  }
end

def rmdir(dirn)
  Dir.foreach(dirn) { |fn|
    if fn == "." or fn == ".."
      next
    end
    fn = File.expand_path(dirn+"/"+fn)
    if File.directory?(fn)
      rmdir(fn)
    else
      File.delete(fn)
    end
  }
  Dir.delete(dirn)
end

appname = 'pdfocr'
version = [0,3,0]
infile = nil
outfile = nil
deletedir = true
deletefiles = true
language = 'eng'
checklang = false
tmp = nil
useocropus = false
usecuneiform = false
usetesseract = false

optparse = OptionParser.new { |opts|
opts.banner = <<-eos
Usage: #{appname} -i input.pdf -o output.pdf
#{appname} adds text to PDF files using the ocropus, cuneiform, or tesseract OCR software
eos

  opts.on("-i", "--input [FILE]", "Specify input PDF file") { |fn|
    infile = fn
  }

  opts.on("-o", "--output [FILE]", "Specify output PDF file") { |fn|
    outfile = fn
  }

  opts.on("-t", "--tesseract", "Use tesseract as the OCR engine (default)") {
    usetesseract = true
  }

  opts.on("-c", "--cuneiform", "Use cuneiform as the OCR engine") {
    usecuneiform = true
  }

  opts.on("-p", "--ocropus", "Use ocropus as the OCR engine") {
    useocropus = true
  }

  opts.on("-l", "--lang [LANG]", "Specify language for the OCR software") { |fn|
    language = fn
    checklang = true
  }

  opts.on("-L", "--nocheck-lang LANG", "Suppress checking of language parameter") { |fn|
    language = fn
    checklang = false
  }

  opts.on("-w", "--workingdir [DIR]", "Specify directory to store temp files in") { |fn|
    deletedir = false
    tmp = fn
  }

  opts.on("-k", "--keep", "Keep temporary files around") {
    deletefiles = false
  }

  opts.on_tail("-h", "--help", "Show this message") {
    puts opts
    exit
  }

  opts.on_tail("-v", "--version", "Show version") {
    puts version.join('.')
    exit
  }

}

optparse.parse!(ARGV)

if not infile or infile == ""
  puts optparse
  puts
  puts "Need to specify an input PDF file"
  exit
end

if infile[-3..-1].casecmp("pdf") != 0
  puts "Input PDF file #{infile} should have a PDF extension"
  exit
end

if not File.file?(infile)
  puts "Input file #{infile} does not exist"
  exit
end

infile = File.expand_path(infile)

if not outfile or outfile == ""
  puts optparse
  puts
  puts "Need to specify an output PDF file"
  exit
end

if outfile[-3..-1].casecmp("pdf") != 0
  puts "Output PDF file should have a PDF extension"
  exit
end

if outfile == infile
  puts "Output PDF file should not be the same as the input PDF file"
  exit
end

if File.file?(outfile)
  puts "Output file #{outfile} already exists"
  exit
end

outfile = File.expand_path(outfile)

if not language or language == ""
  puts "Need to specify a language"
  exit
end

if `which pdfinfo` == ""
  puts "pdfinfo command is missing. Install the poppler-utils package"
  exit
end

if `which pdfseparate` == ""
  puts "pdfseparate command is missing. Install the poppler-utils package"
  exit
end

if `which pdftoppm` == ""
  puts "pdftoppm command is missing. Install the poppler-utils package"
  exit
end

if `which hocr-pdf` == ""
  puts "hocr-pdf command is missing. Install the hocr-tools package"
  exit
end

if `which xsltproc` == ""
  puts "xsltproc command is missing. Install the libxslt package"
  exit
end

if useocropus
  if `which ocroscript` == ""
    puts "The ocroscript command is missing. Install the ocropus package."
    exit
  end
elsif usecuneiform
  if `which cuneiform` == ""
    puts "The cuneiform command is missing. Install the cuneiform package."
    exit
  end
elsif usetesseract
  if `which tesseract` == ""
    puts "The tesseract command is missing. Install the tesseract-ocr package and the"
    puts "language packages you need, e.g. tesseract-ocr-deu, tesseract-ocr-deu-frak,"
    puts "or tesseract-ocr-eng."
    exit
  end
else
  if `which tesseract` != ""
    usetesseract = true
  elsif `which cuneiform` != ""
    usecuneiform = true
  elsif `which ocroscript` != ""
    useocropus = true
  else
    puts "The tesseract command is missing. Install the tesseract-ocr package and the"
    puts "language packages you need, e.g. tesseract-ocr-deu, tesseract-ocr-deu-frak,"
    puts "or tesseract-ocr-eng."
    exit
  end
end

if not deletedir
  if not File.directory?(tmp)
    puts "Working directory #{tmp} does not exist"
    exit
  else
    tmp = File.expand_path(tmp)+"/pdfocr"
    if File.directory?(tmp)
      puts "Directory #{tmp} already exists - remove it"
      exit
    else
      Dir.mkdir(tmp)
    end
  end
else
  tmp = Dir.mktmpdir
end

if checklang
  langlist = []
  if usecuneiform
    begin
      langlist = `cuneiform -l`.split("\n")[-1].split(":")[-1].delete(".").split(" ")
    rescue
      puts "Unable to list supported languages from cuneiform"
    end
  end
  if usetesseract
    begin
      langlist = `tesseract --list-langs 2>&1`.split("\n")[1..-1]
    rescue
      puts "Unable to list supported languages from tesseract"
    end
  end
  if langlist and not langlist.empty?()
    if not langlist.include?(language)
      puts "Language #{language} is not supported or not installed. Please choose from"
      puts langlist.join(' ')
      exit
    end
  end
end

#Find the hocr fixing xslt file
if File.file?(File.expand_path('fix-hocr.xsl'))
  fixhocrpath = File.expand_path('fix-hocr.xsl')
elsif File.file?("/etc/pdfocr/fix-hocr.xsl")
  fixhocrpath = "/etc/pdfocr/fix-hocr.xsl"
else
  puts "Error: fix-hocr.xsl not found"
  exit
end

puts "Input file is #{infile}"
puts "Output file is #{outfile}"
puts "Using working dir #{tmp}"

puts "Getting info from PDF file"

puts

pdfinfo = sh "pdfinfo", infile

if not pdfinfo or pdfinfo == ""
  puts "Error: didn't get info from pdfinfo #{infile}"
  exit
end

puts

begin
  pdfinfo =~ /Pages: \s*(\d+)/
  pagenum = $1.to_i
rescue
  puts "Error: didn't get page count for #{infile} from pdfinfo"
  exit
end

if pagenum == 0
  puts "Error: there are 0 pages in the input PDF file #{infile}"
  exit
end

writef(tmp+"/pdfinfo.txt", pdfinfo)

puts "Converting #{pagenum} pages"

Dir.chdir(tmp+"/") {
  1.upto(pagenum) { |i|
    puts "=========="
    puts "Extracting page #{i}"
    basefn = i.to_s
    interfn = basefn + "-1"
    puts basefn+'.jpg'

    sh "pdfseparate", "-f", "#{i}", "-l", "#{i}", infile, basefn+'.pdf'
    if not File.file?(basefn+'.pdf')
      puts "Error while extracting page #{i}"
      next
    end
    puts "Converting page #{i} to jpg"

    sh "pdftoppm -r 300 #{shell_escape(basefn)}.pdf #{shell_escape(basefn)} -jpeg"
    if not File.file?(interfn+'.jpg')
      puts "Error while converting page #{i} to jpg"
      next
    end
    puts "Running OCR on page #{i}"

    if usetesseract
      #Tesseract can go straight to PDF
      sh "tesseract", interfn+'.jpg', interfn+'-new', "-l", language, "pdf"

      puts "Merging together PDF files"
      sh "pdfunite *-new.pdf #{outfile}"
    else
      if usecuneiform
        sh "cuneiform", "-l", language, "-f", "hocr", "-o", interfn+'.hocr', interfn+'.jpg'
      else
        sh "ocroscript recognize #{shell_escape(interfn)}.jpg > #{shell_escape(interfn)}.hocr"
      end

      if not File.file?(interfn+'.hocr')
        puts "Error while running OCR on page #{i}"
        next
      else
        #Workaround for hocr issues:
        #https://bugs.launchpad.net/cuneiform-linux/+bug/623438
        #https://bugs.launchpad.net/cuneiform-linux/+bug/623438/comments/60
        sh "xsltproc", "-html", "-nonet", "-novalid", "-o", interfn+'.hocr', fixhocrpath, interfn+'.hocr'

        sh "hocr-pdf . > #{outfile}"
      end
    end
  }
}

if deletefiles
  puts "Cleaning up temporary files"
  rmdir(tmp)
end
