#!python

"""
invisible-watermark is the utility to add hidden watermark and extract watermark from image
"""

import argparse
import os
import shutil
import sys
import time
import yaml
import uuid
import cv2
import base64
import time
import pprint

from imwatermark import WatermarkEncoder, WatermarkDecoder

pp = pprint.PrettyPrinter(indent=2)


class CustomArgParseFormatter(
    argparse.ArgumentDefaultsHelpFormatter,
    argparse.RawDescriptionHelpFormatter):
  """A custom formatter that combines the features of multiple base classes.

  This gives us defaults for each argument in the help text, plus it preserves
  whitespace in the description field.
  """
  pass


def main():
    description = __doc__.format()
    parser = argparse.ArgumentParser(description=description,
                                   formatter_class=CustomArgParseFormatter)

    parser.add_argument('-a', '--action', required=True, help='encode|decode')
    parser.add_argument('-t', '--type', default='bits', help='watermark type')
    parser.add_argument('-m', '--method', default='maxDct', help='watermark algorithm')
    parser.add_argument('-w', '--watermark', default='', help='watermark text')
    parser.add_argument('-l', '--length', default=0, type=int,
                        help='watermark length, only valid for bytes watermark')
    parser.add_argument('input',
                        help='The path of input')
    parser.add_argument('-o', '--output', required=False,
                        help='The path of output')

    args = parser.parse_args()
    action = args.action
    inputFile = args.input

    if action == 'encode':
        if not args.output:
            sys.stderr.write('output is required. exiting...\n')
            sys.exit(1)
        if not args.watermark:
            sys.stderr.write('watermark is required. exiting...\n')
            sys.exit(1)

        encoder = WatermarkEncoder()
        wmType = args.type
        if args.method == 'rivaGan':
            WatermarkEncoder.loadModel()
        bgr = cv2.imread(args.input)
        if args.type == 'bytes':
            wm = args.watermark.encode('utf-8')
        elif args.type == 'b16':
            wm = args.watermark.upper().encode('utf-8')
        elif args.type == 'bits':
            wm = [int(c) % 2 for c in args.watermark]
        else:
            wm = args.watermark
        encoder.set_watermark(wmType, wm)
        bgr_encoded = encoder.encode(bgr, args.method)
        cv2.imwrite(args.output, bgr_encoded)
    elif action == 'decode':
        if args.type in ['bytes', 'bits', 'b16']:
            if args.length <= 0:
                sys.stderr.write('length is required for bytes watermark decoding\n')
                sys.exit(1)
            wmType = args.type
            decoder = WatermarkDecoder(wmType, args.length)
        else:
            decoder = WatermarkDecoder(args.type)
        if args.method == 'rivaGan':
            WatermarkDecoder.loadModel()
        bgr = cv2.imread(args.input)
        wm = decoder.decode(bgr, args.method)
        if args.type in ['bytes', 'b16']:
            wm = wm.decode('utf-8')
        print(wm)
    else:
        sys.stderr.write('unknown action\n')
        sys.exit(1)

if __name__ == '__main__':
    sys.exit(main())
