#!python

import argparse
import os
import io
import boto3
import json
import sys
from textractcaller.t_call import Textract_Features, Textract_Types, call_textract
from textractoverlayer.t_overlay import DocumentDimensions, get_bounding_boxes
from textractprettyprinter.t_pretty_print import Pretty_Print_Table_Format, Textract_Pretty_Print, get_string
from textractoverlayer.image_tools import get_filename_from_document
from textracthelper._version import __version__
import textracthelper._version
import logging

logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser()
input_doc_or_example_or_stdin = parser.add_mutually_exclusive_group(
    required=True)
input_doc_or_example_or_stdin.add_argument(
    "--input-document", help="s3 object (s3://) or file from local filesystem")
input_doc_or_example_or_stdin.add_argument(
    "--example",
    dest='use_example',
    action='store_true',
    help="using the example document to call Textract")
input_doc_or_example_or_stdin.set_defaults(use_example=False)
input_doc_or_example_or_stdin.add_argument("--stdin",
                                           dest='inputfromstdin',
                                           action='store_true',
                                           help="receive JSON from stdin")
input_doc_or_example_or_stdin.set_defaults(inputfromstdin=False)

parser.add_argument(
    "--features",
    nargs='+',
    type=str,
    choices=["FORMS", "TABLES"],
    help=
    "features to call Textract with. Will trigger call to AnalyzeDocument instead of DetectDocumentText"
)
parser.add_argument("--pretty-print",
                    nargs='+',
                    type=str,
                    choices=["WORDS", "LINES", "FORMS", "TABLES"],
                    help="")
parser.add_argument(
    "--pretty-print-table-format",
    type=str,
    choices=[
        "csv", "plain", "simple", "github", "grid", "fancy_grid", "pipe",
        "orgtbl", "jira", "presto", "pretty", "psql", "rst", "mediawiki",
        "moinmoin", "youtrack", "html", "unsafehtml", "latex", "latex_raw",
        "latex_booktabs", "latex_longtable", "textile", "tsv"
    ],
    default='github',
    help=
    "which format to output the pretty print information to. Only effects FORMS and TABLES"
)
parser.add_argument(
    "--overlay",
    nargs='+',
    type=str,
    choices=["WORD", "LINE", "FORM", "KEY", "VALUE", "TABLE", "CELL"],
    help="defines what bounding boxes to draw on the output")
parser.add_argument("--pop-up-overlay-output",
                    dest='showoutput',
                    action='store_true',
                    help="shows image with overlay")
parser.set_defaults(showoutput=False)
parser.add_argument("--overlay-output-folder",
                    type=str,
                    default=None,
                    help="output with bounding boxes to folder")
parser.add_argument("--version",
                    action='version',
                    version='%(prog)s {version}'.format(version=__version__),
                    help="print version information")
parser.add_argument("--no-stdout",
                    dest='showstdout',
                    action='store_false',
                    help="no output to stdout")
parser.set_defaults(showstdout=True)
show_logs = parser.add_mutually_exclusive_group(required=False)
show_logs.add_argument("-v",
                       dest='showinfo',
                       action='store_true',
                       help=">=INFO level logging output to stderr")
show_logs.set_defaults(showinfo=False)
show_logs.add_argument("-vv",
                       dest='showdebug',
                       action='store_true',
                       help=">=DEBUG level logging output to stderr")
show_logs.set_defaults(showdebug=False)

args = parser.parse_args()
input_document: str = args.input_document
features_arg = args.features
pretty_print_arg = args.pretty_print
overlay_output_folder = args.overlay_output_folder
overlay_arg = args.overlay
showoutput = args.showoutput
showstdout = args.showstdout
use_example = args.use_example
showdebug = args.showdebug
showinfo = args.showinfo
use_stdin = args.inputfromstdin
pretty_print_table_format_arg = args.pretty_print_table_format

if showinfo or showdebug:
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    handler.setLevel(logging.INFO)
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)
    caller_logger = logging.getLogger('textractcaller')
    caller_logger.propagate = True
    caller_logger.setLevel(logging.INFO)
    caller_logger.addHandler(handler)
    printer_logger = logging.getLogger('textractprettyprinter')
    printer_logger.propagate = True
    printer_logger.setLevel(logging.INFO)
    printer_logger.addHandler(handler)
    logger.info("log level: INFO")

    if showdebug:
        handler.setLevel(logging.DEBUG)
        logger.setLevel(logging.DEBUG)
        caller_logger.setLevel(logging.DEBUG)
        printer_logger.setLevel(logging.DEBUG)
        logger.debug("current log level: DEBUG")

if use_example:
    SCRIPT_DIR = os.path.dirname(
        os.path.abspath(textracthelper._version.__file__))
    input_document = os.path.join(SCRIPT_DIR, "examples/employmentapp.png")

features = None
if features_arg:
    features = [Textract_Features[x] for x in features_arg]
pretty_print_table_format = Pretty_Print_Table_Format[
    pretty_print_table_format_arg]
logger.debug(f"pretty_print_table_format: {pretty_print_table_format}")
pretty_print = None
if pretty_print_arg:
    pretty_print = [Textract_Pretty_Print[x] for x in pretty_print_arg]

if overlay_arg:
    overlay = [Textract_Types[x] for x in overlay_arg]
else:
    overlay = list()

exit_code = 0
if not use_stdin and pretty_print_arg and "FORMS" in pretty_print_arg and (
        not features_arg or "FORMS" not in features_arg):
    print("should pretty-print FORMS but is not requested as --features")
    exit_code = 1
if not use_stdin and pretty_print_arg and "TABLES" in pretty_print_arg and (
        not features_arg or "TABLES" not in features_arg):
    print("should pretty-print TABLES but is not requested as --features")
    exit_code = 1
if not use_stdin and overlay_arg and any([
        x for x in ["CELL", "TABLE"] if x in overlay_arg
]) and (not features_arg or "TABLES" not in features_arg):
    print(
        "should overlay TABLE or CELL but is not requested as --features TABLES"
    )
    exit_code = 1
if not use_stdin and showoutput and overlay_arg and any([
        x for x in ["FORM", "KEY", "VALUE"] if x in overlay_arg
]) and (not features_arg or "FORMS" not in features_arg):
    print(
        "should overlay FORM or KEY or VALUE but FORMS not requested as --features FORMS"
    )
    exit_code = 1
if showoutput and not overlay_arg:
    print(
        "asked for --pop-up-overlay-output or --overlay-output-folder, but not --overlay arguments given"
    )
    exit_code = 1
if exit_code > 0:
    exit(1)

logger.debug(f"calling Textract")
if use_stdin:
    doc = json.load(sys.stdin)
else:
    doc = call_textract(input_document=input_document, features=features)
logger.debug(f"receivedTextract response")
if showstdout:
    if pretty_print:
        print(
            get_string(textract_json=doc,
                       output_type=pretty_print,
                       table_format=pretty_print_table_format))
    else:
        print(json.dumps(doc))

if overlay_output_folder or showoutput:
    colors = {
        "WORD": (128, 128, 0),
        "LINE": (0, 128, 128),
        "FORM": (128, 0, 128),
        "KEY": (255, 0, 0),
        "VALUE": (0, 255, 255),
        "CELL": (255, 0, 255),
        "TABLE": (255, 255, 0)
    }
    from PIL import Image, ImageDraw
    is_s3_document: bool = False
    s3_bucket: str = ""
    s3_key: str = ""
    if len(input_document) > 7 and input_document.lower().startswith("s3://"):
        is_s3_document = True
        input_document = input_document.replace("s3://", "")
        s3_bucket, s3_key = input_document.split("/", 1)
        print(s3_bucket, s3_key)
    ext: str = ""
    _, ext = os.path.splitext(input_document)

    if is_s3_document:
        s3 = boto3.client('s3')
        o = s3.get_object(Bucket=s3_bucket, Key=s3_key)
        input_bytes = o.get('Body').read()
        f = io.BytesIO(input_bytes)
        image: Image.Image = Image.open(f)
        file_name, suffix = os.path.splitext(os.path.basename(s3_key))
    else:
        image = Image.open(input_document)
        file_name, suffix = os.path.splitext(os.path.basename(input_document))
    rgb_im = image.convert('RGB')
    draw = ImageDraw.Draw(rgb_im)
    document_dimension: DocumentDimensions = DocumentDimensions(
        doc_width=image.size[0], doc_height=image.size[1])
    bounding_box_list = get_bounding_boxes(
        textract_json=doc,
        document_dimensions=document_dimension,
        overlay_features=overlay)
    for bbox in bounding_box_list:
        box_color = colors[bbox.box_type.name]
        draw.rectangle(xy=[bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax],
                       outline=box_color,
                       width=2)

    if showoutput:
        rgb_im.show()
    if overlay_output_folder:
        if not os.path.exists(overlay_output_folder):
            os.makedirs(overlay_output_folder, exist_ok=True)
        file_name, suffix = get_filename_from_document(
            input_document=input_document)
        output_types = "_".join(overlay_arg)
        output_file_name = os.path.join(
            overlay_output_folder,
            f"{file_name}_boxed_{output_types}_{suffix}")
        output_format = suffix[1:]
        if output_format.upper() == "JPG":
            output_format = "JPEG"
        rgb_im.save(output_file_name, output_format)
