#! /usr/bin/env python3

from argparse import ArgumentParser
import tempfile
import logging
import os
import sys
import shutil
import re
import subprocess
import base64
import pathlib

__version__ = '2.2'

# a text-list class for accumulating lines of text
class tlist(list):
    def __iadd__(self, item):
        if isinstance(item, list):
            for i in item:
                self.append(i)
        else:
            self.append(item)
        return self


parser = ArgumentParser(
    description="A utility for creating image files from latex snippets..")

parser.add_argument("latex_snippet_or_file",
                    action="store",
                    nargs="*",
                    default=[],
                    help="LaTeX snippet or filename conotainting latex code to compile.")

parser.add_argument("--version",
                    action="store_true",
                    help="Show version.")

parser.add_argument("-v", "--verbose",
                    action="store_true",
                    help="Log information.")

parser.add_argument("-d", "--debug",
                    action="store_true",
                    help="Log debug information.")

parser.add_argument("-k", "--keep-files",
                    action="store_true",
                    help="Keep generated files.")



parser.add_argument("-s", "--font-size",
                    action="store",
                    type=int,
                    default=12,
                    help="Document font size.")

parser.add_argument("-t", "--text-color",
                    action="store",
                    default="black",
                    help="Text color. Color support provided by xcolor. Arguments may be a pre-defined xcolor color, such as 'blue', or specify one of the supported xcolor modes as 'mode:arg', such 'HTML:FF7F00'. See xcolor package for supported modes. ")

parser.add_argument("-b", "--background-color",
                    action="store",
                    default="white",
                    help="Background color. See --text-color for argument types.")

parser.add_argument("-x", "--preamble",
                    action="store",
                    default="",
                    help="Read additional preamble lines from file.")

parser.add_argument("-n", "--no-equation-environment",
                    action="store_true",
                    help="Do not wrap snippet in equation environment. Useful if you want to use your own environment.")

parser.add_argument("-C", "--no-comment",
                    action="store_true",
                    help="Do not embed latex snippet into comment field of output image. Sometimes `convert` does not have 'permission' to uset the comment field.")


parser.add_argument("-a", "--anti-aliasing",
                    action="store",
                    type=int,
                    default=-1,
                    help="Turn anti-aliasing on/off. If -1, anti-aliasing will be turned off if for transparent background, and on otherwise.")

parser.add_argument("-B", "--border",
                    action="store",
                    default=0,
                    help="Size of border, in pixels.")

parser.add_argument("-r", "--resolution",
                    action="store",
                    default="150x150",
                    help="Size of border, in pixels.")

parser.add_argument("-z", "--transparent-background",
                    action="store_true",
                    help="Make background transparent.")

parser.add_argument("-f", "--output-format",
                    action="store",
                    default="png",
                    help="Output image file format. This can be any file extension that ImageMagick's convert command knows about.")

parser.add_argument("-o", "--output-basename",
                    action="store",
                    help="Basename to use for output filename. Image format extension will be appended.")

args = parser.parse_args()

if args.version:
  print("version:",__version__)
  sys.exit()

if args.debug:
    logging.basicConfig(level=logging.DEBUG)
elif args.verbose:
    logging.basicConfig(level=logging.INFO)
else:
    logging.basicConfig(level=logging.ERROR)

def tex2im(args,i=0):
  # create a tmp dir to work in
  curdir = os.getcwd()
  tmpdir = tempfile.mkdtemp(prefix="tex2im-")
  logging.info("created {}".format(tmpdir))


  # check if snippet is actually a filename or not
  latex_snippet = args.latex_snippet_or_file[i]
  input_file = None
  if os.path.isfile(latex_snippet):
      input_file = os.path.abspath(latex_snippet)
      logging.info("{} is a file, reading now".format(input_file))
      with open(input_file) as f:
          # read lines, but remove comments
          lines = filter(lambda x: not re.match('\s*#', x), f.readlines())
          latex_snippet = ''.join(lines).strip()

  logging.info("LaTeX snippet: {}".format(latex_snippet))


  logging.info("creating temp directory")


  os.chdir(tmpdir)


  latex_lines = tlist()

  # preamble
  latex_lines += r'\documentclass[%dpt]{article}' % args.font_size
  latex_lines += r'\usepackage{xcolor}'
  latex_lines += r'\pagestyle{empty}'

  toks = args.text_color.split(':')
  if len(toks) == 1:
    latex_lines += r'\colorlet{text_color}{%s}' % args.text_color
  else:
    latex_lines += r'\definecolor{text_color}{%s}{%s}' % (*toks,)

  toks = args.background_color.split(':')
  if len(toks) == 1:
    latex_lines += r'\colorlet{background_color}{%s}' % args.background_color
  else:
    latex_lines += r'\definecolor{background_color}{%s}{%s}' % (*toks,)

  latex_lines += r'\pagecolor{background_color}'

  # look for a preamble file to load
  for preamble in [args.preamble, os.path.join(curdir, '.tex2im_preamble'), os.path.join(os.getenv("HOME"), '.tex2im_preamble'), os.path.join(os.getenv("HOME"), '.tex2im_header')
                   ]:
      if os.path.isfile(preamble):
          with open(preamble, 'r') as f:
              latex_lines += f.readlines()
          break


  # document
  latex_lines += r'\begin{document}{'
  latex_lines += r'\color{text_color}'
  if not args.no_equation_environment:
      latex_lines += r'\begin{eqnarray*}'

  # snippet
  latex_lines += latex_snippet

  # end document
  if not args.no_equation_environment:
      latex_lines += r'\end{eqnarray*}'
  latex_lines += r'}\end{document}'


  logging.info("writing LaTeX document to {}/out.tex".format(tmpdir))
  with open('out.tex', 'w') as f:
      f.write('\n'.join(latex_lines))

  # TODO: detect image files in the latex snippet make them available in the tempdir
  latex_cmd = "pdflatex -interaction=nonstopmode out.tex"
  logging.info("running LaTeX: {}".format(latex_cmd))
  res = subprocess.run( latex_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT )
  if res.returncode:
    logging.error("LaTeX command returned non-zero status")
    logging.error("LaTeX command: {}".format(latex_cmd))
    logging.error("LaTeX input (out.tex):")
    logging.error("\n".join(latex_lines))
    logging.error("LaTeX output:")
    logging.error(res.stdout.decode('utf-8'))



  # convert to ouput image format
  convert_cmd = "convert -trim -border {BORDER} -bordercolor {BORDERCOLOR} +adjoin -density {DENSITY}".format(BORDER=args.border
                                                                                                             ,BORDERCOLOR=args.background_color
                                                                                                             ,DENSITY=args.resolution
                                                                                                             )
  if args.anti_aliasing == -1 and args.transparent_background:
    convert_cmd += " +antialias"
  elif args.anti_aliasing == -1 and not args.transparent_background:
    convert_cmd += " -antialias"
  elif args.anti_aliasing == 0:
    convert_cmd += " +antialias"
  else:
    convert_cmd += " -antialias"

  if args.transparent_background:
    convert_cmd += " -transparent {}".format(args.background_color)




  # determine output image filename

  output_basename = pathlib.Path(args.output_basename) if args.output_basename is not None else None
  output_dir = pathlib.Path(curdir)

  # check if output_basename is actually a directory
  if output_basename is not None:
    if not output_basename.is_absolute():
      output_basename = output_dir / output_basename

    if output_basename.is_dir():
      output_dir = pathlib.Path(output_basename)
      output_basename = None

  if output_basename is None:
    if input_file is not None:
      output_basename,_ = os.path.splitext( input_file )
    else:
      output_basename = "out"



  if input_file is None and len(args.latex_snippet_or_file) > 1:
    output_basename = str(output_basename)+"-%d"%i

  output_ext = args.output_format
  if output_ext.lower() == "html":
    output_ext = "png"
  if not output_ext.startswith('.'):
    output_ext = "."+output_ext

  if not output_dir.exists():
    pass

  output_file = (output_dir / output_basename).with_suffix(output_ext)

  # # check for repeated extension.
  # if output_file.endswith("{0}.{0}".format(output_ext)):
  #   output_file,_ = os.path.splitext( output_file )






  # split output file name into directory a filename parts
  output_dir,output_file = os.path.split(output_file)
  if output_dir == "":
    output_dir = curdir

  # add comment containing LaTeX
  comment_lines = tlist()
  comment_lines += "# This image was created with tex2im (https://github.com/CD3/tex2im)"
  comment_lines += latex_snippet.replace('\\','\\\\')
  comment_lines += ""
  if not args.no_comment:
    convert_cmd += " -comment '{COMMENT}'".format(COMMENT='\n'.join(comment_lines))

  # convert image create by LaTeX to format requested
  convert_cmd += " out.pdf"
  convert_cmd += " {}".format(output_file)

  logging.info("running ImageMagick: {}".format(convert_cmd))
  res = subprocess.run( convert_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT )
  if res.returncode:
    logging.error("convert command returned non-zero status")
    logging.error("convert command: {}".format(convert_cmd))
    logging.error("convert output:")
    logging.error(res.stdout.decode('utf-8'))

  # if output format was html, convert the png we created
  # to an embedded html tag
  if args.output_format == "html":
    with open( output_file, 'rb' ) as f:
      text = f.read()
    code = base64.b64encode(text).decode('utf-8')
    text  = r'''<img src="data:image/{fmt};base64,{code}" {opts}>'''.format(fmt="png",code=code,opts="")
    basename,_ = os.path.splitext( output_file )
    output_file = "{}.{}".format(basename,"html")
    with open( output_file, 'w' ) as f:
      f.write(text)


  # write latex code to image metadata
  # comment_lines = tlist()
  # comment_lines += "# This image was created with tex2im (https://github.com/CD3/tex2im)"
  # comment_lines += latex_snippet
  # exiftool_cmd = "exiftool -comment='{COMMENT}' {FILE}".format(COMMENT='<newline>'.join(comment_lines),FILE=output_file)
  # res = subprocess.run( exiftool_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT )
  # if res.returncode:
    # logging.error("exiftool command returned non-zero status")
    # logging.error("exiftool command: {}".format(exiftool_cmd))
    # logging.error("exiftool output:")
    # logging.error(res.stdout.decode('utf-8'))

  # copy image to working directory
  logging.info("copying file back to {}".format(os.path.join(curdir,output_file)))
  shutil.copyfile(output_file, os.path.join(output_dir,output_file))


  os.chdir(curdir)
  if args.keep_files:
      print("Generated files were written to {}.".format(tmpdir))
  else:
      shutil.rmtree(tmpdir)
      logging.info("deleted {}".format(tmpdir))


if __name__ == "__main__":
  # TODO: call use multiprocessing to convert multiple images in parallel
  # will need to rewrite the output file naming code.
  for i in range(len(args.latex_snippet_or_file)):
    tex2im(args,i)

