import argparse
import json
import logging
import multiprocessing as mp
import os
import os.path as osp
import shutil
import signal
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Callable, Dict, List, Optional
import threading

import cv2
import numpy as np
import tqdm
import uvicorn
from fastapi import FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse, RedirectResponse
import pyrealsense2 as rs

from markit_rs_recorder.common import (
    CALLBACKS,
    new_realsense_camera_system_from_yaml_file,
    RealsenseSystemModel,
    RealsenseSystemCfg,
    RealsenseCameraCfg,
    get_datetime_tag,
    enumerate_devices_that_supports_advanced_mode,
)

from .post_process import worker as post_processing_worker

app = FastAPI()
logging.basicConfig(level=logging.INFO)

STOP_EV: mp.Event = mp.Event()
FINISH_EV: mp.Event = mp.Event()
READY_EV: mp.Event = mp.Event()
CAPTURE_PROCS: List[mp.Process] = []
POST_PROCESSING_QUEUE: Optional[mp.Queue] = None
ARGS: Optional[argparse.Namespace] = None


class RemoteRecordSeq(RealsenseSystemModel):
    def __init__(self,
                 system_cfg: RealsenseSystemCfg,
                 camera_cfgs: List[RealsenseCameraCfg],
                 callbacks: Dict[str, Callable] = None):
        super().__init__(system_cfg, camera_cfgs, callbacks)
        self.metadata: Dict[Dict] = {cam.friendly_name: [] for cam in self.cameras}

    def insert_meta_data(self, idx, ts, sys_ts, frame_counter, frame_basename):
        self.metadata[idx].append({
            "dev_ts": ts,
            "sys_ts": sys_ts,
            "frame_counter": frame_counter,
            "frame_basename": frame_basename
        })

    def save_meta_data(self):
        save_path = osp.join(self.options.base_dir, "metadata_all.json")
        config_save_path = osp.join(self.options.base_dir, "realsense_config.json")
        bundle = {
            "camera_sn": [cam.option.sn for cam in self.cameras],
            "metadata": self.metadata
        }
        with open(save_path, 'w') as f:
            json.dump(bundle, f, indent=4)

        with open(config_save_path, 'w') as f:
            json.dump({"realsense": {"system": self.options.get_dict(), "cameras": list(map(lambda x: x.get_dict(), self.camera_options))}}, f, indent=4)

    def app(self, stop_ev: mp.Event, finish_ev: mp.Event, ready_ev: mp.Event = None):

        def save_color_frame(path, frame):
            cv2.imwrite(path, frame)

        def save_depth_frame(path, frame):
            np.save(path, frame)

        self.console.log(f"tag is {self.tag}")
        self.console.log(f"save path is {self.options.base_dir}")
        self.console.log(f"number of cameras is {len(self.cameras)}")
        self._set_advanced_mode()
        self.console.log("opening devices")
        self.open()
        self.console.log("starting devices")
        self.start(interval_ms=self.options.interval_ms)

        self.console.log("saving intrinsics")
        for cam in self.cameras:
            '''Convert intrinsic data to dict (readable in Open3D)'''
            mat = [cam.intrinsics.fx, 0, 0, 0, cam.intrinsics.fy, 0, cam.intrinsics.ppx, cam.intrinsics.ppy, 1]
            intrinsics_dict = {'width': cam.intrinsics.width, 'height': cam.intrinsics.height, 'intrinsic_matrix': mat}
            save_path = osp.join(cam.save_path, 'realsense_intrinsic.json')
            if not osp.lexists(save_path):
                with open(save_path, 'w') as f:
                    json.dump(intrinsics_dict, f, sort_keys=False,
                              indent=4,
                              ensure_ascii=False)

        self.console.log("[]start recording")
        progress_bars = [tqdm.tqdm(ncols=0) for _ in range(len(self.cameras))]
        save_workers = ThreadPoolExecutor(max_workers=len(self.cameras) * 2)

        if ready_ev is not None:
            ready_ev.set()

        try:
            while True:
                for idx, cam in enumerate(self.cameras):
                    tic = time.time()
                    color_image, depth_image, ts, sys_ts, frame_counter = cam.get_frames()
                    frame_basename = '{:07d}_{}_{}'.format(frame_counter, ts["backend_t"], sys_ts)
                    self.insert_meta_data(cam.friendly_name, ts, sys_ts, frame_counter, frame_basename)

                    toc = time.time()

                    if color_image is not None:
                        save_workers.submit(save_color_frame, osp.join(cam.color_save_path, frame_basename + '.bmp'), color_image)
                        # i5 12400K save jpg at 17 fps (load 60%), write to PM9A1 at 200 - 500MB/s, memory consumption 1GB/min percamera
                        # i5 12400K save bmp at 20 fps, but write at 1.0GB/s, memory consumption 0GB/min per camera

                        # save_workers.submit(cv2.imwrite, (osp.join(cam.color_save_path, f'{cam.n_frames}_{ts}_{sys_ts}.bmp'), color_image,))
                        # cv2.imwrite(osp.join(cam.color_save_path, f'{cam.n_frames}_{ts}_{sys_ts}.bmp'), color_image)

                    if depth_image is not None:
                        save_workers.submit(save_depth_frame, osp.join(cam.depth_save_path, frame_basename + '.npy'), depth_image)
                        # save_workers.submit(lambda: cv2.imwrite(osp.join(cam.depth_save_path, f'{cam.n_frames}_{ts}_{sys_ts}.png'), depth_image, [cv2.IMWRITE_PNG_COMPRESSION, 0]))

                        # np.save(osp.join(cam.depth_save_path, f'{cam.n_frames}_{ts}_{sys_ts}.npy'), depth_image)
                        # cv2.imwrite(osp.join(cam.depth_save_path, f'{cam.n_frames}_{ts}_{sys_ts}.png'), depth_image, [cv2.IMWRITE_PNG_COMPRESSION, 0])

                    progress_bars[idx].set_description(f'SN={cam.option.sn}, ProcessTime={int((toc - tic) * 1000)}(ms), FrameCounter={frame_counter}')
                    progress_bars[idx].update(1)

                    if self.options.interactive and not self.options.use_bag:
                        if color_image is not None and depth_image is not None:
                            mix = cv2.addWeighted(color_image[..., 0:3], 0.5, cv2.resize(depth_image, (color_image.shape[1], color_image.shape[0])), 0.5, 0)
                        elif color_image is not None and depth_image is None:
                            mix = color_image
                        elif color_image is None and depth_image is not None:
                            mix = depth_image
                        else:
                            break
                        cv2.imshow(cam.window_name, mix)
                        cv2.waitKey(1)
                if stop_ev.is_set():
                    self.console.log('stop recording')
                    finish_ev.set()
                    raise KeyboardInterrupt
        except KeyboardInterrupt as e:
            # raise(e)
            self.console.log("\n" * len(self.cameras))
            self.console.log("stopped in response to KeyboardInterrupt")

            self.console.log("stopping cameras")
            self.stop(interval_ms=self.options.interval_ms)
            for idx, cam in enumerate(self.cameras):
                self.console.log(f"saving last frames of camera {cam.option.sn}")
                while True:
                    color_image, depth_image, ts, sys_ts, frame_counter = cam.get_frames(timeout_ms=50)
                    if frame_counter > 0:
                        self.insert_meta_data(cam.friendly_name, ts, sys_ts, frame_counter, frame_basename)
                        frame_basename = '{:07d}_{}_{}'.format(frame_counter, ts["backend_t"], sys_ts)
                        if color_image is not None:
                            save_workers.submit(save_color_frame, osp.join(cam.color_save_path, frame_basename + '.bmp'), color_image)

                        if depth_image is not None:
                            save_workers.submit(save_depth_frame, osp.join(cam.depth_save_path, frame_basename + '.npy'), depth_image)

                        progress_bars[idx].set_description(f'SN={cam.option.sn}, FrameCounter={frame_counter}')
                        progress_bars[idx].update(1)
                    else:
                        break
            self.close()
            save_workers.shutdown(wait=True)
            self.save_meta_data()
            cv2.destroyAllWindows()
            return

        except Exception as e:
            self.console.log(e)
            raise e
        finally:
            finish_ev.set()
            ready_ev.clear()
            self.console.log(f"finished recording, tag={self.tag}")


def make_response(status_code, **kwargs):
    data = {'code': status_code, 'timestamp': time.time()}
    data.update(**kwargs)
    json_compatible_data = jsonable_encoder(data)
    resp = JSONResponse(content=json_compatible_data, status_code=status_code)
    return resp


def capture_frames(stop_ev: mp.Event,
                   finish_ev: mp.Event,
                   ready_ev: mp.Event,
                   config: str,
                   tag: str,
                   post_processing_queue: mp.Queue = None):
    callbacks = {
        CALLBACKS.tag_cb: lambda: tag,
        CALLBACKS.save_path_cb: lambda cam_cfg, sys_cfg: osp.join(sys_cfg.base_dir, "r" + cam_cfg.sn[-2:]),
        CALLBACKS.camera_friendly_name_cb: lambda cam_cfg, _: "r" + cam_cfg.sn[-2:]
    }

    sys = new_realsense_camera_system_from_yaml_file(RemoteRecordSeq, config, callbacks)

    try:
        sys.app(stop_ev, finish_ev, ready_ev)
    except AttributeError as e:
        logging.error(f"capture_frames failed: {e}")

    if post_processing_queue is not None:
        logging.info(f"adding {sys.options.base_dir} to post processing queue")
        post_processing_queue.put(sys.options.base_dir)


@app.get("/")
def root():
    return RedirectResponse(url='/docs')


@app.get("/v1/status")
def status():
    return make_response(status_code=200, active_processes=[proc.is_alive() for proc in CAPTURE_PROCS].count(True))


@app.get("/v1/ready")
def ready():
    global READY_EV
    if READY_EV is not None:
        if READY_EV.is_set():
            return make_response(status_code=200, ready=True)
        else:
            return make_response(status_code=200, ready=False)
    else:
        return make_response(status_code=500, msg="NOT SUPPORTED")


@app.post("/v1/start")
def start_process(tag: str = None):
    global CAPTURE_PROCS, STOP_EV, FINISH_EV, READY_EV, ARGS, POST_PROCESSING_QUEUE

    # Wait until last capture ends
    if len(CAPTURE_PROCS) > 0:
        if STOP_EV.is_set():
            FINISH_EV.wait(timeout=5)
            if FINISH_EV.is_set():
                [proc.join(timeout=3) for proc in CAPTURE_PROCS]
                if any([proc.is_alive() for proc in CAPTURE_PROCS]):
                    logging.warning("[realsense] join timeout")
                    [os.kill(proc.pid, signal.SIGTERM) for proc in CAPTURE_PROCS if proc.is_alive()]
                # clean up resources
                CAPTURE_PROCS = []
                READY_EV.clear()
            else:
                return make_response(status_code=500, msg="NOT FINISHED")
        else:
            return make_response(status_code=500, msg="RUNNING")

    if len(CAPTURE_PROCS) == 0:
        STOP_EV.clear()
        FINISH_EV.clear()

        if tag is None:
            tag = get_datetime_tag()
        logging.info(f"[realsense] tag={tag}")

        if len(CAPTURE_PROCS) <= 0:
            CAPTURE_PROCS = [mp.Process(target=capture_frames,
                                        args=(STOP_EV,
                                              FINISH_EV,
                                              READY_EV,
                                              ARGS.config,
                                              tag,
                                              POST_PROCESSING_QUEUE))]
            DELAY_S = 2  # Magic delay duration to avoid U3V communication error
            [(proc.start(), time.sleep(DELAY_S
                                       )) for proc in CAPTURE_PROCS]

        return make_response(status_code=200, msg="START OK", subpath=tag)


@app.post("/v1/stop")
def stop_process():
    global CAPTURE_PROCS, STOP_EV
    logging.info("[realsense] Stop")

    if len(CAPTURE_PROCS) > 0 and any([proc.is_alive() for proc in CAPTURE_PROCS]):
        STOP_EV.set()
        return make_response(status_code=200, msg=f"STOP OK: {len([None for proc in CAPTURE_PROCS if proc.is_alive()])} procs are running")
    else:
        return make_response(status_code=500, msg="NOT RUNNING")


@app.post("/v1/reset")
def reset_camera():
    global CAPTURE_PROCS, STOP_EV
    if len(CAPTURE_PROCS) > 0 and any([proc.is_alive() for proc in CAPTURE_PROCS]):
        return make_response(status_code=500, msg="RUNNING")
    else:
        for dev in enumerate_devices_that_supports_advanced_mode(rs.context()):
            dev.hardware_reset()
        time.sleep(3)
        return make_response(status_code=200, msg="RESET OK")


@app.post("/v1/kill")
def kill_process():
    global CAPTURE_PROCS, STOP_EV, FINISH_EV
    logging.info("[realsense] kill")

    if len(CAPTURE_PROCS) and any([proc.is_alive() for proc in CAPTURE_PROCS]) > 0:
        STOP_EV.set()
        FINISH_EV.wait(timeout=1)
        [proc.join(timeout=1) for proc in CAPTURE_PROCS]
        if any([proc.is_alive() for proc in CAPTURE_PROCS]):
            logging.warning("[realsense] join timeout, force kill all processes")
            [os.kill(proc.pid, signal.SIGTERM) for proc in CAPTURE_PROCS if proc.is_alive()]
        return make_response(status_code=200, msg="KILL OK")
    else:
        return make_response(status_code=500, msg="NOT RUNNING")


def post_processing_thread():
    global POST_PROCESSING_QUEUE, ARGS
    logging.info("post processing thread started")
    while True:
        if not POST_PROCESSING_QUEUE.empty():
            base_dir = POST_PROCESSING_QUEUE.get()
            if ARGS.calibration is not None:
                try:
                    _calibration_file_name = "calibration.json"
                    shutil.copy(osp.join(ARGS.calibration, _calibration_file_name), base_dir)
                    logging.info(f"copy calibration file from {osp.join(ARGS.calibration, _calibration_file_name)} to {base_dir}")
                except Exception as e:
                    logging.error(f"copy calibration file failed: {e}")
                try:
                    post_processing_worker(base_dir)
                except Exception as e:
                    logging.error(f"post processing failed: {e}")
                # p = mp.Process(target=post_processing_worker, args=(sys.options.base_dir,))
        else:
            time.sleep(1)


def main(args: argparse.Namespace):
    global ARGS, POST_PROCESSING_QUEUE
    POST_PROCESSING_QUEUE = mp.Queue()
    ARGS = args
    # Prepare system
    logging.info('the server listens at port {}'.format(args.port))

    if not osp.exists(args.config):
        logging.error(f"config file {args.config} does not exist")
        exit(1)

    if args.calibration is not None and not osp.exists(args.calibration):
        logging.error(f"calibration file {args.calibration} does not exist")
        exit(1)

    try:
        t = threading.Thread(target=post_processing_thread)
        t.start()
        uvicorn.run(app=app, port=args.port)
    except KeyboardInterrupt:
        logging.info(f"main() got KeyboardInterrupt")
        os._exit(1)


def entry_point(argv):
    parser = argparse.ArgumentParser(description='Recorder')
    # parser.add_argument('--app', type=str, help='', default='')
    parser.add_argument('--config', type=str, help='The realsense system configuration', default='./realsense_config.yaml')
    parser.add_argument('--calibration', type=str, help='The realsense system calibration', default=None)
    parser.add_argument('--port', type=int, help="Port to listen", default=5050)
    parser.add_argument('--debug', action='store_true', help='Toggle Debug mode')
    args = parser.parse_args(argv)
    main(args)


if __name__ == '__main__':
    import sys

    entry_point(sys.argv[1:])
