#  Copyright 2023-2024 Amazon.com, Inc. or its affiliates.

import argparse
import logging
from secrets import token_hex
from typing import Dict, List, Optional, Union

from flask import Flask
from osgeo import gdal

# Enable exceptions for GDAL
gdal.UseExceptions()


def setup_server(app: Flask):
    """
    The assumption is that this script will be the ENTRYPOINT for the inference
    container. SageMaker will launch the container with the "serve" argument. We
    also have the option of using multiple models from this single container;
    only one model will be active at a time (i.e., this is not a Multi Model Server),
    so it can be selected by name using the "model" parameter.

    :param app: The flask application to set up
    :return: None
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("command", nargs="?", default="serve")
    parser.add_argument("-v", "--verbose", default=False)
    args = parser.parse_args()

    # Set up the logging for the Flask application
    # and log the startup information
    configure_logging(args.verbose)
    logging.info("Initializing REST Model Server...")
    for arg, value in sorted(vars(args).items()):
        logging.info("Argument %s: %r", arg, value)

    # Start the simple web application server using Waitress.
    # Flask's app.run() is only intended to be used in development
    #  mode, so this provides a solution for hosting the application.
    from waitress import serve

    serve(app, host="0.0.0.0", port=8080, clear_untrusted_proxy_headers=True)


def configure_logging(verbose: bool = False) -> None:
    """
    Configure application logging. Note the timestamp for the
    log record is available in CloudWatch if this is used in a
    non-AWS environment, you can add %(asctime)s to the start of
    the format string.

    :param verbose: True, if the DEBUG log level should be used, defaults to False
    :return: None
    """
    logging_level = logging.INFO
    if verbose:
        logging_level = logging.DEBUG

    logging.basicConfig(level=logging_level, format="%(levelname)-8s : %(message)s")


def detect_to_feature(
    fixed_object_bbox: List[float],
    fixed_object_mask: Optional[List[List[float]]] = None,
    detection_score: Optional[float] = 1.0,
    detection_type: Optional[str] = "sample_object",
) -> Dict[str, Union[str, list]]:
    """
    Converts the bbox object into a sample GeoJSON formatted detection.
    The model container does not normally provide the world coordinates,
    so they're defaulted to 0,0 here since GeoJSON features require a geometry.

    :param detection_type: Class assigned to the detection
    :param detection_score: Confidence score assigned to the detection
    :param fixed_object_mask: Polygon version of mask generated by detectron2
    :param fixed_object_bbox: Bounding box to transform into a geojson feature
    :return: dict: Dictionary representation of a geojson feature
    """
    feature = {
        "type": "Feature",
        "geometry": {"coordinates": [0.0, 0.0], "type": "Point"},
        "id": token_hex(16),
        "properties": {
            "bounds_imcoords": fixed_object_bbox,
            "detection_score": detection_score,
            "feature_types": {detection_type: detection_score},
            "image_id": token_hex(16),
        },
    }

    if fixed_object_mask is not None:
        feature["properties"]["geom_imcoords"] = fixed_object_mask

    return feature
