# app/synchronizer.py
import logging
import time
import threading
from typing import Dict, List, Optional, Tuple, Union, Any, Callable
from datetime import datetime, timedelta 
import queue
import uuid

class Synchronizer:
    def __init__(self):
        """Initialize the Synchronizer for coordinating agent interactions."""
        self.logger = logging.getLogger(__name__)
        
        # Synchronization barriers
        self.barriers = {}
        
        # Completion signals
        self.completion_signals = {}
        
        # Dependency tracking
        self.dependencies = {}
        
        # Locks for thread safety
        self.barrier_lock = threading.RLock()
        self.signal_lock = threading.RLock()
        
        # Event queues for agents
        self.event_queues = {}
        
        # Callback registry
        self.callbacks = {}
        
    def register_agent(self, agent_id: str) -> None:
        """Register an agent with the synchronizer."""
        with self.barrier_lock:
            if agent_id not in self.event_queues:
                self.event_queues[agent_id] = queue.Queue()
                self.logger.info(f"Registered agent: {agent_id}")
    
    def create_barrier(self, barrier_id: str, participants: List[str]) -> str:
        """
        Create a synchronization barrier that multiple agents must reach.
        Returns a unique barrier instance ID.
        """
        instance_id = f"{barrier_id}_{uuid.uuid4().hex[:8]}"
        
        with self.barrier_lock:
            self.barriers[instance_id] = {
                "participants": set(participants),
                "arrived": set(),
                "event": threading.Event(),
                "created_at": datetime.now().isoformat()
            }
            
        self.logger.info(f"Created barrier {instance_id} with {len(participants)} participants")
        return instance_id
    
    def arrive_at_barrier(self, barrier_id: str, agent_id: str, data: Any = None) -> bool:
        """
        Signal that an agent has arrived at a barrier.
        Returns True if all participants have arrived, False otherwise.
        """
        with self.barrier_lock:
            if barrier_id not in self.barriers:
                self.logger.warning(f"Barrier {barrier_id} not found")
                return False
                
            barrier = self.barriers[barrier_id]
            
            # Check if agent is a participant
            if agent_id not in barrier["participants"]:
                self.logger.warning(f"Agent {agent_id} is not a participant in barrier {barrier_id}")
                return False
                
            # Record arrival
            barrier["arrived"].add(agent_id)
            
            # Store data if provided
            if data is not None:
                if "data" not in barrier:
                    barrier["data"] = {}
                barrier["data"][agent_id] = data
                
            # Check if all participants have arrived
            all_arrived = barrier["arrived"] == barrier["participants"]
            
            if all_arrived:
                # Signal all waiting threads
                barrier["event"].set()
                self.logger.info(f"All participants arrived at barrier {barrier_id}")
                
                self.logger.info(f" barrier participants: {barrier['participants']}")
                
                # Notify agents via their event queues
                for participant in barrier["participants"]:
                    if participant in self.event_queues:
                        self.event_queues[participant].put({
                            "type": "barrier_complete",
                            "barrier_id": barrier_id,
                            "timestamp": datetime.now().isoformat()
                        })
                
                # Execute callbacks if registered
                callback_key = f"barrier:{barrier_id}"
                if callback_key in self.callbacks:
                    try:
                        self.callbacks[callback_key](barrier_id, barrier.get("data", {}))
                    except Exception as e:
                        self.logger.error(f"Error executing callback for {callback_key}: {e}")
                
            return all_arrived
    
    def wait_for_barrier(self, barrier_id: str, timeout: Optional[float] = None) -> bool:
        """
        Wait until all participants arrive at the barrier.
        Returns True if all arrived, False on timeout.
        """
        with self.barrier_lock:
            if barrier_id not in self.barriers:
                self.logger.warning(f"Barrier {barrier_id} not found")
                return False
                
            barrier = self.barriers[barrier_id]
            event = barrier["event"]
            
        # Release lock before waiting
        return event.wait(timeout)
    
    def get_barrier_data(self, barrier_id: str) -> Dict[str, Any]:
        """Get data provided by agents at a barrier."""
        with self.barrier_lock:
            if barrier_id not in self.barriers:
                self.logger.warning(f"Barrier {barrier_id} not found")
                return {}
                
            barrier = self.barriers[barrier_id]
            return barrier.get("data", {}).copy()
    
    def signal_completion(self, agent_id: str, task_id: str, result: Any = None) -> None:
        """Signal that an agent has completed a task."""
        with self.signal_lock:
            key = f"{agent_id}:{task_id}"
            self.completion_signals[key] = {
                "completed_at": datetime.now().isoformat(),
                "result": result
            }
            
            # Check if any dependencies are waiting on this
            for dep_key, deps in self.dependencies.items():
                if key in deps["waiting_on"]:
                    deps["waiting_on"].remove(key)
                    
                    # If no more dependencies, signal completion
                    if not deps["waiting_on"]:
                        waiting_agent, waiting_task = dep_key.split(":", 1)
                        
                        # Notify the waiting agent
                        if waiting_agent in self.event_queues:
                            self.event_queues[waiting_agent].put({
                                "type": "dependencies_met",
                                "task_id": waiting_task,
                                "timestamp": datetime.now().isoformat()
                            })
                        
                        # Execute callback if registered
                        callback_key = f"dependency:{dep_key}"
                        if callback_key in self.callbacks:
                            try:
                                self.callbacks[callback_key](dep_key)
                            except Exception as e:
                                self.logger.error(f"Error executing callback for {callback_key}: {e}")
            
        self.logger.info(f"Agent {agent_id} completed task {task_id}")
    
    def wait_for_completion(self, agent_id: str, task_id: str, 
                           timeout: Optional[float] = None) -> Tuple[bool, Any]:
        """
        Wait for an agent to complete a task.
        Returns (completed, result) tuple.
        """
        key = f"{agent_id}:{task_id}"
        end_time = time.time() + timeout if timeout is not None else None
        
        while True:
            with self.signal_lock:
                if key in self.completion_signals:
                    return True, self.completion_signals[key].get("result")
            
            # Check timeout
            if end_time is not None and time.time() >= end_time:
                return False, None
                
            # Sleep briefly to avoid busy waiting
            time.sleep(0.1)
    
    def register_dependencies(self, agent_id: str, task_id: str, 
                             dependencies: List[Tuple[str, str]]) -> None:
        """
        Register that a task depends on completion of other tasks.
        Dependencies are specified as [(agent_id, task_id), ...].
        """
        waiting_on = set()
        for dep_agent, dep_task in dependencies:
            dep_key = f"{dep_agent}:{dep_task}"
            
            # If dependency is already completed, don't wait on it
            with self.signal_lock:
                if dep_key not in self.completion_signals:
                    waiting_on.add(dep_key)
        
        key = f"{agent_id}:{task_id}"
        
        with self.signal_lock:
            self.dependencies[key] = {
                "waiting_on": waiting_on,
                "registered_at": datetime.now().isoformat()
            }
            
        self.logger.info(f"Registered {len(waiting_on)} dependencies for {key}")
    
    def are_dependencies_met(self, agent_id: str, task_id: str) -> bool:
        """Check if all dependencies for a task are met."""
        key = f"{agent_id}:{task_id}"
        
        with self.signal_lock:
            if key not in self.dependencies:
                return True  # No dependencies registered
                
            return len(self.dependencies[key]["waiting_on"]) == 0
    
    def register_callback(self, event_type: str, event_id: str, callback: Callable) -> None:
        """Register a callback function for a specific event."""
        key = f"{event_type}:{event_id}"
        self.callbacks[key] = callback
        self.logger.info(f"Registered callback for {key}")
    
    def get_events(self, agent_id: str, timeout: Optional[float] = 0) -> List[Dict[str, Any]]:
        """
        Get events for an agent from its queue.
        Non-blocking by default (timeout=0).
        """
        if agent_id not in self.event_queues:
            return []
            
        events = []
        queue = self.event_queues[agent_id]
        
        try:
            # Try to get at least one event
            event = queue.get(block=(timeout is not None), timeout=timeout)
            events.append(event)
            
            # Get any additional events that are already in the queue
            while not queue.empty():
                try:
                    event = queue.get_nowait()
                    events.append(event)
                except queue.Empty:
                    break
        except queue.Empty:
            pass
            
        return events
    
    def send_event(self, agent_id: str, event_type: str, data: Any = None) -> bool:
        """Send a custom event to an agent."""
        if agent_id not in self.event_queues:
            self.logger.warning(f"Agent {agent_id} not registered")
            return False
            
        event = {
            "type": event_type,
            "timestamp": datetime.now().isoformat(),
            "data": data
        }
        
        self.event_queues[agent_id].put(event)
        self.logger.info(f"Sent {event_type} event to {agent_id}")
        return True
    
    def cleanup(self, older_than_hours: Optional[int] = None) -> Dict[str, int]:
        """
        Clean up old synchronization artifacts.
        Returns counts of items removed.
        """
        if older_than_hours is None:
            # Clear everything
            with self.barrier_lock:
                barriers_count = len(self.barriers)
                self.barriers = {}
                
            with self.signal_lock:
                signals_count = len(self.completion_signals)
                dependencies_count = len(self.dependencies)
                self.completion_signals = {}
                self.dependencies = {}
                
            return {
                "barriers": barriers_count,
                "signals": signals_count,
                "dependencies": dependencies_count
            }
        
        # Calculate cutoff time
        cutoff = datetime.now() - timedelta(hours=older_than_hours)
        cutoff_str = cutoff.isoformat()
        
        # Clean up old barriers
        barriers_removed = 0
        with self.barrier_lock:
            to_remove = []
            for barrier_id, barrier in self.barriers.items():
                if barrier["created_at"] < cutoff_str:
                    to_remove.append(barrier_id)
                    
            for barrier_id in to_remove:
                del self.barriers[barrier_id]
                barriers_removed += 1
        
        # Clean up old signals and dependencies
        signals_removed = 0
        dependencies_removed = 0
        
        with self.signal_lock:
            # Clean signals
            to_remove = []
            for key, signal in self.completion_signals.items():
                if signal["completed_at"] < cutoff_str:
                    to_remove.append(key)
                    
            for key in to_remove:
                del self.completion_signals[key]
                signals_removed += 1
                
            # Clean dependencies
            to_remove = []
            for key, dep in self.dependencies.items():
                if dep["registered_at"] < cutoff_str:
                    to_remove.append(key)
                    
            for key in to_remove:
                del self.dependencies[key]
                dependencies_removed += 1
        
        return {
            "barriers": barriers_removed,
            "signals": signals_removed,
            "dependencies": dependencies_removed
        }