
import numpy as np
import time
from multiprocessing import Process, Queue, Manager
import threading
from common_code import kwq, util, file_names, vector_utils, video_meta

from datetime import datetime
from clickhouse_connect import get_client


topic_subscribe = 'upload_to_ch_db'
group_id = "vision_v3_uploaders"

def get_score_file(f):
    return f.key().decode().strip('"')



def insert_into_ch(file_path, client):
    cset = util.get_cset_match(file_path)
    ff = np.load(cset['.detection.npz'], allow_pickle = True)
    det_results = dict(ff)
    del_idx = list()

    filt_scores = det_results['final_scores']
    filt_frames = det_results['final_frames']
    spec_label_idx = ff['species_label_map'].tolist()

    ftp_vid = video_meta.FTPVideo(cset)

    frame_num_off = ftp_vid.sorted_frames_info
    frame_period = ftp_vid.frameperiod
    frame_rate = ftp_vid.framerate

    indexes = np.searchsorted(ff['scored_frames'], filt_frames)
    dff = np.diff(ff['scored_frames'])

    head = np.insert(dff, 0, 0)
    tail = np.append(dff, 0)
    all_frame_durations = frame_period*np.mean(np.stack([head,tail]), axis=0)


    frame_file = [ftp_vid.cpath.split('/')[-1].replace('_reduced','').replace('.mp4','')]*len(filt_frames)
    frame_bboxes = ff['final_boxes']
    frame_numbers = filt_frames
    frame_scores = filt_scores
    frame_species = [spec_label_idx[int(x)] for x in ff['final_labels']]
    frame_durations = all_frame_durations[indexes]
    frame_tstamps = [frame_num_off[x]['time'] for x in filt_frames]
    db_insert_times = [datetime.now()]*len(filt_frames)

    a = frame_numbers
    starts = np.r_[0, np.flatnonzero(a[1:] != a[:-1]) + 1]
    det_count_num = np.arange(len(a)) - np.repeat(starts, np.diff(np.r_[starts, len(a)]))
    camera =  [ftp_vid.camera_name] * len(filt_frames)


    result = client.insert('nuggets.cameras_object_detection_results',
                  column_oriented = True,
                  column_names = ['camera','video_filename','frame_number','detection_number','classification_score','object_name','bbox_xyxy','frame_duration','frame_timestamp','db_insert_time'],
                  data = [camera, frame_file, frame_numbers, det_count_num, frame_scores, frame_species, frame_bboxes.tolist(), frame_durations, frame_tstamps, db_insert_times]
                  )


# %%


def watch_and_upload(rank, health_queue):
    client = get_client(host='192.168.1.242', compression=True)
    start_time = time.time()
    client_suffix = '_'+str(rank)
    client_id = 'uploader_to_ch_db'+client_suffix
    current_time = time.time()    
    health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
    health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
    
    queue = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'earliest')
    queue.subscribe([topic_subscribe])

    last_heartbeat = time.time()
    
    while True:
        current_time = time.time()
        
        # # Exit after 5 minutes
        # if (current_time - start_time) > 60*5:
        #     print('###########################')
        #     print('######### EXITING #########')
        #     print('###########################')            
        #     return

        # Send heartbeat every 10 seconds
        if current_time - last_heartbeat > 5:
            try:
                health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
                last_heartbeat = current_time
            except:
                # Queue might be full, continue working
                pass
        
        result = queue.poll(timeout=0.05)    
        if result is not None:
            try:
                score_this = get_score_file(result)                
                insert_into_ch(score_this, client)
                print('#' + str(rank) + ' Uploaded to ' + score_this)
                # Send success heartbeat
                try:
                    health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'working'})
                except:
                    pass
            except Exception as e:
                print('Failed ' + score_this + ':' + str(e))
                
            queue.commit()

def health_monitor(processes, health_queues, num_workers):
    """Monitor process health and restart hung processes"""
    last_heartbeat = {}
    HEALTH_TIMEOUT = 120  # Consider process hung if no heartbeat for 60 seconds
    
    while True:
        time.sleep(5)  # Check every 5 seconds
        current_time = time.time()
        
        for rank in range(num_workers):
            # Drain the health queue for this process
            latest_heartbeat = None
            try:
                while True:
                    try:
                        heartbeat = health_queues[rank].get_nowait()
                        if heartbeat['rank'] == rank:
                            latest_heartbeat = heartbeat['timestamp']
                    except:
                        break
            except:
                pass
            
            # Update last known heartbeat
            if latest_heartbeat:
                last_heartbeat[rank] = latest_heartbeat
            
            # Check if process is alive and responsive
            process = processes[rank]
            last_seen = last_heartbeat.get(rank, 0)
            
            if not process.is_alive():
                print(f"Process {rank} has died. Restarting...")
                processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
                processes[rank].start()
                last_heartbeat[rank] = current_time
            elif current_time - last_seen > HEALTH_TIMEOUT:
                print(f"Process {rank} appears hung (no heartbeat for {current_time - last_seen:.1f}s). Restarting...")
                process.terminate()
                process.join(timeout=5)
                if process.is_alive():
                    process.kill()
                    process.join()
                
                # Create new process
                processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
                processes[rank].start()
                last_heartbeat[rank] = current_time

# Main execution
if __name__ == "__main__":
    num_workers = 4
    
    # Create health queues for each worker
    health_queues = [Queue(maxsize=10) for _ in range(num_workers)]
    
    # Create and start worker processes
    processes = []
    for rank in range(num_workers):
        p = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
        p.start()
        processes.append(p)
    
    # Start health monitor in a separate thread
    monitor_thread = threading.Thread(
        target=health_monitor, 
        args=(processes, health_queues, num_workers),
        daemon=True
    )
    monitor_thread.start()
    
    try:
        # Wait for all processes to complete (they will exit after 5 minutes)
        for p in processes:
            p.join()
        print("All worker processes completed")
    except KeyboardInterrupt:
        print("Shutting down...")
        for p in processes:
            p.terminate()
            p.join(timeout=5)
            if p.is_alive():
                p.kill()
                p.join()

