import datetime
import os
import sys
import threading

sys.path.append('../')  # NOQA: E402

from protos import aerial_pb2 as pb
from flask import Flask, request
import pika


class MyServer(Flask):
    def __init__(self, *args, **kwargs):
        Flask.__init__(self, *args, **kwargs)
        self.message_received: int = 0
        self._def_routes()
        self.devices = {}
        self.route('/api', methods=['POST'])(self.api)
        self.route('/', methods=['GET'])(self.hello)

    def init_rabbmitmq(self, host: str, user: str, password: str, queue: str):
        credentials = pika.PlainCredentials(user, password)
        parameters = pika.ConnectionParameters(
            credentials=credentials, host=host)
        connection = pika.BlockingConnection(parameters)
        self.channel = connection.channel()
        self.channel.queue_declare(queue=queue, durable=True)
        self.channel.basic_qos(prefetch_count=1)
        self.channel.basic_consume(
            on_message_callback=self._on_request, queue=queue)

    def _def_routes(self):
        self.routes = {
            pb.SysCommand.SYS_COMMAND: self._on_sys_command,
            pb.SysCommand.SENSOR_COMMAND: self._on_sensor_command,
            pb.SysCommand.FIVE_G_COMMAND: self._on_five_g_command,
            pb.SysCommand.MICRO_WAVE_COMMAND: self._on_micro_wave_command,
        }

    def _get_device(self, device_id: str):
        device = self.devices.get(device_id, None)
        if device is None:
            start = datetime.datetime.now()
            if not os.path.exists('./Images'):
                os.mkdir('./Images')
            path = f"./Images/{device_id}-{start.strftime('%Y-%m-%d-%H-%M-%S')}"
            if not os.path.exists(path):
                os.mkdir(path)
            device = dict(power_on=True, srart=start, path=path, images=0, micro_wave=0)
            self.devices[device_id] = device
        return device

    def _on_sys_command(self, req: pb.RequestCommand, resp: pb.ResponseCommand):
        if pb.SubCommand.POWER_ON_SUB_COMMAND == req.sub_command:
            print('power on message received')
            self._get_device(req.device_id)
        elif pb.SubCommand.POWER_OFF_SUB_COMMAND == req.sub_command:
            device = self.devices[req.device_id]
            device['power_on'] = False
            print('power off message received')
        return

    def _on_sensor_command(self, req: pb.RequestCommand, resp: pb.ResponseCommand):
        print('sensor command')

    def _write_image(self, device: dict, idx: int, image: bytes):
        if (image is None) | (len(image) == 0):
            return
        filename = f'{device["path"]}/{idx}-{device["images"]}.jpg'
        with open(filename, 'wb') as f:
            f.write(image)

    def _on_five_g_command(self, req: pb.RequestCommand, resp: pb.ResponseCommand):
        # print('five g command')
        if pb.SubCommand.FIVE_G_SEND_IMAGE_SUB_COMMAND == req.sub_command:
            print(
                f'send image sub command received, length={req.data1.length}')
            if req.data1.length > 0:
                device = self._get_device(req.device_id)
                img_msg = pb.ImagesMessage()
                img_msg.ParseFromString(req.data1.data)
                # print(img_msg.image1)
                count = device["images"]
                self._write_image(device, 1, img_msg.image1)
                self._write_image(device, 2, img_msg.image2)
                self._write_image(device, 3, img_msg.image3)
                self._write_image(device, 4, img_msg.image4)
                self._write_image(device, 5, img_msg.image5)
                device['images'] += 1

    def _on_micro_wave_command(self, req: pb.RequestCommand, resp: pb.ResponseCommand):
        print('micro wave command')
        if pb.SubCommand.MICRO_WAVE_SUB_COMMAND == req.sub_command:
            print('micro wave sub command received')
            device = self._get_device(req.device_id)
            micro_wave_msg = pb.MicroWaveMessage()
            if (micro_wave_msg.length == len(micro_wave_msg.data)):
                filename = f'{device["path"]}/{device["micro_wave"]}.micro_wave'
                with open(filename, 'wb') as f:
                    f.write(micro_wave_msg.data)
                device['micro_wave'] += 1

    def _default(self, req: pb.RequestCommand, resp: pb.ResponseCommand):
        print('onknown command')

    def _on_request(self, ch, method, properties, body):
        req = pb.RequestCommand()
        req.ParseFromString(body)
        dispatcher = self.routes.get(req.sys_command, self._default)
        resp = pb.ResponseCommand()
        dispatcher(req, resp)
        resp.checksum = 12345
        resp.state = pb.ResponseState.SUCCESS
        if properties.reply_to != None:
            ch.basic_publish(exchange='',
                             routing_key=properties.reply_to,  # 收消息的队列
                             properties=pika.BasicProperties(
                                 correlation_id=properties.correlation_id),  # 返回消息的队列
                             body=resp.SerializeToString())  # 返回结果数据
        ch.basic_ack(delivery_tag=method.delivery_tag)  # 确保消息被 客户端接收
        self.message_received += 1
        return

    def api(self):
        req = pb.RequestCommand()
        req.ParseFromString(request.data)
        dispatcher = self.routes.get(req.sys_command, self._default)
        resp = pb.ResponseCommand()
        dispatcher(req, resp)
        resp.checksum = 12345
        response = Flask.response_class(
            resp.SerializeToString(),
            mimetype='application/x-protobuf'
        )
        return response, 200

    def hello(self):
        return 'Hello, World!'

    def consuming(self):
        print(' [*] RabbitMQ Server started. Waiting for messages.')
        self.channel.start_consuming()

    def run_rabbitmq(self):
        t = threading.Thread(target=self.consuming)
        t.start()

    def run_all(self, host="localhost", port=5000):
        self.run_rabbitmq()
        self.run(debug=False, host=host, port=port)


if __name__ == "__main__":
    server = MyServer(__name__)
    server.init_rabbmitmq(host='localhost', user='chai', password='password123', queue='aerial_rpc')
    server.run_all(host="0.0.0.0", port=5000)