from common_code.vector_utils import upload_vectors_to_db
from common_code import kwq, util, file_names, vector_utils
import time
from multiprocessing import Process, Queue, Manager
import threading

topic_subscribe = 'upload_to_vector_db'
group_id = "vision_v3_uploaders"

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

def watch_and_upload(rank, health_queue):
    start_time = time.time()
    client_suffix = '_'+str(rank)
    client_id = 'uploader_to_vector_db'+client_suffix

    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:
            score_this = get_score_file(result)
            try:
                upload_vectors_to_db(score_this)
                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 = 15  # 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()

