"""
BGX Task module - Task dataclass and lifecycle state machine.

This module defines the Task dataclass and TaskState enum that represent
the core unit of work in BGX with proper state transitions.
"""

import enum
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Optional, Union


class TaskState(enum.Enum):
    """Task lifecycle states."""
    PENDING = "pending"
    QUEUED = "queued"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    RETRYING = "retrying"
    DEAD = "dead"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"


class Priority(enum.Enum):
    """Task priority levels."""
    HIGH = 0
    NORMAL = 1
    LOW = 2


@dataclass(frozen=True)
class RetryPolicy:
    """Configuration for task retry behavior."""
    max_attempts: int = 3
    strategy: str = "exponential"  # fixed, linear, exponential, fibonacci
    base_delay: float = 1.0
    max_delay: float = 60.0
    jitter: bool = True
    retry_on: list[type] = field(default_factory=lambda: [Exception])
    abort_on: list[type] = field(default_factory=list)


@dataclass
class Task:
    """Represents a unit of work in BGX."""
    
    # Core task identification and function
    func: Callable
    args: tuple = field(default_factory=tuple)
    kwargs: dict = field(default_factory=dict)
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    
    # State and lifecycle
    state: TaskState = TaskState.PENDING
    priority: Priority = Priority.NORMAL
    tag: Optional[str] = None
    retry_policy: Optional[RetryPolicy] = None
    timeout: Optional[float] = None
    
    # Timestamps
    created_at: float = field(default_factory=time.time)
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    
    # Execution tracking
    attempt: int = 1
    result: Any = None
    error: Optional[Exception] = None
    
    # Hooks
    on_success: Optional[Callable] = None
    on_failure: Optional[Callable] = None
    on_timeout: Optional[Callable] = None
    on_retry: Optional[Callable] = None
    
    def __post_init__(self):
        """Validate task configuration after creation."""
        if not callable(self.func):
            raise TypeError("func must be callable")
        
        if not isinstance(self.priority, Priority):
            raise TypeError("priority must be a Priority enum value")
        
        if self.timeout is not None and self.timeout <= 0:
            raise ValueError("timeout must be positive")
        
        if self.retry_policy is not None:
            if self.retry_policy.max_attempts < 1:
                raise ValueError("max_attempts must be at least 1")
            if self.retry_policy.base_delay <= 0:
                raise ValueError("base_delay must be positive")
            if self.retry_policy.max_delay <= 0:
                raise ValueError("max_delay must be positive")
    
    def transition_to(self, new_state: TaskState) -> None:
        """Transition task to a new state with validation."""
        valid_transitions = {
            TaskState.PENDING: {TaskState.QUEUED, TaskState.CANCELLED},
            TaskState.QUEUED: {TaskState.RUNNING, TaskState.CANCELLED, TaskState.PENDING},
            TaskState.RUNNING: {TaskState.SUCCESS, TaskState.FAILED, TaskState.TIMEOUT, TaskState.CANCELLED},
            TaskState.FAILED: {TaskState.RETRYING, TaskState.DEAD},
            TaskState.RETRYING: {TaskState.QUEUED, TaskState.PENDING},
            TaskState.DEAD: set(),
            TaskState.SUCCESS: set(),
            TaskState.TIMEOUT: set(),
            TaskState.CANCELLED: set(),
        }
        
        if new_state not in valid_transitions.get(self.state, set()):
            raise ValueError(f"Invalid state transition from {self.state} to {new_state}")
        
        old_state = self.state
        self.state = new_state
        
        # Update timestamps
        if new_state == TaskState.RUNNING:
            self.started_at = time.time()
        elif new_state in {TaskState.SUCCESS, TaskState.FAILED, TaskState.DEAD, TaskState.TIMEOUT, TaskState.CANCELLED}:
            self.finished_at = time.time()
    
    def can_retry(self) -> bool:
        """Check if task can be retried based on policy and attempt count."""
        if self.retry_policy is None:
            return False
        
        if self.attempt >= self.retry_policy.max_attempts:
            return False
        
        if self.error is None:
            return False
        
        # Check if error type is in retry_on list
        for retry_type in self.retry_policy.retry_on:
            if isinstance(self.error, retry_type):
                # Check if error type is in abort_on list
                for abort_type in self.retry_policy.abort_on:
                    if isinstance(self.error, abort_type):
                        return False
                return True
        
        return False
    
    def compute_retry_delay(self) -> float:
        """Compute delay before next retry based on policy."""
        if self.retry_policy is None:
            return 0.0
        
        delay = 0.0
        base = self.retry_policy.base_delay
        
        if self.retry_policy.strategy == "fixed":
            delay = base
        elif self.retry_policy.strategy == "linear":
            delay = base * self.attempt
        elif self.retry_policy.strategy == "exponential":
            delay = base * (2 ** (self.attempt - 1))
        elif self.retry_policy.strategy == "fibonacci":
            delay = base * self._fibonacci(self.attempt)
        else:
            raise ValueError(f"Unknown retry strategy: {self.retry_policy.strategy}")
        
        # Apply max delay cap
        delay = min(delay, self.retry_policy.max_delay)
        
        # Add jitter if enabled
        if self.retry_policy.jitter:
            import random
            jitter_amount = delay * 0.25
            delay += random.uniform(0, jitter_amount)
        
        return delay
    
    def _fibonacci(self, n: int) -> int:
        """Compute nth Fibonacci number."""
        if n <= 0:
            return 0
        elif n == 1:
            return 1
        elif n == 2:
            return 1
        else:
            a, b = 1, 1
            for _ in range(3, n + 1):
                a, b = b, a + b
            return b
    
    def wait(self, timeout: Optional[float] = None) -> Any:
        """Wait for task completion and return result."""
        import threading
        import time
        
        start_time = time.time()
        check_interval = 0.01  # 10ms
        
        while self.state not in {TaskState.SUCCESS, TaskState.FAILED, TaskState.DEAD, TaskState.TIMEOUT, TaskState.CANCELLED}:
            if timeout is not None and (time.time() - start_time) > timeout:
                raise TimeoutError(f"Task {self.id} did not complete within {timeout} seconds")
            time.sleep(check_interval)
        
        if self.state == TaskState.SUCCESS:
            return self.result
        elif self.error is not None:
            raise self.error
        elif self.state == TaskState.TIMEOUT:
            raise TimeoutError(f"Task {self.id} timed out")
        elif self.state == TaskState.CANCELLED:
            raise RuntimeError(f"Task {self.id} was cancelled")
        else:
            raise RuntimeError(f"Task {self.id} failed without error")
    
    def __repr__(self) -> str:
        """String representation of task."""
        return f"Task(id={self.id[:8]}..., state={self.state.value}, func={getattr(self.func, '__name__', 'unknown')})"
