#!/usr/bin/python
'''
Copyright (c) 2017 noriah vix@noriah.dev

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 isfurnished
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.
'''

import sys
import re
import argparse
from PyVDF import PyVDF

def FlattenArray(array):
  out = list()
  for a in array.keys():
    b = a
    if re.search("\.", b) != None: b = "[" + b + "]"
    if re.search(" ", b) != None: b = '"' + b + '"'
    if isinstance(array[a], dict):
      out.append(b)
      out += ['{}.{}'.format(b,i) for i in FlattenArray(array[a])]
    else: out.append('{}="{}"'.format(b,array[a]))
  return out

def main():

  parser = argparse.ArgumentParser(prog="vdfedit",
                                   description="Read and Write VDF KeyValue Files")

  parser.add_argument('infile', type=str,
                      help="The file to Read",
                      metavar="FILE")

  parser.add_argument('-o', '--out', type=str,
                      help="The file to Write out to Defaults to the infile",
                      metavar="FILE", dest="outfile")
  
  parser.add_argument('-p', '--print', action='store_true',
                      help="Print the file as groups of Paths.Key=Value",
                      dest="print_flat")

  parser.add_argument('-g', '--get', action="append", type=str,
                      help="Add a key to Search For",
                      metavar="key", dest="gets")

  parser.add_argument('-s', '--set', action="append", type=str,
                      help="Add/set a key-value pair in the vdf",
                      metavar="key=value", dest="sets")

  parser.add_argument('-w', '--write', action='store_true',
                      help='Write out the data regardless of setting a value.',
                      dest="write_file")

  parser.add_argument('-c', '--check', action='store_true',
                      help='Check the file for validity',
                      dest="check_file")

  parser.add_argument('-len', '--token-length', default=1200, type=int,
                      help="Set the maximum token length. Larger values cause slower searching, however smaller values can result in errors. Defaults to 1200",
                      metavar="int", dest="tokenLength")

  parser.add_argument('-d', '--delim', default=",", type=str,
                      help="Set the string used to separate found values. Defaults to ','",
                      metavar="str", dest="delim")

  parser.add_argument('--indent', default='\t', type=str,
                      help="Set the indention used when writing out a file. Defaults to \\t",
                      metavar="str", dest="indent")

  parser.add_argument('--spacing', default='\t\t', type=str,
                      help="Set the spacing used when writing out a file. Defaults to \\t\\t",
                      metavar="str", dest="spacing")

  parser.add_argument('--condense', '--condensed', '--use-condensed', action='store_true',
                      help='Use consensed output when writing a file',
                      dest="use_condensed")

  parser.add_argument('--fast', '--faster-reading', action='store_true',
                      help='''Use faster reading. **Note** If you write a file
                      while this option is set, your output will not be in the
                      same order as the original file''',
                      dest="use_speed")

  if not sys.stdin.isatty():
    import shlex
    var2 = sys.stdin.read()
    var2 = re.sub('\\n', ' ', var2)
    sys.argv += shlex.split(var2)

  args = parser.parse_args()

  PyVDF.setCondensed(args.use_condensed)
  PyVDF.setIndentation(args.indent.replace("\\t", "\t"))
  PyVDF.setSpacing(args.spacing.replace("\\t", "\t"))
  PyVDF.useFastDict(args.use_speed)
  PyVDF.setMaxTokenLength(args.tokenLength)

  args.outfile = args.outfile if args.outfile else args.infile if args.infile else None

  if args.gets is not None or args.sets is not None or args.print_flat or args.write_file or args.check_file:
    vdf, gets, sets = PyVDF(), args.gets, args.sets
    vdf.load(args.infile)

    if args.print_flat and args.gets is not None:
      print(" --get cannot be accompanied by --print")
      return
      
    if args.gets is not None:
      gets = vdf.findMany(args.gets)
      for g in gets:
        if isinstance(g, dict):
          gets.remove(g)
          print("\n".join(FlattenArray(g)).decode("string-escape"))
      if len(gets):
        print(args.delim.join(gets).decode("string-escape"))
    
    if args.sets is not None:
      sets = list()
      setsa = sets.append
      for string in args.sets:
        try:
          setsa((string.split("=", 1)))
        except ValueError:
          parser.print_usage()
          print("vdfedit: error: set statement missing value\n '{}'".format(string))
          return
      vdf.editMany(sets)
    
    if args.sets is not None or args.write_file:
      vdf.write_file(args.outfile)

    if args.print_flat:
      print("\n".join(FlattenArray(vdf.getData())))

  else:
    parser.print_usage()
    print("vdfedit: error: at least one of -p, -g, -s, -w or -c required")
    return


if __name__ == "__main__":
  main()
