import argparse
from coco.storage import Annotation, COCOStorage
import logging
import coloredlogs
import tornado.ioloop
import tornado.web
import tornado.websocket
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("server")


class JsonEncoder(json.JSONEncoder):
    def default(self, obj):
        method = getattr(obj, "forJson", None)
        if callable(method):
            return obj.forJson()
        # Let the base class default method raise the TypeError
        return json.JSONEncoder.default(self, obj)


class AnnotationHandler(tornado.web.RequestHandler):
    def initialize(self, storage: COCOStorage):
        self.storage = storage
        self.set_header("Content-Type", "application/json")

    def get(self, *params):
        self.write(json.dumps(self.getData(*params), cls=JsonEncoder))

    def getData(self) -> Annotation:
        # get specific annotation
        annotation_id = self.get_argument('id', None)
        annotation_id = None if not annotation_id else int(annotation_id)

        # get by category id
        category_id = self.get_argument('category', None)
        category_id = None if not category_id else int(category_id)

        normalise = self.get_argument('normalise', False)
        normalise = int(normalise) if normalise is not False else False

        min_area = self.get_argument('min_area', None)
        min_area = int(min_area) if min_area is not None else None

#         category_id = None if not category_id else int(category_id)

        logger.debug(
            f'Get annotation id: {annotation_id}, category: {category_id}, normalised: {normalise}, min_area: {min_area}')

        annotation = self.storage.getRandomAnnotation(
            annotation_id=annotation_id, category_id=category_id, min_area=min_area)
        if normalise:
            return annotation.getNormalised(normalise, normalise)
        return annotation

class AnnotationSvgHandler(AnnotationHandler):
    def initialize(self, storage: COCOStorage):
        self.storage = storage
        self.set_header("Content-Type", "image/svg+xml")

    def get(self, *params):
        annotation = self.getData(*params)
        dwg = annotation.asSvg(None, clip_image=True, image_dir='/dataset')
        xml = dwg.tostring()
        print(xml)
        self.write(xml)


class StaticFileWithHeaderHandler(tornado.web.StaticFileHandler):
    def set_extra_headers(self, path):
        """For subclass to add extra headers to the response"""
        if path[-5:] == '.html':
            self.set_header("Access-Control-Allow-Origin", "*")


def make_app(storage, debug):
    return tornado.web.Application([
        (r"/annotation.json", AnnotationHandler, {'storage': storage}),
        (r"/annotation.svg", AnnotationSvgHandler, {'storage': storage}),
        (r"/dataset/(.*)", StaticFileWithHeaderHandler,
         {"path": 'dataset', "default_filename": 'index.html'}),
        (r"/(.*)", StaticFileWithHeaderHandler,
         {"path": 'www', "default_filename": 'index.html'}),
    ], debug=debug)


if __name__ == "__main__":
    argParser = argparse.ArgumentParser(
        description='Server for COCO web interface')
    argParser.add_argument(
        '--port',
        '-P',
        type=int,
        default=8888,
        help='Port to listen on'
    )
    argParser.add_argument(
        '--db',
        type=COCOStorage,
        metavar='DATABASE',
        dest='storage',
        help='SQLite db filename, will be created if not existing',
        default='dataset/instances_val2017.db'
    )
    argParser.add_argument(
        '--verbose',
        '-v',
        action='store_true',
        help='Increase log level'
    )
    args = argParser.parse_args()

    loglevel = logging.DEBUG if args.verbose else logging.INFO
    coloredlogs.install(
        level=loglevel,
        fmt="%(asctime)s %(hostname)s %(name)s[%(process)d] %(levelname)s %(message)s"
    )

    app = make_app(args.storage, debug=args.verbose)
    app.listen(args.port)
    logger.info(f"Listening on {args.port}")
    tornado.ioloop.IOLoop.current().start()