"""
BGX Executor module - Smart thread/process pool manager.

This module provides intelligent task execution with priority queues,
concurrency limiting per tag, and auto-detection of execution mode.
"""

import concurrent.futures
import queue
import threading
import time
from collections import defaultdict
from typing import Callable, Dict, Optional, Set

from .task import Task, TaskState, Priority


class TaskExecutor:
    """Smart executor with priority queues and concurrency control."""
    
    def __init__(
        self,
        workers: int = None,
        process_workers: int = None,
        mode: str = "auto"
    ):
        """Initialize the task executor.
        
        Args:
            workers: Number of thread workers (default: CPU count * 2)
            process_workers: Number of process workers (default: CPU count)
            mode: Execution mode - "thread", "process", or "auto"
        """
        import os
        
        # Validate and set workers
        if workers is None:
            self.workers = (os.cpu_count() or 1) * 2
        elif workers <= 0:
            self.workers = (os.cpu_count() or 1) * 2  # Use default for invalid values
        else:
            self.workers = workers
            
        if process_workers is None:
            self.process_workers = (os.cpu_count() or 1)
        elif process_workers <= 0:
            self.process_workers = (os.cpu_count() or 1)  # Use default for invalid values
        else:
            self.process_workers = process_workers
            
        self.mode = mode
        
        # Thread pool for IO-bound tasks
        self.thread_executor = concurrent.futures.ThreadPoolExecutor(
            max_workers=self.workers,
            thread_name_prefix="bgx-thread"
        )
        
        # Process pool for CPU-bound tasks
        self.process_executor = concurrent.futures.ProcessPoolExecutor(
            max_workers=self.process_workers
        )
        
        # Priority queue for pending tasks
        self._task_queue = queue.PriorityQueue()
        self._sequence_counter = 0  # For FIFO ordering within same priority
        
        # Concurrency control per tag
        self._tag_semaphores: Dict[str, threading.Semaphore] = defaultdict(
            lambda: threading.Semaphore(self.workers)
        )
        self._tag_limits: Dict[str, int] = {}
        
        # Worker thread for processing the queue
        self._worker_thread = threading.Thread(
            target=self._queue_worker,
            name="bgx-queue-worker",
            daemon=True
        )
        self._shutdown_event = threading.Event()
        self._worker_thread.start()
        
        # Track active tasks per tag
        self._active_tasks: Dict[str, Set[str]] = defaultdict(set)
        self._active_lock = threading.Lock()
    
    def set_tag_concurrency(self, tag: str, max_concurrency: int) -> None:
        """Set maximum concurrent tasks for a specific tag."""
        if max_concurrency <= 0:
            raise ValueError("max_concurrency must be positive")
        
        self._tag_limits[tag] = max_concurrency
        # Update semaphore if it already exists
        if tag in self._tag_semaphores:
            current_count = self._tag_semaphores[tag]._value
            self._tag_semaphores[tag] = threading.Semaphore(max_concurrency)
    
    def submit(self, task: Task, hint: str = None) -> None:
        """Submit a task for execution.
        
        Args:
            task: Task to execute
            hint: Execution hint - "cpu" for process pool, None for auto
        """
        if task.state != TaskState.PENDING:
            raise ValueError(f"Task {task.id} is not in PENDING state")
        
        # Store execution hint on task
        task._execution_hint = hint
        
        # Add to priority queue
        priority_value = task.priority.value
        sequence = self._next_sequence()
        self._task_queue.put((priority_value, sequence, task))
        
        # Transition to QUEUED state
        task.transition_to(TaskState.QUEUED)
    
    def _next_sequence(self) -> int:
        """Get next sequence number for FIFO ordering."""
        self._sequence_counter += 1
        return self._sequence_counter
    
    def _queue_worker(self) -> None:
        """Worker thread that processes the priority queue."""
        while not self._shutdown_event.is_set():
            try:
                # Get next task from queue with timeout
                priority, sequence, task = self._task_queue.get(timeout=0.1)
                
                # Check if task was cancelled while in queue
                if task.state == TaskState.CANCELLED:
                    self._task_queue.task_done()
                    continue
                
                # Wait for tag concurrency limit
                tag = task.tag or "default"
                semaphore = self._tag_semaphores[tag]
                
                # Acquire semaphore with timeout to allow shutdown
                if not semaphore.acquire(timeout=0.1):
                    # Put task back in queue and try again
                    self._task_queue.put((priority, sequence, task))
                    self._task_queue.task_done()
                    continue
                
                try:
                    # Submit to appropriate executor
                    self._submit_to_executor(task)
                    
                    # Track active task
                    with self._active_lock:
                        self._active_tasks[tag].add(task.id)
                        
                except Exception as e:
                    # Release semaphore on error
                    semaphore.release()
                    task.error = e
                    task.transition_to(TaskState.FAILED)
                    
                finally:
                    self._task_queue.task_done()
                    
            except queue.Empty:
                # No tasks available, continue loop
                continue
            except Exception as e:
                # Log error but continue processing
                print(f"Queue worker error: {e}")
    
    def _submit_to_executor(self, task: Task) -> None:
        """Submit task to the appropriate executor."""
        # Determine executor type
        executor = self._choose_executor(task)
        
        # Submit future
        future = executor.submit(self._execute_task, task)
        
        # Add callback to handle completion
        future.add_done_callback(lambda f: self._task_completed(task, f))
        
        # Transition to RUNNING (will be done in _execute_task)
    
    def _choose_executor(self, task: Task) -> concurrent.futures.Executor:
        """Choose appropriate executor based on mode and hint."""
        hint = getattr(task, '_execution_hint', None)
        
        if self.mode == "process":
            return self.process_executor
        elif self.mode == "thread":
            return self.thread_executor
        elif self.mode == "auto":
            if hint == "cpu":
                return self.process_executor
            else:
                return self.thread_executor
        else:
            raise ValueError(f"Unknown execution mode: {self.mode}")
    
    def _execute_task(self, task: Task) -> None:
        """Execute a task and handle its lifecycle."""
        # Transition to RUNNING now that we're actually executing
        task.transition_to(TaskState.RUNNING)
        
        try:
            # Set timeout if specified
            if task.timeout is not None:
                # Use signal-based timeout for thread execution
                # Note: This is a simplified approach, proper timeout handling
                # would require more complex mechanisms
                start_time = time.time()
                
                # Execute the function
                result = task.func(*task.args, **task.kwargs)
                
                # Check for timeout
                if time.time() - start_time > task.timeout:
                    raise TimeoutError(f"Task exceeded timeout of {task.timeout} seconds")
                
                task.result = result
                task.transition_to(TaskState.SUCCESS)
                
                # Store result in vault (handled by engine)
                
                # Call success hook
                if task.on_success:
                    try:
                        task.on_success(result)
                    except Exception as e:
                        print(f"Success hook error for task {task.id}: {e}")
                        
            else:
                # Execute without timeout
                task.result = task.func(*task.args, **task.kwargs)
                task.transition_to(TaskState.SUCCESS)
                
                # Store result in vault (handled by engine)
                
                # Call success hook
                if task.on_success:
                    try:
                        task.on_success(task.result)
                    except Exception as e:
                        print(f"Success hook error for task {task.id}: {e}")
                        
        except Exception as e:
            task.error = e
            task.transition_to(TaskState.FAILED)
            
            # Store error in vault (handled by engine)
            
            # Call failure hook
            if task.on_failure:
                try:
                    task.on_failure(e)
                except Exception as hook_error:
                    print(f"Failure hook error for task {task.id}: {hook_error}")
    
    def _task_completed(self, task: Task, future: concurrent.futures.Future) -> None:
        """Handle task completion and cleanup."""
        # Release tag semaphore
        tag = task.tag or "default"
        self._tag_semaphores[tag].release()
        
        # Remove from active tracking
        with self._active_lock:
            self._active_tasks[tag].discard(task.id)
        
        # Handle timeout if future raised TimeoutError
        if future.cancelled():
            task.transition_to(TaskState.CANCELLED)
        elif future.exception():
            # Exception already handled in _execute_task
            pass
        
        # Store result in vault via callback if available
        # This will be handled by the engine's callback system
        if hasattr(task, '_vault_store_callback'):
            try:
                task._vault_store_callback(task)
            except Exception as e:
                print(f"Vault storage error for task {task.id}: {e}")
    
    def get_queue_size(self) -> int:
        """Get number of tasks currently queued."""
        return self._task_queue.qsize()
    
    def get_active_count(self, tag: str = None) -> int:
        """Get number of currently active tasks, optionally filtered by tag."""
        with self._active_lock:
            if tag is None:
                return sum(len(tasks) for tasks in self._active_tasks.values())
            else:
                return len(self._active_tasks.get(tag or "default", set()))
    
    def get_tag_info(self) -> Dict[str, Dict[str, int]]:
        """Get information about all tags."""
        with self._active_lock:
            info = {}
            for tag, active_tasks in self._active_tasks.items():
                info[tag] = {
                    "active": len(active_tasks),
                    "limit": self._tag_limits.get(tag, self.workers),
                    "available": self._tag_semaphores[tag]._value
                }
            return info
    
    def shutdown(self, wait: bool = True) -> None:
        """Shutdown the executor and wait for completion."""
        self._shutdown_event.set()
        
        if wait:
            # Wait for queue worker to finish
            self._worker_thread.join(timeout=5.0)
            
            # Wait for all tasks to complete
            self.thread_executor.shutdown(wait=True)
            self.process_executor.shutdown(wait=True)
        else:
            # Cancel pending tasks
            while not self._task_queue.empty():
                try:
                    _, _, task = self._task_queue.get_nowait()
                    task.transition_to(TaskState.CANCELLED)
                except queue.Empty:
                    break
            
            # Shutdown executors without waiting
            self.thread_executor.shutdown(wait=False)
            self.process_executor.shutdown(wait=False)
