import time
import traceback
from datetime import datetime
from multiprocessing import Process, Queue 
import logging
import cv2
import numpy as np
from nodes.Node import Node
import ffmpeg

import os

from nodes.FFVideoWriterNode import FFVideoWriterNode

class FishFFVideoWriterNode(FFVideoWriterNode):
    def __init__(self, name="FishFFVideoWriterNode", parent=None):
        super().__init__(name, parent)

        self._videoFileName = None

        self._currVideoId = 0
        self._currObjectId = 0
        self._currFrameIndex = 0
        self._lastFrameIndex = 0
        self._startFrameIndex = 0

        self._logFileName="../output/fish.txt"
        self._logFile = None
        self.is_opened = False

    def open(self):
        self._startFrameIndex = self._currFrameIndex
        filename = "{}VIDEO_{:04}_{:05}.mp4".format(self._videoPath,self._currVideoId,self._currObjectId)
        self._videoFileName = filename

        logging.info(f'VideoWriter: FileName:{filename}')

        #fourcc = cv2.VideoWriter_fourcc('H','E','V','C')
        #fourcc = cv2.VideoWriter_fourcc('H','2','6','4')
        #fourcc = cv2.VideoWriter_fourcc(*'H264')

        #self.writer = cv2.VideoWriter(filename,cv2.VideoWriter_fourcc('M', 'P', 'E', 'G'), 
        #                                int(self.framerate), (int(self.framesize[0]), int(self.framesize[1])))

        #self.writer = cv2.VideoWriter(filename, fourcc, int(self._frameRate),(int(self._frameSize[0]), int(self._frameSize[1])))

        # Set up the FFMPEG process
        width = self._frameSize[0]
        height = self._frameSize[1]

        self._process = (
            ffmpeg
            .input('pipe:0', format='rawvideo', s=f'{width}x{height}', framerate=29, pix_fmt='rgb24')
            .output(filename, vcodec='h264_nvenc', pix_fmt='yuv420p')
            .run_async(pipe_stdin=True)
        )

        logging.info("Video Writer Opened.")
        self.is_opened = True

    def setLogFileName(self, logFile):
        self._logFileName=logFile

    def openLog(self):
        self._logFile=open(self._logFileName,"w")

    def close(self):
        time.sleep(3)
        #self.writer.release()
        # Close stdin and wait for the process to finish
        self._logFile.write("{}\t{}\t{}\t{}\t{}\n".format(self._videoFileName,self._currVideoId,self._currObjectId,self._startFrameIndex, self._lastFrameIndex))
        self._logFile.flush()
        if self._process is not None:
            self._process.stdin.close()
            self._process.wait()
        logging.info("FFVideoWriter:Close Video Writer.")
        self.is_opened = False


    def run(self):
        #self.open()
        self.setStatus(Node.NODE_STATUS_RUNNING)
        self.is_stopped = False

        empty_counter = 0

        while True:
            if self.is_stopped:
                if self.is_opened:
                    self.close()
                break
            else:
                if( not self._framePool.empty()):

                    try:
                        #self.evalStart()
                        frame, videoIndex, frameIndex, tracks = self._framePool.get()

                        self._currFrameIndex = frameIndex

                        if videoIndex != self._currVideoId:
                            if self.is_opened:
                                self.close()
                            self._currVideoId = videoIndex

                        if len(tracks) < 1:
                            empty_counter += 1
                            if empty_counter > 5 and self.is_opened:
                                self.close()
                            continue
                        else:
                            empty_counter = 0
                            if not self.is_opened:
                                # find smallest object id
                                identities = tracks[:, 8]
                                sid = identities[0]
                                for id in identities:
                                    if id < sid:
                                        sid = id
                                self._currObjectId = int(sid)
                                self._startFrameIndex = frameIndex

                                self.open()

                        self._lastFrameIndex = frameIndex
                        # print(f"StartFrameIndex:{self._startFrameIndex}, LastFrameIndex:{self._lastFrameIndex}")
                        # Write frame data to FFMPEG process
                        self._process.stdin.write(frame)
                        #self.evalEnd()
                        logging.info("VideoWriterNode:Writer consume a frame.")
                    except Exception as e:
                        logging.error("Writing Video Error:")
                        traceback.print_exc()
                        self.close()
                        self.open()
                        self._videoStartTime=time.time()
                else:
                    if self.getParentStatus()==Node.NODE_STATUS_EXITED:
                        self.setStatus(Node.NODE_STATUS_EXITED)
                        logging.info("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
                        logging.info("FishFFVideoWriter is completed.")
                        break