"""
BGX Engine module - Central orchestrator BackgroundEngine.

This module provides the BackgroundEngine class that coordinates all
BGX components and exposes the public API.
"""

import threading
import time
from typing import Any, Callable, Dict, List, Optional, Union

from .executor import TaskExecutor
from .vault import ResultVault
from .retry import RetryManager
from .scheduler import TaskScheduler
from .task import Task, TaskState, Priority, RetryPolicy
from .group import TaskGroup, GroupPattern


class BackgroundEngine:
    """Central orchestrator for BGX background task execution."""
    
    def __init__(
        self,
        workers: int = None,
        process_workers: int = None,
        mode: str = "auto",
        vault_ttl: float = 300.0,
        vault_max_size: int = 1000
    ):
        """Initialize the background engine.
        
        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"
            vault_ttl: Result TTL in seconds for the vault
            vault_max_size: Maximum number of results in the vault
        """
        # Core components
        self.executor = TaskExecutor(
            workers=workers,
            process_workers=process_workers,
            mode=mode
        )
        self.vault = ResultVault(ttl=vault_ttl, max_size=vault_max_size)
        self.retry_manager = RetryManager()
        self.scheduler = TaskScheduler(submit_callback=self._submit_scheduled_task)
        
        # Task tracking
        self._tasks: Dict[str, Task] = {}
        self._tasks_lock = threading.RLock()
        
        # Setup retry callback
        self.retry_manager.set_submit_callback(self._submit_retry)
        
        # Engine state
        self._shutdown = False
        self._shutdown_lock = threading.Lock()
    
    def fire(
        self,
        func: Callable,
        *args,
        priority: Priority = Priority.NORMAL,
        tag: str = None,
        timeout: float = None,
        retry: RetryPolicy = None,
        on_success: Callable = None,
        on_failure: Callable = None,
        on_timeout: Callable = None,
        on_retry: Callable = None,
        **kwargs
    ) -> None:
        """Submit a fire-and-forget task.
        
        Args:
            func: Function to execute
            *args: Positional arguments for the function
            priority: Task priority
            tag: Task tag for grouping and circuit breaking
            timeout: Maximum execution time in seconds
            retry: Retry policy configuration
            on_success: Hook called on successful completion
            on_failure: Hook called on failure
            on_timeout: Hook called on timeout
            on_retry: Hook called before each retry
            **kwargs: Keyword arguments for the function
        """
        self._check_shutdown()
        
        if not callable(func):
            raise TypeError("func must be callable")
        
        task = Task(
            func=func,
            args=args,
            kwargs=kwargs,
            priority=priority,
            tag=tag,
            timeout=timeout,
            retry_policy=retry,
            on_success=on_success,
            on_failure=on_failure,
            on_timeout=on_timeout,
            on_retry=on_retry
        )
        
        self._submit_task(task)
    
    def submit(
        self,
        func: Callable,
        *args,
        priority: Priority = Priority.NORMAL,
        tag: str = None,
        timeout: float = None,
        retry: RetryPolicy = None,
        on_success: Callable = None,
        on_failure: Callable = None,
        on_timeout: Callable = None,
        on_retry: Callable = None,
        **kwargs
    ) -> Task:
        """Submit a task with tracking and return Task object.
        
        Args:
            func: Function to execute
            *args: Positional arguments for the function
            priority: Task priority
            tag: Task tag for grouping and circuit breaking
            timeout: Maximum execution time in seconds
            retry: Retry policy configuration
            on_success: Hook called on successful completion
            on_failure: Hook called on failure
            on_timeout: Hook called on timeout
            on_retry: Hook called before each retry
            **kwargs: Keyword arguments for the function
            
        Returns:
            Task object for tracking
        """
        self._check_shutdown()
        
        if not callable(func):
            raise TypeError("func must be callable")
        
        task = Task(
            func=func,
            args=args,
            kwargs=kwargs,
            priority=priority,
            tag=tag,
            timeout=timeout,
            retry_policy=retry,
            on_success=on_success,
            on_failure=on_failure,
            on_timeout=on_timeout,
            on_retry=on_retry
        )
        
        self._submit_task(task)
        return task
    
    def _submit_task(self, task: Task) -> None:
        """Internal method to submit a task to the executor."""
        with self._tasks_lock:
            self._tasks[task.id] = task
        
        # Set up vault storage callback
        task._vault_store_callback = self._store_task_result
        
        # Submit to executor
        self.executor.submit(task)
    
    def _store_task_result(self, task: Task) -> None:
        """Store task result or error in the vault."""
        if task.state == TaskState.SUCCESS:
            self.vault.store(task.id, result=task.result)
        elif task.state in {TaskState.FAILED, TaskState.DEAD, TaskState.TIMEOUT}:
            self.vault.store(task.id, error=task.error)
    
    def _submit_scheduled_task(self, func: Callable, *args, **kwargs) -> None:
        """Submit a task from the scheduler."""
        self.fire(func, *args, **kwargs)
    
    def schedule(self, func: Callable, delay: float, *args, **kwargs):
        """Schedule a one-shot task to run after delay.
        
        Args:
            func: Function to execute
            delay: Delay in seconds before execution
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function
            
        Returns:
            ScheduledHandle for managing the scheduled task
        """
        self._check_shutdown()
        return self.scheduler.schedule(func, delay, *args, **kwargs)
    
    def every(self, func: Callable, interval: float, *args, **kwargs):
        """Schedule a recurring task to run at fixed intervals.
        
        Args:
            func: Function to execute
            interval: Interval in seconds between executions
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function
            
        Returns:
            ScheduledHandle for managing the scheduled task
        """
        self._check_shutdown()
        return self.scheduler.every(func, interval, *args, **kwargs)
    
    def cron(self, func: Callable, expression: str, *args, **kwargs):
        """Schedule a recurring task using cron expression.
        
        Args:
            func: Function to execute
            expression: Cron expression (5 fields: minute hour day month weekday)
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function
            
        Returns:
            ScheduledHandle for managing the scheduled task
        """
        self._check_shutdown()
        return self.scheduler.cron(func, expression, *args, **kwargs)
    
    def create_group(self) -> TaskGroup:
        """Create a new task group.
        
        Returns:
            TaskGroup instance for managing task dependencies
        """
        self._check_shutdown()
        return TaskGroup(self)
    
    def get_group(self, group_id: str) -> Optional[TaskGroup]:
        """Get a task group by ID (placeholder for future implementation)."""
        # This is a placeholder - in a real implementation, you'd store
        # and retrieve taskGroup instances
        return None
    
    def _submit_retry(self, task_id: str) -> None:
        """Submit a task for retry."""
        with self._tasks_lock:
            if task_id not in self._tasks:
                return
            
            task = self._tasks[task_id]
            
            # Reset task state for retry
            task.transition_to(TaskState.PENDING)
            
            # Submit to executor
            self.executor.submit(task)
    
    def get_result(self, task_id: str, timeout: Optional[float] = None) -> Any:
        """Get the result of a task.
        
        Args:
            task_id: Unique task identifier
            timeout: Maximum time to wait for result
            
        Returns:
            Task result
            
        Raises:
            KeyError: If task_id is not found
            TimeoutError: If timeout is reached
            Exception: Re-raises the original task error
        """
        # Try to get from vault first
        try:
            return self.vault.wait_for(task_id, timeout=timeout)
        except KeyError:
            # Task not in vault, check if it exists
            with self._tasks_lock:
                if task_id not in self._tasks:
                    raise KeyError(f"Task {task_id} not found")
            
            # Task exists but result not ready, wait for it
            return self.vault.wait_for(task_id, timeout=timeout)
    
    def get_task(self, task_id: str) -> Optional[Task]:
        """Get a task by ID.
        
        Args:
            task_id: Unique task identifier
            
        Returns:
            Task object or None if not found
        """
        with self._tasks_lock:
            return self._tasks.get(task_id)
    
    def list_tasks(self, state: TaskState = None, tag: str = None) -> List[Task]:
        """List tasks with optional filtering.
        
        Args:
            state: Filter by task state
            tag: Filter by task tag
            
        Returns:
            List of matching tasks
        """
        with self._tasks_lock:
            tasks = list(self._tasks.values())
        
        # Apply filters
        if state is not None:
            tasks = [t for t in tasks if t.state == state]
        
        if tag is not None:
            tasks = [t for t in tasks if t.tag == tag]
        
        return tasks
    
    def cancel(self, task_id: str) -> bool:
        """Cancel a task.
        
        Args:
            task_id: Unique task identifier
            
        Returns:
            True if task was cancelled, False if not found or already terminal
        """
        with self._tasks_lock:
            if task_id not in self._tasks:
                return False
            
            task = self._tasks[task_id]
        
        # Can only cancel pending or queued tasks
        if task.state in {TaskState.PENDING, TaskState.QUEUED}:
            task.transition_to(TaskState.CANCELLED)
            
            # Cancel retry if scheduled
            self.retry_manager.cancel_retry(task_id)
            
            return True
        
        return False
    
    def stats(self) -> Dict[str, Any]:
        """Get engine statistics.
        
        Returns:
            Dictionary containing engine statistics
        """
        with self._tasks_lock:
            tasks = list(self._tasks.values())
        
        # Count tasks by state
        state_counts = {}
        for state in TaskState:
            state_counts[state.value] = sum(1 for t in tasks if t.state == state)
        
        return {
            "total_tasks": len(tasks),
            "active": state_counts.get("running", 0),
            "queued": self.executor.get_queue_size(),
            "success": state_counts.get("success", 0),
            "failed": state_counts.get("failed", 0),
            "dead": state_counts.get("dead", 0),
            "retrying": self.retry_manager.get_retry_count(),
            "vault_stats": self.vault.get_stats(),
            "tag_info": self.executor.get_tag_info()
        }
    
    def set_tag_concurrency(self, tag: str, max_concurrency: int) -> None:
        """Set maximum concurrent tasks for a specific tag.
        
        Args:
            tag: Task tag
            max_concurrency: Maximum concurrent tasks
        """
        self.executor.set_tag_concurrency(tag, max_concurrency)
    
    def _check_shutdown(self) -> None:
        """Check if engine is shutdown and raise error if so."""
        with self._shutdown_lock:
            if self._shutdown:
                raise RuntimeError("Engine is shutdown")
    
    def shutdown(self, wait: bool = True) -> None:
        """Shutdown the engine.
        
        Args:
            wait: Whether to wait for tasks to complete
        """
        with self._shutdown_lock:
            if self._shutdown:
                return
            
            self._shutdown = True
        
        # Shutdown components
        self.scheduler.shutdown()
        self.executor.shutdown(wait=wait)
        self.retry_manager.shutdown()
        self.vault.shutdown()
        
        # Clear task tracking
        with self._tasks_lock:
            self._tasks.clear()
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.shutdown()


# Convenience functions for creating engines
def create_engine(
    workers: int = None,
    process_workers: int = None,
    mode: str = "auto",
    vault_ttl: float = 300.0,
    vault_max_size: int = 1000
) -> BackgroundEngine:
    """Create a new BackgroundEngine instance.
    
    Args:
        workers: Number of thread workers
        process_workers: Number of process workers
        mode: Execution mode
        vault_ttl: Result TTL in seconds
        vault_max_size: Maximum vault size
        
    Returns:
        BackgroundEngine instance
    """
    return BackgroundEngine(
        workers=workers,
        process_workers=process_workers,
        mode=mode,
        vault_ttl=vault_ttl,
        vault_max_size=vault_max_size
    )


def io_engine(workers: int = None, **kwargs) -> BackgroundEngine:
    """Create an engine optimized for IO-bound tasks.
    
    Args:
        workers: Number of thread workers
        **kwargs: Additional engine arguments
        
    Returns:
        BackgroundEngine configured for IO-bound work
    """
    import os
    if workers is None:
        workers = (os.cpu_count() or 1) * 4  # More workers for IO-bound
    
    return BackgroundEngine(
        workers=workers,
        mode="thread",
        **kwargs
    )


def cpu_engine(process_workers: int = None, **kwargs) -> BackgroundEngine:
    """Create an engine optimized for CPU-bound tasks.
    
    Args:
        process_workers: Number of process workers
        **kwargs: Additional engine arguments
        
    Returns:
        BackgroundEngine configured for CPU-bound work
    """
    import os
    if process_workers is None:
        process_workers = os.cpu_count() or 1
    
    return BackgroundEngine(
        process_workers=process_workers,
        mode="process",
        **kwargs
    )
