import json
import logging
import os
import shutil
from urllib.error import HTTPError
import tornado.ioloop
import tornado.web
import tornado.websocket
from urllib.parse import urlparse
import uuid
import datetime
import html
import argparse
import coloredlogs
import glob
import filelock
import svganim.strokes
import svganim.uimethods
import cairosvg


logger = logging.getLogger("svganim.webserver")


class DateTimeEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, datetime.datetime):
            return o.isoformat(timespec="milliseconds")

        return super().default(self, o)


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", "*")
        if path[-4:] == ".svg":
            self.set_header("Content-Type", "image/svg+xml")


class WebSocketHandler(tornado.websocket.WebSocketHandler):
    """
    Websocket from the drawing
    """

    # CORS_ORIGINS = ['localhost']
    connections = set()

    def initialize(self, config):
        self.config = config
        self.strokes = []
        self.hasWritten = False
        self.prev_file = None
        self.prev_file_duration = 0
        self.dimensions = [None, None]

    # def check_origin(self, origin):
    #     parsed_origin = urlparse(origin)
    #     # parsed_origin.netloc.lower() gives localhost:3333
    #     valid = any([parsed_origin.hostname.endswith(origin) for origin in self.CORS_ORIGINS])
    #     return valid

    # the client connected
    def open(self, p=None):
        self.__class__.connections.add(self)
        self.prefix = datetime.datetime.now().strftime("%Y-%m-%d-")
        self.filename = (
            self.prefix + str(self.check_filenr()) + "-" + uuid.uuid4().hex[:6]
        )
        logger.info(f"{self.filename=}")
        self.write_message(json.dumps({"filename": self.filename}))

    def check_filenr(self):
        files = glob.glob(os.path.join(self.config.storage, self.prefix + "*"))
        return len(files) + 1

    def appendEvent(self, row):
        if not self.hasWritten and self.prev_file and 'event' in row and row['event'] == 'viewbox':
            # ignore canvas movement after
            return

        # write to an appendable json format. So basically a file that should be wrapped in [] to be json-parsable
        # TODO use jsonlines -- which is not so much different but (semi-)standardized
        with open(
            os.path.join(self.config.storage, self.filename +
                         ".json_appendable"), "a"
        ) as fp:
            if not self.hasWritten:
                if self.prev_file:
                    # TODO WIP
                    with open(
                        self.prev_file, 'r'
                    ) as fprev:
                        wrote = False
                        for line in fprev:
                            wrote = True
                            fp.write(line)
                        if wrote:
                            fp.write(",\n")

                # metadata to first row, but only on demand
                fp.write(
                    json.dumps(
                        [
                            datetime.datetime.now().strftime("%Y-%m-%d %T"),
                            self.dimensions[0],
                            self.dimensions[1],
                        ]
                    )
                )
                # writer.writerow()
                self.hasWritten = True

            fp.write(",\n")
            # first column is color, rest is points
            fp.write(json.dumps(row))

    def preloadFile(self, file):
        if self.hasWritten:
            logger.error("Cannot preload when already written content")
            return False

        logger.info(f"load {file}")
        # TODO, make sure file doesn't load file outside of storage
        prev_file = os.path.join(
            self.config.storage, file + ".json_appendable")
        if not os.path.exists(prev_file):
            logger.error(f"Cannot preload non-existent file: {prev_file}")
            self.write_message(json.dumps(
                {"error": f"Non-existent file: {file}"}))
            return False

        self.prev_file = prev_file

        metadata = self.getFileMetadata(self.prev_file)
        self.prev_file_duration = self.getLastTimestampInFile(self.prev_file)
        logger.info(
            "Previous file set. {self.prev_file} {metadata=} time: {self.prev_file_duration}")

        self.write_message(json.dumps(
            {"preloaded_svg": f"/drawing/{file}", "dimensions": [metadata[1], metadata[2]], "time": self.prev_file_duration}))

    def getFileMetadata(self, filename):
        with open(filename, "r") as fp:
            first_line = fp.readline().strip()
            if first_line.endswith(","):
                first_line = first_line[:-1]

            metadata = json.loads(first_line)

        return metadata

    def getLastTimestampInFile(self, filename):
        with open(filename, "r") as fp:
            for line in fp:
                pass  # loop until the end
            last_line = line.strip()
            if last_line.endswith(","):
                last_line = last_line[:-1]

        data = json.loads(last_line)
        if type(data) is list:
            raise Exception("Oddly, the file ends with merely metadata")

        if data['event'] == 'stroke':
            return data['points'][-1][3]
        elif data['event'] == 'viewbox':
            return data['viewboxes'][-1]['t']
        else:
            raise Exception("Unknown last event")

    # the client sent the message
    def on_message(self, message):
        logger.info(f"recieve: {message}")

        try:
            msg = json.loads(message)
            if msg["event"] == "stroke":
                logger.info("stroke")
                for i in range(len(msg['points'])):
                    msg['points'][i][3] += self.prev_file_duration
                self.appendEvent(msg)
            elif msg["event"] == "dimensions":
                self.dimensions = [int(msg["width"]), int(msg["height"])]
                logger.info(f"{self.dimensions=}")
            elif msg["event"] == "viewbox":
                logger.info("move or resize")
                if len(msg['viewboxes']) == 0:
                    logger.warn("Empty viewbox array")
                else:
                    for i in range(len(msg['viewboxes'])):
                        msg['viewboxes'][i]['t'] += self.prev_file_duration
                    self.appendEvent(msg)
            elif msg["event"] == "preload":
                self.preloadFile(msg["file"])
            else:
                # self.send({'alert': 'Unknown request: {}'.format(message)})
                logger.warn("Unknown request: {}".format(message))

        except Exception as e:
            # self.send({'alert': 'Invalid request: {}'.format(e)})
            logger.exception(e)

    # client disconnected
    def on_close(self):
        self.__class__.rmConnection(self)

        logger.info(f"Client disconnected: {self.request.remote_ip}")

    @classmethod
    def rmConnection(cls, client):
        if client not in cls.connections:
            return
        cls.connections.remove(client)

    @classmethod
    def hasConnection(cls, client):
        return client in cls.connections


class AudioListingHandler(tornado.web.RequestHandler):
    def initialize(self, config):
        self.config = config
        self.audiodir = os.path.join(self.config.storage, "audio")

    def get(self):
        # filename = self.get_argument("file", None)
        self.set_header("Content-Type", "application/json")
        if not os.path.exists(self.audiodir):
            names = []
        else:
            names = sorted(
                [
                    f"/audio/{name}"
                    for name in os.listdir(self.audiodir)
                    if name not in [".gitignore"]
                ]
            )
            print(names)
        self.write(json.dumps(names))


class AnimationHandler(tornado.web.RequestHandler):
    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index
    
    async def get(self, filename):
        # filename = self.get_argument("file", None)
        if filename == "":
            self.set_header("Content-Type", "application/json")
            files = []
            names = [
                name
                for name in os.listdir(self.config.storage)
                if name.endswith("json_appendable")
            ]
            for name in names:
                fn = os.path.join(self.config.storage, name)
                stat = os.stat(fn)
                if stat.st_size == 0:
                    continue

                with open(fn, "r") as fp:
                    first_line = fp.readline().strip()
                    if first_line.endswith(","):
                        first_line = first_line[:-1]

                    drawing_specs = json.loads(first_line)
                    drawing_id = name[:-16]
                    md =  self.index.drawings[drawing_id].get_metadata() if drawing_id in self.index.drawings else {}
                    title = md['title'] if 'title' in md else None
                    files.append(
                        {
                            "name": f"/files/{drawing_id}",
                            "id": drawing_id,
                            "title": title,
                            "ctime": drawing_specs[0],
                            "mtime": datetime.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %T"),
                            "dimensions": [drawing_specs[1], drawing_specs[2]],
                            "svg": f"/drawing/{drawing_id}.svg",
                        }
                    )

            files.sort(key=lambda k: k["mtime"])
            self.write(json.dumps(files))
        else:
            if filename[-4:] == ".svg":
                extension = "svg"
                filename = filename[:-4]
            elif filename[-4:] == ".png":
                extension = "png"
                filename = filename[:-4]
            elif filename[-4:] == ".mp3":
                extension = "mp3"
                filename = filename[:-4]
            elif filename[-4:] == ".wav":
                extension = "wav"
                filename = filename[:-4]
            else:
                extension = None

            logger.info(f"file {filename=}, {extension=}")
            # if annotation_id not in self.index.annotations:
            #     raise tornado.web.HTTPError(404)

            # annotation = self.index.annotations[annotation_id]


            t_in = self.get_argument('t_in', None)
            t_out = self.get_argument('t_out', None)

            animation = self.index.drawings[filename].get_animation()

            if t_in is not None and t_out is not None:
                animation = animation.getSlice(float(t_in), float(t_out))


            if extension == "svg":
                self.set_header("Content-Type", "image/svg+xml")
                self.write(animation.get_as_svg())
            elif extension == "png":
                self.set_header("Content-Type", "image/png")
                svgstring = animation.get_as_svg()
                self.write(cairosvg.svg2png(bytestring=svgstring))
            elif extension == "mp3":
                self.set_header("Content-Type", "audio/mp3")
                audio = await animation.audio.export(format="mp3")
                self.write(audio.read())
            elif extension == "wav":
                self.set_header("Content-Type", "audio/wav")
                audio = await animation.audio.export(format="wav")
                self.write(audio.read())
            else:
                self.set_header("Content-Type", "application/json")
                self.write(json.dumps(animation.asDict()))
            

class TagHandler(tornado.web.RequestHandler):
    """List all tags"""

    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index
        self.metadir = os.path.join(self.config.storage, "metadata")

    def get(self):
        self.set_header("Content-Type", "application/json")
        tags = self.index.tags.keys()
        self.write(json.dumps(list(tags)))


class TagAnnotationsHandler(tornado.web.RequestHandler):
    """List all annotations for given tag"""

    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index
        self.metadir = os.path.join(self.config.storage, "metadata")

    def get(self, tag):
        if not self.index.has_tag(tag):
            raise tornado.web.HTTPError(404)

        self.set_header("Content-Type", "application/json")
        # annotations = self.index.tags[tag]
        # self.write(json.dumps(list([a.id for a in annotations])))
        annotations = self.index.get_nested_annotations_for_tag(tag)
        self.write(json.dumps([{
            "id": annotation.id,
            "tag": annotation.tag,
            "id_hash": svganim.uimethods.annotation_hash(input=annotation.id),
            "url": annotation.getJsonUrl(),
            "comment": annotation.comment,
            "drawing": annotation.drawing.get_url()
        } for annotation in annotations]))


class AnnotationHandler(tornado.web.RequestHandler):
    """Get annotation as svg"""

    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index
        self.metadir = os.path.join(self.config.storage, "metadata")

    def get(self, annotation_id):
        if annotation_id[-4:] == ".svg":
            extension = "svg"
            annotation_id = annotation_id[:-4]
        elif annotation_id[-4:] == ".png":
            extension = "png"
            annotation_id = annotation_id[:-4]
        elif annotation_id[-4:] == ".mp3":
            extension = "mp3"
            annotation_id = annotation_id[:-4]
        elif annotation_id[-4:] == ".wav":
            extension = "wav"
            annotation_id = annotation_id[:-4]
        else:
            extension = None

        logger.info(f"annotation {annotation_id=}, {extension=}")
        if annotation_id not in self.index.annotations:
            raise tornado.web.HTTPError(404)

        annotation = self.index.annotations[annotation_id]

        if extension == "svg":
            self.set_header("Content-Type", "image/svg+xml")
            self.set_header("Cache-Control", "max-age=31536000, immutable")
            
            self.write(annotation.get_as_svg())
        elif extension == "png":
            self.set_header("Content-Type", "image/png")
            svgstring = annotation.get_as_svg()
            self.write(cairosvg.svg2png(bytestring=svgstring))
        elif extension == "mp3":
            self.set_header("Content-Type", "audio/mp3")
            self.write(annotation.getAnimationSlice(
            ).audio.export(format="mp3").read())
        elif extension == "wav":
            self.set_header("Content-Type", "audio/wav")
            self.write(annotation.getAnimationSlice(
            ).audio.export(format="wav").read())
        else:
            self.set_header("Content-Type", "application/json")
            self.write(json.dumps({
                "id": annotation.id,
                "tag": annotation.tag,
                "audio": f"/annotation/{annotation.id}.mp3",
            }))
    
    def post(self, annotation_id):
        """change tag for given annotation"""
        if annotation_id not in self.index.annotations:
            raise tornado.web.HTTPError(404)
        
        # might be set on file level, but let's try to avoid issues by keeping it simple
        lock = filelock.FileLock("metadata_write.lock", timeout=10)
        with lock:
            newTagId = self.get_argument('tag_id')
            if not self.index.has_tag(newTagId):
                raise tornado.web.HTTPError(400)
            
            annotation: svganim.strokes.Annotation = self.index.annotations[annotation_id]

            logger.info(f"change tag from {annotation.tag} to {newTagId}")
            
            # change metadata and reload index
            metadata = annotation.drawing.get_metadata()
            change = False
            for idx, ann in enumerate(metadata['annotations']):
                if ann['t_in'] == annotation.t_in and ann['t_out'] == annotation.t_out and annotation.tag == ann['tag']:
                    #found!?
                    metadata['annotations'][idx]['tag'] = newTagId
                    change = True
                    break
            
            if change == False:
                raise HTTPError(409)
            
            with open(annotation.drawing.metadata_fn, "w") as fp:
                logger.info(f"save tag in {annotation.drawing.metadata_fn}")
                json.dump(metadata, fp)

            self.index.refresh()


class DrawingHandler(tornado.web.RequestHandler):
    """Get drawing as svg"""

    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index
        self.metadir = os.path.join(self.config.storage, "metadata")

    def get(self, drawing_id):
        if drawing_id[-4:] == ".svg":
            extension = "svg"
            drawing_id = drawing_id[:-4]
        elif drawing_id[-4:] == ".png":
            extension = "png"
            drawing_id = drawing_id[:-4]
        elif drawing_id[-4:] == ".mp3":
            extension = "mp3"
            drawing_id = drawing_id[:-4]
        elif drawing_id[-4:] == ".wav":
            extension = "wav"
            drawing_id = drawing_id[:-4]
        else:
            extension = None

        logger.info(f"drawing {drawing_id=}, {extension=}")
        if drawing_id not in self.index.drawings:
            self.index.refresh()
            # double check
            if drawing_id not in self.index.drawings:
                raise tornado.web.HTTPError(404)

        drawing = self.index.drawings[drawing_id]

        if extension == "svg":
            self.set_header("Content-Type", "image/svg+xml")
            self.write(drawing.get_animation().get_as_svg())
        elif extension == "png":
            self.set_header("Content-Type", "image/png")
            svgstring =drawing.get_animation().get_as_svg()
            self.write(cairosvg.svg2png(bytestring = svgstring))
        elif extension == "mp3":
            self.set_header("Content-Type", "audio/mp3")
            self.write(drawing.get_animation(
            ).audio.export(format="mp3").read())
        elif extension == "wav":
            self.set_header("Content-Type", "audio/wav")
            self.write(drawing.get_animation(
            ).audio.export(format="wav").read())
        else:
            self.set_header("Content-Type", "application/json")
            self.write(json.dumps({
                "id": drawing.id,
                "annotations_url": drawing.get_annotations_url(),
                "audio": f"/drawing/{drawing.id}.mp3",
                "svg": f"/drawing/{drawing.id}.svg",
            }))


class AnnotationsHandler(tornado.web.RequestHandler):
    def initialize(self, config):
        self.config = config
        self.metadir = os.path.join(self.config.storage, "metadata")

    def prepare(self):
        if self.request.headers.get("Content-Type", "").startswith("application/json"):
            self.json_args = json.loads(self.request.body)
        else:
            self.json_args = None

    def get_filenames(self):
        return [
            name[:-16]
            for name in os.listdir(self.config.storage)
            if name.endswith("json_appendable")
        ]

    def get(self, filename):
        self.set_header("Content-Type", "application/json")
        filenames = self.get_filenames()

        print(filenames, filename)

        if filename not in filenames:
            raise tornado.web.HTTPError(404)

        meta_file = os.path.join(self.metadir, filename + ".json")
        if not os.path.exists(meta_file):
            self.set_status(404)
            return

        with open(meta_file, "r") as fp:
            self.write(json.load(fp))

    def post(self, filename):
        # filename = self.argument("file", None)

        filenames = self.get_filenames()
        print(filenames, filename)

        if filename not in filenames:
            raise tornado.web.HTTPError(404)

        if not os.path.exists(self.metadir):
            os.mkdir(self.metadir)

        meta_file = os.path.join(self.metadir, filename + ".json")
        with open(meta_file, "w") as fp:
            json.dump(self.json_args, fp)

class TagsHandler(tornado.web.RequestHandler):
    def initialize(self, config, index: svganim.strokes.AnnotationIndex) -> None:
        self.config = config
        self.index = index
    
    def get(self):
        self.set_header("Content-Type", "application/json")
        self.write(self.index.root_tag.toJson(with_counts=True))
        # with open('www/tags.json', 'r') as fp:
        #     # TODO: enrich with counts
        #     self.write(fp.read())

    def put(self):
        # data = json.loads(self.request.body)
        tree = svganim.strokes.loadTagFromJson(self.request.body)
        logger.info(f"New tag tree:\n{tree}")
        newTagsContent = tree.toJson()
        # save at minute resolution
        now = datetime.datetime.utcnow().isoformat(timespec='minutes')

        backup_dir = os.path.join(self.config.storage, 'tag_versions')
        if not os.path.exists(backup_dir):
            logger.warning(f"Creating tags backupdir {backup_dir}")
            os.mkdir(backup_dir)
        
        bakfile = os.path.join(backup_dir, f'tags.{now}.json')
        logger.info(f"Creating tags backup {bakfile}" )
        shutil.copyfile('www/tags.json', bakfile)

        with open('www/tags.json', 'w') as fp:
            fp.write(newTagsContent)
        
        # update as to load new tag into cache
        self.index.refresh()

        self.set_status(204)
        # print()
        

class IndexHandler(tornado.web.RequestHandler):
    """Get annotation as svg"""

    def initialize(self, config, index: svganim.strokes.AnnotationIndex):
        self.config = config
        self.index = index

    def get(self):
        do_refresh = bool(self.get_query_argument('refresh', False))
        if do_refresh:
            logger.info("Reloading Annotation Index")
            self.index.refresh()
            logger.info("\treloaded annotation index")

        self.render("templates/index.html", index=self.index)


class Server:
    """
    Server for HIT -> plotter events
    As well as for the Status interface
    """

    loop = None

    def __init__(self, config, logger):
        self.config = config
        self.logger = logger

        # self.config['server']['port']
        self.web_root = os.path.join("www")

        if not os.path.exists(self.config.storage):
            raise NotADirectoryError("Provided files directory doesn't exist.")

        self.index = svganim.strokes.AnnotationIndex(
            os.path.join(self.config.storage,"annotation_index.shelve"), self.config.storage, os.path.join(self.config.storage,"metadata")
        )
        self.logger.info("Loading Annotation Index")
        self.index.refresh()
        self.logger.info("\tloaded annotation index")

    def start(self):
        application = tornado.web.Application(
            [
                (
                    r"/ws(.*)",
                    WebSocketHandler,
                    {
                        "config": self.config,
                    },
                ),
                (r"/files/(.*)", AnimationHandler, {"config": self.config, "index": self.index}),
                (
                    r"/audio/(.+)",
                    tornado.web.StaticFileHandler,
                    {"path": os.path.join(self.config.storage, "audio")},
                ),
                (r"/audio", AudioListingHandler, {"config": self.config}),
                (r"/annotations/(.+)", AnnotationsHandler,
                 {"config": self.config}),
                (r"/tags", TagHandler,
                 {"config": self.config, "index": self.index}),
                (
                    r"/tags/(.+)",
                    TagAnnotationsHandler,
                    {"config": self.config, "index": self.index},
                ),
                (
                    r"/annotation/(.+)",
                    AnnotationHandler,
                    {"config": self.config, "index": self.index},
                ),
                (
                    r"/drawing/(.+)",
                    DrawingHandler,
                    {"config": self.config, "index": self.index},
                ),

                (r"/index", IndexHandler,
                 {"config": self.config, "index": self.index}),

                (r"/tags.json", TagsHandler,
                 {"config": self.config, "index": self.index}),

                (r"/(.*)", StaticFileWithHeaderHandler,
                 {"path": self.web_root, 'default_filename': 'index.html'}),
            ],
            debug=True,
            autoreload=True,
            ui_methods= svganim.uimethods
        )
        application.listen(self.config.port)
        tornado.ioloop.IOLoop.current().start()


if __name__ == "__main__":
    argParser = argparse.ArgumentParser(
        description="Start up the vector animation server"
    )
    # argParser.add_argument(
    #     '--config',
    #     '-c',
    #     required=True,
    #     type=str,
    #     help='The yaml config file to load'
    # )
    argParser.add_argument("--port", type=int, default=7890, help="Port")
    argParser.add_argument(
        "--storage", type=str, default="files", help="directory name for output files"
    )
    argParser.add_argument(
        "--logfile", type=str, default=None, help="log file to output to"
    )
    argParser.add_argument("--verbose", "-v", action="count", default=0)

    args = argParser.parse_args()

    loglevel = (
        logging.NOTSET
        if args.verbose > 1
        else logging.DEBUG
        if args.verbose > 0
        else logging.INFO
    )

    coloredlogs.install(
        level=loglevel,
        #       default: "%(asctime)s %(hostname)s %(name)s[%(process)d] %(levelname)s %(message)s"
        fmt="%(asctime)s %(hostname)s %(name)s[%(process)d,%(threadName)s] %(levelname)s %(message)s",
    )

    # File logging
    if args.logfile is not None:
        formatter = logging.Formatter(
            fmt="%(asctime)s %(module)s:%(lineno)d %(levelname)8s | %(message)s",
            datefmt="%Y/%m/%d %H:%M:%S",
        )  # %I:%M:%S %p AM|PM format
        logFileHandler = logging.handlers.RotatingFileHandler(
            args.logfile, maxBytes=1024 * 512, backupCount=5
        )
        logFileHandler.setFormatter(formatter)

        logger.addHandler(logFileHandler)
    

    logger.info(f"Start server: http://localhost:{args.port}")

    server = Server(args, logger)
    server.start()