"""
Distributed processing module for Toxo - handles scalability and distributed computing.
"""

import asyncio
import redis
import json
import uuid
from typing import Dict, Any, List, Optional, Callable, Union
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
import multiprocessing as mp
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import pickle
import logging
from pathlib import Path

from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError


@dataclass
class Task:
    """Represents a distributed task."""
    task_id: str
    task_type: str
    data: Dict[str, Any]
    priority: int = 0
    created_at: datetime = None
    timeout: int = 300  # 5 minutes default
    retries: int = 0
    max_retries: int = 3
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()


@dataclass
class TaskResult:
    """Represents a task result."""
    task_id: str
    success: bool
    result: Any = None
    error: Optional[str] = None
    processed_at: datetime = None
    processing_time: float = 0.0
    
    def __post_init__(self):
        if self.processed_at is None:
            self.processed_at = datetime.now()


class DistributedTaskQueue:
    """
    Distributed task queue using Redis for scalability.
    """
    
    def __init__(
        self,
        redis_url: str = "redis://localhost:6379",
        queue_name: str = "toxo_tasks",
        result_ttl: int = 3600  # 1 hour
    ):
        self.redis_url = redis_url
        self.queue_name = queue_name
        self.result_ttl = result_ttl
        self.logger = get_logger(__name__)
        
        try:
            self.redis_client = redis.from_url(redis_url)
            self.redis_client.ping()
            self.logger.info(f"Connected to Redis at {redis_url}")
        except Exception as e:
            self.logger.warning(f"Redis connection failed: {e}")
            self.redis_client = None
    
    async def enqueue_task(self, task: Task) -> bool:
        """
        Enqueue a task for distributed processing.
        
        Args:
            task: Task to enqueue
            
        Returns:
            True if successful
        """
        if not self.redis_client:
            self.logger.warning("Redis not available, cannot enqueue task")
            return False
        
        try:
            task_data = asdict(task)
            task_data['created_at'] = task.created_at.isoformat()
            
            # Add to priority queue
            self.redis_client.zadd(
                self.queue_name,
                {json.dumps(task_data): task.priority}
            )
            
            self.logger.debug(f"Enqueued task {task.task_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to enqueue task: {e}")
            return False
    
    async def dequeue_task(self, timeout: int = 10) -> Optional[Task]:
        """
        Dequeue a task from the distributed queue.
        
        Args:
            timeout: Timeout in seconds
            
        Returns:
            Task if available, None otherwise
        """
        if not self.redis_client:
            return None
        
        try:
            # Get highest priority task
            result = self.redis_client.bzpopmax(self.queue_name, timeout=timeout)
            if not result:
                return None
            
            _, task_data, _ = result
            task_dict = json.loads(task_data)
            
            # Convert created_at back to datetime
            task_dict['created_at'] = datetime.fromisoformat(task_dict['created_at'])
            
            return Task(**task_dict)
            
        except Exception as e:
            self.logger.error(f"Failed to dequeue task: {e}")
            return None
    
    async def store_result(self, result: TaskResult) -> bool:
        """
        Store task result.
        
        Args:
            result: Task result to store
            
        Returns:
            True if successful
        """
        if not self.redis_client:
            return False
        
        try:
            result_data = asdict(result)
            result_data['processed_at'] = result.processed_at.isoformat()
            
            # Store result with TTL
            key = f"{self.queue_name}:result:{result.task_id}"
            
            # Serialize result data
            serialized_result = pickle.dumps(result_data)
            
            self.redis_client.setex(
                key,
                self.result_ttl,
                serialized_result
            )
            
            self.logger.debug(f"Stored result for task {result.task_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to store result: {e}")
            return False
    
    async def get_result(self, task_id: str) -> Optional[TaskResult]:
        """
        Get task result.
        
        Args:
            task_id: Task ID
            
        Returns:
            TaskResult if available
        """
        if not self.redis_client:
            return None
        
        try:
            key = f"{self.queue_name}:result:{task_id}"
            result_data = self.redis_client.get(key)
            
            if not result_data:
                return None
            
            # Deserialize result data
            result_dict = pickle.loads(result_data)
            result_dict['processed_at'] = datetime.fromisoformat(result_dict['processed_at'])
            
            return TaskResult(**result_dict)
            
        except Exception as e:
            self.logger.error(f"Failed to get result: {e}")
            return None
    
    def get_queue_size(self) -> int:
        """Get current queue size."""
        if not self.redis_client:
            return 0
        
        try:
            return self.redis_client.zcard(self.queue_name)
        except:
            return 0


class DistributedWorker:
    """
    Distributed worker that processes tasks from the queue.
    """
    
    def __init__(
        self,
        task_queue: DistributedTaskQueue,
        worker_id: str = None,
        max_workers: int = None
    ):
        self.task_queue = task_queue
        self.worker_id = worker_id or str(uuid.uuid4())
        self.max_workers = max_workers or mp.cpu_count()
        self.logger = get_logger(__name__)
        self.task_handlers: Dict[str, Callable] = {}
        self.running = False
        
        # Thread pool for I/O bound tasks
        self.thread_executor = ThreadPoolExecutor(max_workers=self.max_workers)
        
        # Process pool for CPU bound tasks
        self.process_executor = ProcessPoolExecutor(max_workers=self.max_workers)
    
    def register_handler(self, task_type: str, handler: Callable):
        """
        Register a task handler.
        
        Args:
            task_type: Type of task
            handler: Handler function
        """
        self.task_handlers[task_type] = handler
        self.logger.info(f"Registered handler for task type: {task_type}")
    
    async def start(self):
        """Start the worker."""
        self.running = True
        self.logger.info(f"Starting worker {self.worker_id}")
        
        while self.running:
            try:
                # Get task from queue
                task = await self.task_queue.dequeue_task()
                if not task:
                    continue
                
                # Process task
                await self._process_task(task)
                
            except Exception as e:
                self.logger.error(f"Worker error: {e}")
                await asyncio.sleep(1)
    
    async def stop(self):
        """Stop the worker."""
        self.running = False
        self.thread_executor.shutdown(wait=True)
        self.process_executor.shutdown(wait=True)
        self.logger.info(f"Stopped worker {self.worker_id}")
    
    async def _process_task(self, task: Task):
        """
        Process a single task.
        
        Args:
            task: Task to process
        """
        start_time = datetime.now()
        
        try:
            # Check if we have a handler for this task type
            if task.task_type not in self.task_handlers:
                raise ToxoError(f"No handler for task type: {task.task_type}")
            
            handler = self.task_handlers[task.task_type]
            
            # Execute task
            if asyncio.iscoroutinefunction(handler):
                result = await handler(task.data)
            else:
                # Run in thread pool
                result = await asyncio.get_event_loop().run_in_executor(
                    self.thread_executor,
                    handler,
                    task.data
                )
            
            # Store successful result
            processing_time = (datetime.now() - start_time).total_seconds()
            task_result = TaskResult(
                task_id=task.task_id,
                success=True,
                result=result,
                processing_time=processing_time
            )
            
            await self.task_queue.store_result(task_result)
            self.logger.debug(f"Processed task {task.task_id} successfully")
            
        except Exception as e:
            # Store error result
            processing_time = (datetime.now() - start_time).total_seconds()
            task_result = TaskResult(
                task_id=task.task_id,
                success=False,
                error=str(e),
                processing_time=processing_time
            )
            
            await self.task_queue.store_result(task_result)
            self.logger.error(f"Failed to process task {task.task_id}: {e}")
            
            # Retry if possible
            if task.retries < task.max_retries:
                task.retries += 1
                await self.task_queue.enqueue_task(task)
                self.logger.info(f"Retrying task {task.task_id} (attempt {task.retries})")


class DistributedProcessingManager:
    """
    Manages distributed processing for Toxo.
    """
    
    def __init__(
        self,
        redis_url: str = "redis://localhost:6379",
        num_workers: int = None
    ):
        self.redis_url = redis_url
        self.num_workers = num_workers or mp.cpu_count()
        self.logger = get_logger(__name__)
        
        # Initialize components
        self.task_queue = DistributedTaskQueue(redis_url)
        self.workers: List[DistributedWorker] = []
        self.running = False
        
        # Register default task handlers
        self._register_default_handlers()
    
    def _register_default_handlers(self):
        """Register default task handlers."""
        # These would be implemented based on Toxo's needs
        pass
    
    async def submit_task(
        self,
        task_type: str,
        data: Dict[str, Any],
        priority: int = 0,
        timeout: int = 300
    ) -> str:
        """
        Submit a task for distributed processing.
        
        Args:
            task_type: Type of task
            data: Task data
            priority: Task priority (higher = more important)
            timeout: Task timeout in seconds
            
        Returns:
            Task ID
        """
        task_id = str(uuid.uuid4())
        task = Task(
            task_id=task_id,
            task_type=task_type,
            data=data,
            priority=priority,
            timeout=timeout
        )
        
        success = await self.task_queue.enqueue_task(task)
        if not success:
            raise ToxoError("Failed to submit task")
        
        return task_id
    
    async def get_result(self, task_id: str, timeout: int = 30) -> Optional[TaskResult]:
        """
        Get task result, waiting if necessary.
        
        Args:
            task_id: Task ID
            timeout: Timeout in seconds
            
        Returns:
            TaskResult if available
        """
        start_time = datetime.now()
        
        while (datetime.now() - start_time).total_seconds() < timeout:
            result = await self.task_queue.get_result(task_id)
            if result:
                return result
            
            await asyncio.sleep(0.1)
        
        return None
    
    async def start_workers(self):
        """Start distributed workers."""
        self.running = True
        
        for i in range(self.num_workers):
            worker = DistributedWorker(
                self.task_queue,
                worker_id=f"worker-{i}",
                max_workers=2  # Limit threads per worker
            )
            
            # Register handlers
            for task_type, handler in self._get_task_handlers().items():
                worker.register_handler(task_type, handler)
            
            self.workers.append(worker)
        
        # Start all workers
        tasks = [worker.start() for worker in self.workers]
        await asyncio.gather(*tasks)
    
    async def stop_workers(self):
        """Stop all workers."""
        self.running = False
        
        for worker in self.workers:
            await worker.stop()
        
        self.workers.clear()
    
    def _get_task_handlers(self) -> Dict[str, Callable]:
        """Get task handlers for workers."""
        return {
            "generate_response": self._handle_generate_response,
            "process_document": self._handle_process_document,
            "train_model": self._handle_train_model,
            "rerank_responses": self._handle_rerank_responses,
        }
    
    async def _handle_generate_response(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle response generation task."""
        # Implementation would integrate with ToxoLayer
        return {"response": "Generated response", "confidence": 0.9}
    
    async def _handle_process_document(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle document processing task."""
        # Implementation would integrate with DocumentProcessor
        return {"processed": True, "chunks": []}
    
    async def _handle_train_model(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle model training task."""
        # Implementation would integrate with ToxoTrainer
        return {"trained": True, "metrics": {}}
    
    async def _handle_rerank_responses(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Handle response reranking task."""
        # Implementation would integrate with ResponseReranker
        return {"ranked_responses": [], "best_response": ""}
    
    def get_stats(self) -> Dict[str, Any]:
        """Get processing statistics."""
        return {
            "queue_size": self.task_queue.get_queue_size(),
            "num_workers": len(self.workers),
            "running": self.running,
            "redis_connected": self.task_queue.redis_client is not None
        }


# Global distributed processing manager
_distributed_manager = None

def get_distributed_manager() -> DistributedProcessingManager:
    """Get global distributed processing manager."""
    return DistributedProcessingManager()


@dataclass 
class ClusterNode:
    """Represents a node in a distributed cluster."""
    node_id: str
    host: str
    port: int
    status: str = "active"
    capabilities: List[str] = None
    load: float = 0.0
    last_heartbeat: datetime = None
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.capabilities is None:
            self.capabilities = []
        if self.metadata is None:
            self.metadata = {}
        if self.last_heartbeat is None:
            self.last_heartbeat = datetime.now()


class DistributedManager:
    """
    Main distributed computing manager for Toxo.
    
    Coordinates distributed tasks, manages cluster nodes, and handles
    load balancing across the distributed system.
    """
    
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self.redis_url = redis_url
        self.logger = get_logger(__name__)
        self.processing_manager = DistributedProcessingManager(redis_url)
        self.cluster_nodes: Dict[str, ClusterNode] = {}
        self.task_queue = DistributedTaskQueue(redis_url)
        
    def add_node(self, node: ClusterNode):
        """Add a node to the cluster."""
        self.cluster_nodes[node.node_id] = node
        self.logger.info(f"Added cluster node {node.node_id}")
        
    def remove_node(self, node_id: str):
        """Remove a node from the cluster."""
        if node_id in self.cluster_nodes:
            del self.cluster_nodes[node_id]
            self.logger.info(f"Removed cluster node {node_id}")
    
    def get_active_nodes(self) -> List[ClusterNode]:
        """Get list of active cluster nodes."""
        return [node for node in self.cluster_nodes.values() if node.status == "active"]
    
    async def submit_distributed_task(self, task_type: str, data: Dict[str, Any], **kwargs) -> str:
        """Submit a task for distributed processing."""
        return await self.processing_manager.submit_task(task_type, data, **kwargs)
    
    async def get_task_result(self, task_id: str, timeout: int = 30) -> Optional[TaskResult]:
        """Get result of a distributed task."""
        return await self.processing_manager.get_result(task_id, timeout)
    
    async def start_cluster(self):
        """Start the distributed cluster."""
        await self.processing_manager.start_workers()
        self.logger.info("Distributed cluster started")
    
    async def stop_cluster(self):
        """Stop the distributed cluster."""
        await self.processing_manager.stop_workers()
        self.logger.info("Distributed cluster stopped")
    
    def get_cluster_stats(self) -> Dict[str, Any]:
        """Get cluster statistics."""
        return {
            "total_nodes": len(self.cluster_nodes),
            "active_nodes": len(self.get_active_nodes()),
            "processing_stats": self.processing_manager.get_stats(),
            "queue_size": self.task_queue.get_queue_size()
        }


class TaskDistributor:
    """
    Distributes tasks across cluster nodes based on load and capabilities.
    """
    
    def __init__(self, distributed_manager: DistributedManager):
        self.distributed_manager = distributed_manager
        self.logger = get_logger(__name__)
    
    def select_optimal_node(self, task_requirements: Dict[str, Any] = None) -> Optional[ClusterNode]:
        """Select the optimal node for a task based on load and capabilities."""
        active_nodes = self.distributed_manager.get_active_nodes()
        
        if not active_nodes:
            return None
        
        # Simple load-based selection for now
        return min(active_nodes, key=lambda node: node.load)
    
    async def distribute_task(self, task_type: str, data: Dict[str, Any], **kwargs) -> str:
        """Distribute a task to the optimal node."""
        optimal_node = self.select_optimal_node()
        
        if not optimal_node:
            self.logger.warning("No active nodes available for task distribution")
            return await self.distributed_manager.submit_distributed_task(task_type, data, **kwargs)
        
        # For now, just submit to the distributed manager
        # In a full implementation, we would route to specific nodes
        self.logger.info(f"Distributing task to node {optimal_node.node_id}")
        return await self.distributed_manager.submit_distributed_task(task_type, data, **kwargs) 