"""AgentRuntime - Manages agent lifecycle, threading, and graceful shutdown."""

import logging
import time
import json
import threading
from typing import Optional, Dict, Any, Callable


class AgentRuntime:
    """Manages agent runtime lifecycle and threading."""
    
    def __init__(
        self,
        name: str,
        operational_cycle: Callable,
        perception_cycle: Optional[Callable] = None,
        lifecycle_hooks: Optional[Any] = None,
        executor: Optional[Any] = None,
    ):
        """Initialize agent runtime.
        
        Args:
            name: Agent name
            operational_cycle: Callable for main operational cycle
            perception_cycle: Optional callable for perception cycle
            lifecycle_hooks: LifecycleHooks instance
            executor: ThreadPoolExecutor instance (optional)
        """
        self.name = name
        self.operational_cycle = operational_cycle
        self.perception_cycle = perception_cycle
        self.lifecycle_hooks = lifecycle_hooks
        self.executor = executor
        
        self.worker_thread: Optional[threading.Thread] = None
        self.perception_thread: Optional[threading.Thread] = None
    
    def start(self) -> None:
        """Start agent worker threads."""
        logging.info(f"[RUNTIME] {self.name}: Starting agent...")
        
        # Invoke on_start hook
        if self.lifecycle_hooks:
            self.lifecycle_hooks.invocar("on_start", {
                "agent": self.name,
                "timestamp": time.time()
            })
        
        # Start worker thread
        self.worker_thread = threading.Thread(
            target=self.operational_cycle,
            name=f"{self.name}_Worker",
            daemon=True
        )
        self.worker_thread.start()
        logging.info(f"[RUNTIME] {self.name}: Worker thread started")
        
        # Start perception thread (if defined)
        if self.perception_cycle:
            self.perception_thread = threading.Thread(
                target=self.perception_cycle,
                name=f"{self.name}_Perception",
                daemon=True
            )
            self.perception_thread.start()
            logging.info(f"[RUNTIME] {self.name}: Perception thread started")
    
    def is_running(self) -> bool:
        """Check if agent threads are running."""
        if self.worker_thread is None:
            return False
        return self.worker_thread.is_alive()
    
    def graceful_shutdown(
        self,
        inbox: Any,
        timeout: float = 30.0,
        metrics: Optional[Dict[str, Any]] = None,
        outbox: Optional[Any] = None,
        incident_log: Optional[Any] = None,
        health_probe: Optional[Any] = None,
        auditor: Optional[Any] = None,
        metrics_collector: Optional[Any] = None,
        alive_flag_getter: Optional[Callable[[], bool]] = None,
        capacity: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Gracefully shutdown agent.
        
        Steps:
        1. Stop signal (alive=False)
        2. Wait for queue to empty (until timeout)
        3. Wait for threads to finish
        4. Save final state
        5. Clean resources
        
        Args:
            inbox: PriorityQueue instance
            timeout: Maximum seconds to wait (default 30s)
            metrics: Global metrics dict
            outbox: Outbox deque
            incident_log: Incident log deque
            health_probe: Health probe instance
            auditor: Auditor instance
            metrics_collector: Metrics collector instance
            alive_flag_getter: Callable that returns alive status
            capacity: Queue capacity for readiness probe
        
        Returns:
            Dict with shutdown statistics
        """
        logging.info(f"[SHUTDOWN] {self.name}: Starting graceful shutdown...")
        shutdown_start = time.time()
        
        # 1. Signal stop (via alive_flag_getter if provided)
        # Note: actual alive flag should be set by caller
        
        # Invoke on_stop hook
        if self.lifecycle_hooks:
            self.lifecycle_hooks.invocar("on_stop", {
                "agent": self.name,
                "timestamp": time.time()
            })
        
        # 2. Wait for queue to empty (with timeout)
        logging.info(f"[SHUTDOWN] {self.name}: Waiting for queue to drain...")
        drain_start = time.time()
        while not inbox.empty() and (time.time() - drain_start) < (timeout * 0.5):
            queue_size = inbox.qsize()
            if queue_size > 0 and queue_size % 10 == 0:
                logging.debug(f"[SHUTDOWN] {self.name}: Queue has {queue_size} pending tasks...")
            time.sleep(0.1)
        
        # 3. Wait for worker threads to finish
        logging.info(f"[SHUTDOWN] {self.name}: Waiting for threads to finish...")
        wait_time = timeout - (time.time() - shutdown_start)
        
        if self.worker_thread and self.worker_thread.is_alive():
            self.worker_thread.join(timeout=wait_time * 0.5)
            if self.worker_thread.is_alive():
                logging.warning(f"[SHUTDOWN] {self.name}: Worker thread did not finish in time")
        
        if self.perception_thread and self.perception_thread.is_alive():
            self.perception_thread.join(timeout=wait_time * 0.5)
            if self.perception_thread.is_alive():
                logging.warning(f"[SHUTDOWN] {self.name}: Perception thread did not finish in time")
        
        # 4. Shutdown executor
        if self.executor:
            logging.info(f"[SHUTDOWN] {self.name}: Shutting down executor...")
            self.executor.shutdown(wait=True, cancel_futures=False)
        
        # 5. Save final state
        logging.info(f"[SHUTDOWN] {self.name}: Saving final state...")
        try:
            if outbox:
                success_list = [dict(r) for r in outbox]
                with open(f"successes_{self.name}_shutdown.json", "w") as f:
                    json.dump(success_list, f, indent=2)
            
            if incident_log:
                incident_list = [dict(r) for r in incident_log]
                with open(f"incidents_{self.name}_shutdown.json", "w") as f:
                    json.dump(incident_list, f, indent=2)
            
            if auditor:
                with open(f"audit_{self.name}_shutdown.json", "w") as f:
                    f.write(auditor.exportar_json())
        
        except Exception as e:
            logging.error(f"[SHUTDOWN] {self.name}: Error saving state: {e}")
        
        # Calculate final statistics
        total_time = time.time() - shutdown_start
        stats = {
            "agent": self.name,
            "timestamp_shutdown": time.time(),
            "shutdown_seconds": round(total_time, 2),
            "successful_tasks": metrics.get("successes", 0) if metrics else 0,
            "failed_tasks": metrics.get("fatal_failures", 0) if metrics else 0,
            "recovered_tasks": metrics.get("recovered_failures", 0) if metrics else 0,
            "rejected_tasks": metrics.get("rejected", 0) if metrics else 0,
            "dead_letter_size": 0,  # Caller should compute this
            "pending_queue": inbox.qsize(),
            "active_threads": sum([
                1 if self.worker_thread and self.worker_thread.is_alive() else 0,
                1 if self.perception_thread and self.perception_thread.is_alive() else 0
            ]),
        }
        
        if metrics_collector:
            stats["complete_metrics"] = metrics_collector.obtener_estado()
        
        logging.info(f"[SHUTDOWN] {self.name}: Shutdown completed in {total_time:.2f}s")
        return stats
