import time
import cv2
from confluent_kafka import Producer
from confluent_kafka.admin import AdminClient, NewTopic
import torch.multiprocessing as mp
import zmq
import socket as s
import numpy as np
import torch
import argparse
import configparser
from PIL import Image
import io
import datetime
import threading
from multiprocessing import Process
import os

def creator():
    folder_name = "/tmp/images/fire/"

    try:
        os.makedirs(folder_name)
        print(f"Folder '{folder_name}' created successfully.")
    except FileExistsError:
        print(f"Folder '{folder_name}' already exists.")


conf = {
    'bootstrap.servers': 'localhost:29092'
}
producer = Producer(conf)
def main():
    topic = 'fire'
    try:
        mp.set_start_method('spawn', force=True)
    except RuntimeError:
        pass
    port = []
    config.read('/usr/lib/python3.8/dist-packages/firedetection/fire_config.ini')
    for section in config.sections():
        if section.startswith('camera_'):
            value = config.get(section, 'port')
            port.append(value)
    print(port)
    creator()
    check_kafka(topic, len(port))
    TIME_ = float(config['timeout']['value'])
    PRODUCERS = [StreamVideo(index, port, TIME_) for index, port in enumerate(port)]
    for p in PRODUCERS:
        p.start()

    for p in PRODUCERS:
        p.join()

def check_kafka(topic, num):
    a = AdminClient({'bootstrap.servers': 'localhost:29092'})
    new_topics = [NewTopic(topic, num_partitions=num,
                           replication_factor=1) for topic in [str(topic)]]
    fs = a.create_topics(new_topics)
    for topic, f in fs.items():
        try:
            f.result()
            print("Topic {} created".format(topic))
        except Exception as e:
            print("Failed to create topic {}: {}".format(topic, e))


def delivery_report(err, msg):
    if err is not None:
        print('Message delivery failed: {}'.format(err))
    else:
        print('Message delivered to {} [{}]'.format(
            msg.topic(), msg.partition()))



DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S-%f"
config = configparser.ConfigParser()


def event(producer, cam_num):
    producer.produce('fire',  str("Camera  " + str(cam_num+1)+ "Обнаружен огонь").encode('utf-8'), partition=cam_num,callback=delivery_report)


def streamer(camera, socket):
    import base64
    try:
        obj = socket.recv_string()
        source_id = obj[:len(str(camera))]
        frame = obj[len(str(camera)):]
        if source_id == str(camera):
            img = base64.b64decode(frame)
            npimg = np.frombuffer(img, dtype=np.uint8)
            source = cv2.imdecode(npimg, 1)
            return source
    except Exception as e:
        print(str(e))


def process(producer, id, frame ,model):
        print(str(id) + " scan")
        #frame = streamer(id, socket)
        results = model(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        if 'fire' in results.pandas().xyxy[0]['name'].to_list():
                    event(producer, id)
                    results.render()
                    now_time = datetime.datetime.now().strftime(DATETIME_FORMAT)
                    img_savename = f"FIRE_{now_time}.png"
                    Image.fromarray(results.ims[0]).save("/tmp/images/fire/"+img_savename)
    

     

class StreamVideo(Process):

    def __init__(self, cam_num, port, interval):
        super(StreamVideo, self).__init__()
        self.interval = interval
        self.cam_num = cam_num
        self.port = port
        self.ret = True
        self.frame = None
        self.socket = None

    def run(self):
        self.model = torch.hub.load('ultralytics/yolov5', 'custom',
                            path='/usr/lib/python3.8/dist-packages/firedetection/train.pt', force_reload=True)
        context = zmq.Context()
        self.socket =  context.socket(zmq.SUB)
        self.socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self.socket.connect('tcp://127.0.0.1:'+str(self.port))
        print('Init' + ' Camera '+str(self.cam_num))
        self.frame = streamer(self.cam_num, self.socket)
        self.later=time.time()
        while True:
            self.frame = streamer(self.cam_num, self.socket)
            if int(abs(time.time()-self.later))==self.interval:
                process(producer,self.cam_num,self.socket,self.frame,self.model)
                self.later=time.time()
                

