import sys
import argparse
import logging
import faulthandler
from pathlib import Path

from nodes.Node import Node
from nodes.SourceNode import SourceNode, VideoSourceNode, RtspSourceNode
from nodes.SinkNode import VideoSinkNode, TestSinkNode, PlayerSinkNode, RtmpSinkNode
from nodes.FFVideoWriterNode import FFVideoWriterNode
from nodes.YOLONode import YOLONode
from nodes.SORTTrackNode import SORTTrackNode
from nodes.FlaskSinkNode import FlaskSinkNode
from nodes.VideoListSourceNode import VideoListSourceNode
from nodes.utils.shm import *
from FishVideoListSourceNode import FishVideoListSourceNode
from FishSORTTrackNode import FishSORTTrackNode
from FishFFVideoWriterNode import FishFFVideoWriterNode
from FishYOLONode import FishYOLONode
from FishBatchYOLONode import FishBatchYOLONode
from FishFilterNode import FishFilterNode
from utils import *
#from FlaskFetcher import fishFetcher

level = logging.ERROR
logging.basicConfig(level=level)

faulthandler.enable()

def createAppDict(key,nc):
    dict ={}
    for i in range(nc):
        dict[i]=0

    CreateAppSharedMemory(key)

    StoreAppDict(key, dict)

def getParams():
    # Create an ArgumentParser object
    parser = argparse.ArgumentParser(description="A simple argument parser")

    # Define the arguments you expect
    parser.add_argument('--input', type=str, help='A video folder or a video list file.', default="../input")
    parser.add_argument('--output', type=str, help='A output folder.', default="../output")
    parser.add_argument('--appkey', type=str, help='A unique string for communication', default="/FisherKey")
    parser.add_argument('--port', type=int, help="Video port.", default=8042)

    # Parse the arguments
    args = parser.parse_args()

    return args

def main():
    #CreateSharedMemory(SHM_VIDEO_NAME);
    args = getParams()

    strinput = args.input
    stroutput = args.output
    fport = args.port
    appkey = args.appkey

    createAppDict(appkey, 5)
    dict = LoadAppDict(appkey)
    print('AppDict:', dict)

    processes = []

    source = FishVideoListSourceNode()
    source.listVideos(strinput)
    source.setId(0)
    source.setAppKey(appkey)
    source.setStatus(1)
    source.saveVideoList(stroutput + "/videolist.txt")
    source.setSourceCloseAction(source.closeAction)

    yolo = FishBatchYOLONode()
    yolo.setId(1)
    yolo.setAppKey(appkey)
    yolo.setBatchSize(50)
    yolo.updatePoolSize(3000)
    yolo.setModelFile("../FishZHX20241204A.pt")

    sort = FishSORTTrackNode()
    sort.setId(2)
    sort.setAppKey(appkey)
    sort.updatePoolSize(1000)

    filter = FishFilterNode()
    filter.setId(3)
    filter.setAppKey(appkey)
    filter.updatePoolSize(1000)

    #fsink = FlaskSinkNode()
    #fsink.updatePoolSize(500)
    #fsink.setPort(fport)

    vsink = FishFFVideoWriterNode()
    vsink.setId(4)
    vsink.setAppKey(appkey)
    vsink.updatePoolSize(30000)
    vsink.setFrameSize((960,540))
    vsink.setVideoPath(stroutput)
    vsink.setLogFileName(stroutput + "/fish.txt")
    vsink.openLog()


    #fsink.setParent(filter)
    vsink.setParent(filter)

    filter.setParent(sort)

    sort.setParent(yolo)
    yolo.setParent(source)

    #fsink.start()
    vsink.start()
    processes.append(vsink)
    filter.start()
    processes.append(filter)
    sort.start()
    processes.append(sort)
    yolo.start()
    processes.append(yolo)
    source.start()
    processes.append(source)

    # Wait for any process to finish
    while processes:
        for process in processes:
            if not process.is_alive():  # Check if the process has finished
                process.join()  # Wait for it to finish
                processes.remove(process)  # Remove the finished process
                print(f"Process {process.getNodeName()} has completed")
                break  # Exit after the first process finishes

    print("Main program finished")

    while(True):
        print("Wait source to join")
        source.join(timeout=10)  # Wait for a maximum of 10 seconds
 
        if source.is_alive():
            print("Process is still alive, something might be wrong.")
        else:
            print(">>>>>>>>>>>>>>>>Process Source completed successfully.<<<<<<<<<<<<<<<<<<")
            break

    while(True):
        print("Wait YOLO to join")
        yolo.join(timeout=10)  # Wait for a maximum of 10 seconds

        if yolo.is_alive():
            print("Process is still alive, something might be wrong.")
        else:
            print(">>>>>>>>>>>>>>>>Process YOLO completed successfully.<<<<<<<<<<<<<<<<<<")
            break

  
    #print(">>>>>>>>>>>>>>>>Process YOLO completed successfully.<<<<<<<<<<<<<<<<<<")
    sort.join()
    print(">>>>>>>>>>>>>>>>Process SORT completed successfully.<<<<<<<<<<<<<<<<<<")
    filter.join()
    print(">>>>>>>>>>>>>>>>Process Filter completed successfully.<<<<<<<<<<<<<<<<<<")
    vsink.join()
    print(">>>>>>>>>>>>>>>>Process SINK completed successfully.<<<<<<<<<<<<<<<<<<")
    #fsink.join()
    dict = LoadAppDict(appkey)
    print('AppDict2:', dict)
    ClearAppSharedMemory(appkey)

if __name__ == "__main__":
    main()