import time
import json
import logging
import threading
import numpy as np

from nxva.v11 import Detector
from nxva.streaming import MultiStreaming
from nxva.va.shm_utils import SHMHandler
from collections import deque
from flask import Flask, request, jsonify

app = Flask(__name__)

class DetectionServer:
    def __init__(self, max_detection_num, model_list, allowed_class_names, all_cameras_ids, interval=0, namespace=None):
        '''
        Args:
            max_detection_num: int
                max number of detection in a frame
            model_list: list
                list of model names
                model_list = ['./configs/detection.yaml', ...]
            allowed_class_names: list
                list of class names which the task needs
                allowed_class_names = ['people', 'car', ...]
            all_cameras_ids: list
                append all cameras.yaml where in cameras's folder
            interval: int
                interval time for detection, default is 0
            namespace: str
                namespace for logger, default is None

        Attributes:
            max_detection_num: int
                max number of detection in a frame
            all_cameras_ids: list
                append all cameras.yaml where in cameras's folder
            stream: MultiStreaming
                streaming
            detector: DetectorBuilder
                detector
            model_dict: dict
                run models by model_name
                model_dict[model_name] = model
            class_name_dict: dict
                send class names to client after label shift
                class_name_dict = {0: 'people', ...}
            label_shift: dict
                merge all labels
                label_shift[model_name] = shift
            task_list: list
                list of task names
                task_list = [task1, task2, ...]
            mix_name_classes: dict
                mix_name_classes = {class_name: class_id, ...}
            frames_results_queue
                frames_results_queue = queue.Queue()
            queue_lock: threading.Lock
                lock for frames_results_queue
        '''
        self.logger = logging.getLogger(namespace) if namespace else logging.getLogger()

        self.app = Flask(__name__)
        self.setup_routes()

        self.max_detection_num     = max_detection_num
        self.interval              = interval

        self.all_cameras_ids       = all_cameras_ids
        self.cameras_num           = len(all_cameras_ids) # camera數量

        # streaming
        self.stream = MultiStreaming(config='./configs/streaming.yaml', 
                                    reconnect_interval=30,
                                    verbose=True
                                    )
        self.stream.run()

        # model yaml
        self.mix_class_names  = {}
        self.model_dict       = {}
        self.label_shift      = {}
        shift                 = 0
        
        model_name = 0
        for cfg in model_list:
            model = Detector(cfg)
            model.classes = []
            model.max_det = max_detection_num
            model.name_classes = {name: class_id for class_id, name in model.class_names.items()}
            self.model_dict[model_name] = model
            model_name += 1

        for cls_name in allowed_class_names:
            for model_name, model in reversed(self.model_dict.items()):
                class_id = model.name_classes.get(cls_name)
                if class_id is not None:
                    model.classes.append(class_id)

        shift = 0
        for model_name, model in self.model_dict.items():
            for class_id, cls_name in model.class_names.items():
                self.mix_class_names[class_id + shift] = cls_name
            self.label_shift[model_name] = shift
            shift = len(self.mix_class_names)

        self.mix_name_classes     = {name: class_id for class_id, name in self.mix_class_names.items()}
        
        # task for debug
        self.task_list            = []

        # queue
        self.frames_results_queue = deque(maxlen=1)

        # lock
        self.queue_lock           = threading.Lock()

        # SHMHandler
        self.shm_handler          = SHMHandler(self.stream.w, self.stream.h, self.cameras_num, max_detection_num, namespace=namespace)

        # shared memory
        self.shm_handler.build_shared_memory()

        # init control block
        self.shm_handler.write_control_info(slot_idx=0)

        # stop event
        self.stop_event           = threading.Event()
        self.interval_stop_event  = threading.Event()

        threading.Thread(target=self.api_run, daemon=True).start()
    
    
    def setup_routes(self):
        @app.route('/api/registration_server', methods=['POST'])
        def registration_server():
            '''
            Registration server.
            Listen to the client and send the shared memory info.
            body:
                {
                "camera_ids"    : camera_ids,
                "class_names"   : class_names,
                }
            '''
            self.logger.info('Waiting for connection...')

            message       = request.get_json()
            cameras_ids   = message['camera_ids']
            class_names   = message['class_names']

            task_count = len(self.task_list)
            name = f"task_{task_count}"
            
            # find camera_ids in all_cameras_ids
            camera_indices_list = [
                self.all_cameras_ids.index(task_camera) 
                for task_camera in cameras_ids
                if task_camera in self.all_cameras_ids
            ]

            # find class_id in class_names
            class_name_dict = {
                self.mix_name_classes[name]: name
                for name in class_names
                if name in self.mix_name_classes
            }

            self.task_list.append(name)
            self.logger.info(f"Register task: {self.task_list}")

            respone = {
                'shm_name':                 "frames_results_shm",
                'total_cameras_num':        self.cameras_num,
                'camera_ids':               camera_indices_list,
                'width':                    self.stream.w,
                'height':                   self.stream.h,
                'class_name_dict':          class_name_dict,
                'max_detection_num':        self.max_detection_num,
            }

            return jsonify(respone), 200

        @app.route('/api/watchdog_server', methods=['POST'])
        def watchdog_server():
            '''
            Check if the cameras is running.
            body:
                {
                "device_id": "CAMERA-201-1"
                }
            '''
            data      = request.get_json()
            device_id = data['device_id']

            camera_index  = self.all_cameras_ids.index(device_id)
            camera_status = self.stream.status[camera_index]

            server_available    = []
            server_unavailable  = []
            if camera_status:
                server_available.append(device_id)
            else:
                server_unavailable.append(device_id)

            data = {
                'server_available': server_available,
                'server_unavailable': server_unavailable
                }

            return jsonify(data), 200
        
    def api_run(self):
        app.run(host='0.0.0.0', port=9453)
    
    def run(self):
        '''
        Run get_frames_results_queue threading.
        Get frames and results from stream, and put them into queue.
        '''
        self.stop_event.clear()
        num_frames = self.stream.num
        threading.Thread(target=self._write_frames_results_shm, daemon=True).start()
        next_time = time.perf_counter() + self.interval
        while not self.stop_event.is_set():
            merged_results = [np.empty((0, 6), dtype=float) for _ in range(num_frames)]
            frames, status = self.stream.get_frames(status=True)

            for model_name, model in self.model_dict.items():
                results = model(frames)

                # Shift class_id and merge results
                for frame_idx, result in enumerate(results):
                    if result.shape[1] >= 6:
                        result[:, 5] += self.label_shift[model_name]
                    merged_results[frame_idx % num_frames] = np.vstack((merged_results[frame_idx % num_frames], result))

            with self.queue_lock:
                self.frames_results_queue.append((frames, merged_results))

            # Calculate the next time to send the next frame
            next_time += self.interval
            delay = next_time - time.perf_counter()
            if delay > 0:
                self.interval_stop_event.wait(delay)
            else:
                time.sleep(0.001)
            
            time.sleep(0.01)
    
    def _write_frames_results_shm(self):
        '''
        If queue is not empty, get frames and results from queue, and put them into frames_buffer and results_buffer.

        1) Get frames and results from queue
        2) Write frames and results to shared memory
        3) Write control info to shared memory
        '''

        frame_counter = 0
        while not self.stop_event.is_set():
            if len(self.frames_results_queue) > 0:
                with self.queue_lock:
                    frames, results = self.frames_results_queue.popleft()
                
                for cam_id in range(self.cameras_num):

                    # write frames and metas to shared memory
                    old_slot_idx = self.shm_handler.read_control_info()

                    # buffer index (0->1 or 1->0)
                    next_slot_idx = 1 - old_slot_idx

                    self.shm_handler.write_frame_and_bbox(cam_id, next_slot_idx, frames[cam_id], results[cam_id])

                # frame_counter limit, results
                if frame_counter <= self.stream.h*self.stream.w:
                    frame_counter += 1
                else:
                    frame_counter = 0

                meta_dict = {
                    "camera_id": cam_id,
                    "frame_counter": frame_counter,
                    "timestamp": time.time()
                }
                meta_str = json.dumps(meta_dict)
                
                self.shm_handler.write_control_info(next_slot_idx)
                self.shm_handler.write_meta_info(meta_str)

                    
            time.sleep(0.001)

    def stop(self):
        '''
        Stop the detection server.
        '''
        self.stop_event.set()
        self.logger.info('Detection server stopped.')


if __name__ == "__main__":
    from utilities import load_configs

    task_list, hyper_config, streaming_config, cameras_dict, cameras_configs_dict, all_cameras_ids = load_configs()

    max_detection_num = hyper_config['shared_memory']['max_detection_num']
    model_list        = hyper_config['model']['model_list']
    class_names  = [config['class'] for _, config in hyper_config.items() if 'class' in config]
    allowed_class_names = list(set(name for class_name in class_names for name in class_name))

    dection_server = DetectionServer(max_detection_num, model_list, allowed_class_names, all_cameras_ids)
    dection_server.run()