import argparse
import json
import sys
import time
import multiprocessing
from typing import List, Dict, Union

import mapnik
from shapely.geometry import shape

from utils import load_style, Utils


class RenderWorkerProcess(multiprocessing.Process):
    def __init__(self, args: argparse.Namespace, worker_num: int):
        super().__init__()
        self._args = args
        self._worker_number = worker_num
        self._print("Initialization done.", color="green")

    @property
    def args(self):
        return self._args

    @property
    def worker_number(self):
        return self._worker_number

    def _print(self, msg: str, color: str = ""):
        prefix = "[WORKER {}] ".format(self.worker_number)
        print("\033[9{}m{}\033[0m".format(color, prefix + msg))

    def _render_tile(self, z: int, x: int, y: int) -> bytes:
        m = mapnik.Map(self.args.tile_size, self.args.tile_size)
        m.load_font_set("../fonts.list")
        m.loads(self.args.style_filename)

        ds = mapnik.Datasource(
            plugin="shape",
            parameter={"file": self.args.input_data_filename},
            bbox=(0, 0, 1, 1),
        )

        layer = m.layers[0]
        layer.datasource = ds

        proj_str = "+proj=merc +lon_0=0 +k=1 +x_0=0 +y_0=0 +ellps=WGS84 +datum=WGS84 +units=m +no_defs"
        mercer = pyproj.Transformer.from_crs("epsg:4326", proj_str)

        left, bottom, right, top = self.args.bbox[z][x][y]
        ul_longitude, ul_latitude = mercer.transform(left, top)
        ll_longitude, ll_latitude = mercer.transform(left, bottom)
        ur_longitude, ur_latitude = mercer.transform(right, top)

        bounding_box = mapnik.Box2d(ul_longitude, ul_latitude, ur_longitude, ur_latitude)
        m.zoom_all_to(bounding_box)

        im = mapnik.Image(self.args.tile_size, self.args.tile_size)
        mapnik.render(m, im, mapnik.Envelope(*bounding_box.bounds()))

        return im.tostring()

    def render(self, queue):
        cnt = 0

        while True:
            params = queue.get()
            if params == "EXIT":
                queue.task_done()
                break

            z, x, y = params
            result = self._render_tile(z, x, y)
            with open(f"rendered/{self.args.base_name}-{z}_{x}_{y}", "wb") as f:
                f.write(result)

            self._print(f"Successfully processed {cnt}: ({z},{x},{y}).", color="blue")
            cnt += 1
            queue.task_done()


def parse_args(args):
    parser = argparse.ArgumentParser(description="Generate Tile Images.")
    mandatory_opts = parser.add_argument_group("Required arguments")

    mandatory_opts.add_argument("-i", "--input-data-filename", required=True, metavar="FILE",
                                dest="input_data_filename",
                                help="Path to OSM XML/PBF data file.")

    mandatory_opts.add_argument("-s", "--style-filename", required=True, metavar="FILE",
                                dest="style_filename",
                                help="Path to the customized Mapnik style XML file.")

    parser.add_argument("-w", "--workers", type=int, default=2, metavar="COUNT",
                        dest="process_count",
                        help="Number of parallel workers to spawn for rendering tiles."
                             "(Default: 2)")

    parser.add_argument("--base-name", default="rendering", metavar="STR",
                        dest="base_name",
                        help="Base name for resulting tile images. Default: \"rendering\".")

    parser.add_argument("--tile-size", type=int, default=256, metavar="PIXELS",
                        dest="tile_size",
                        help="Width & Height dimensions for square tiles. Default: 256 pixels.")

    parser.add_argument("--cores", action="store_true",
                        dest="use_cpu_count",
                        help="Automatically detect CPU core count for workers."
                             "Overridden by --workers argument. Can also force enable by passing -w N -c flag.")

    return parser.parse_args(args)


if __name__ == "__main__":
    START = time.time()

    ARGS = parse_args(sys.argv[1:])
    WORKERS_COUNT = min(ARGS.process_count, multiprocessing.cpu_count())
    USE_CPU_COUNT = False if ARGS.process_count != -1 else True

    u = Utils()
    MAX_ZOOM, MIN_ZOOM, BB = u.calc_params(ARGS.input_data_filename)
    QUEUE = multiprocessing.JoinableQueue(maxsize=WORKERS_COUNT)

    PROCS = []
    for i in range(WORKERS_COUNT):
        proc = RenderWorkerProcess(ARGS, i)
        PROCS.append(proc)
        proc.start()

    for z in range(MIN_ZOOM, MAX_ZOOM + 1):
        # Modifying inner loop behavior below leads to significant speedups in certain cases
        # Experiment with this block to fit personal preferences
        for x in range(0, 2**z, 2):
            for y in range(0, 2**z, 2):
                QUEUE.put((z, x, y))

    QUEUE.join()

    print("\033[93mClosing active workers...\033[0m")
    for _ in range(WORKERS_COUNT):
        QUEUE.put("EXIT")

    for proc in PROCS:
        proc.join()

    END = time.time()
    print(f"\033[92mCompleted generation in {END - START:.2f} seconds!\033[0m")