"""
Campfire class for orchestrating multimodal LLM workflows.
"""

import asyncio
import logging
from typing import List, Dict, Any, Optional, Callable, Union
from datetime import datetime, timedelta

from .torch import Torch
from .camper import Camper
from ..party_box.box_driver import BoxDriver
from ..mcp.protocol import MCPProtocol


logger = logging.getLogger(__name__)


class Campfire:
    """
    A Campfire orchestrates a group of Campers (models/tools) to process
    torches (data) and pass results via MCP to other campfires.
    
    Each campfire has a specific purpose and contains one or more campers
    that work together to process incoming torches and generate new ones.
    """
    
    def __init__(
        self,
        name: str,
        campers: List[Camper],
        party_box: BoxDriver,
        mcp_protocol: Optional[MCPProtocol] = None,
        config: Dict[str, Any] = None
    ):
        """
        Initialize a Campfire.
        
        Args:
            name: Unique name for this campfire
            campers: List of campers in this campfire
            party_box: Storage driver for assets
            mcp_protocol: MCP protocol for communication
            config: Additional configuration
        """
        self.name = name
        self.campers = campers
        self.party_box = party_box
        self.mcp_protocol = mcp_protocol
        self.config = config or {}
        
        # State management
        self.is_running = False
        self.processed_torches: Dict[str, Torch] = {}
        self.torch_queue: asyncio.Queue = asyncio.Queue()
        
        # Configuration
        self.max_concurrent_tasks = self.config.get('max_concurrent_tasks', 3)
        self.torch_ttl = self.config.get('torch_ttl_hours', 24)
        self.auto_cleanup = self.config.get('auto_cleanup', True)
        
        # Callbacks
        self.on_torch_processed: Optional[Callable[[Torch], None]] = None
        self.on_error: Optional[Callable[[Exception, Torch], None]] = None
        
        # Setup campers
        for camper in self.campers:
            camper.set_party_box(self.party_box)
            camper.set_campfire_name(self.name)
    
    async def start(self) -> None:
        """Start the campfire processing loop."""
        if self.is_running:
            logger.warning(f"Campfire {self.name} is already running")
            return
        
        self.is_running = True
        logger.info(f"Starting campfire: {self.name}")
        
        # Start processing tasks
        tasks = []
        for i in range(self.max_concurrent_tasks):
            task = asyncio.create_task(self._processing_loop())
            tasks.append(task)
        
        # Start cleanup task if enabled
        if self.auto_cleanup:
            cleanup_task = asyncio.create_task(self._cleanup_loop())
            tasks.append(cleanup_task)
        
        # Subscribe to MCP channels if protocol is available
        if self.mcp_protocol:
            await self._setup_mcp_subscriptions()
        
        # Wait for all tasks to complete
        try:
            await asyncio.gather(*tasks)
        except asyncio.CancelledError:
            logger.info(f"Campfire {self.name} processing cancelled")
        finally:
            self.is_running = False
    
    async def stop(self) -> None:
        """Stop the campfire processing."""
        logger.info(f"Stopping campfire: {self.name}")
        self.is_running = False
        
        # Cancel any pending tasks
        for task in asyncio.all_tasks():
            if not task.done() and task.get_name().startswith(f"campfire_{self.name}"):
                task.cancel()
    
    async def process_torch(self, torch: Torch) -> List[Torch]:
        """
        Process a single torch through all campers.
        
        Args:
            torch: Input torch to process
            
        Returns:
            List of output torches generated by campers
        """
        logger.info(f"Processing torch {torch.torch_id} in campfire {self.name}")
        
        try:
            output_torches = []
            
            # Process through each camper
            for camper in self.campers:
                logger.debug(f"Processing torch with camper: {camper.__class__.__name__}")
                
                try:
                    # Process the torch
                    result_torches = await camper.process(torch)
                    
                    if result_torches:
                        if isinstance(result_torches, Torch):
                            result_torches = [result_torches]
                        
                        for result_torch in result_torches:
                            # Set source campfire
                            result_torch.source_campfire = self.name
                            result_torch.metadata['processed_by'] = camper.__class__.__name__
                            result_torch.metadata['parent_torch_id'] = torch.torch_id
                            
                            output_torches.append(result_torch)
                            
                            # Send via MCP if available
                            if self.mcp_protocol:
                                await self._send_torch_via_mcp(result_torch)
                
                except Exception as e:
                    logger.error(f"Error processing torch with {camper.__class__.__name__}: {e}")
                    if self.on_error:
                        self.on_error(e, torch)
                    continue
            
            # Store processed torch
            self.processed_torches[torch.torch_id] = torch
            
            # Call callback if set
            if self.on_torch_processed:
                self.on_torch_processed(torch)
            
            logger.info(f"Generated {len(output_torches)} output torches from {torch.torch_id}")
            return output_torches
            
        except Exception as e:
            logger.error(f"Error processing torch {torch.torch_id}: {e}")
            if self.on_error:
                self.on_error(e, torch)
            return []
    
    async def add_torch(self, torch: Torch) -> None:
        """
        Add a torch to the processing queue.
        
        Args:
            torch: Torch to add to queue
        """
        await self.torch_queue.put(torch)
        logger.debug(f"Added torch {torch.torch_id} to queue")
    
    async def add_torch_from_data(
        self,
        claim: str,
        path: Optional[str] = None,
        confidence: float = 1.0,
        metadata: Dict[str, Any] = None,
        channel: str = "default"
    ) -> Torch:
        """
        Create and add a torch from raw data.
        
        Args:
            claim: Main claim or content
            path: Optional path to asset
            confidence: Confidence score (0.0 to 1.0)
            metadata: Additional metadata
            channel: MCP channel for the torch
            
        Returns:
            Created torch
        """
        torch = Torch(
            claim=claim,
            path=path,
            confidence=confidence,
            metadata=metadata or {},
            channel=channel,
            source_campfire=self.name
        )
        
        await self.add_torch(torch)
        return torch
    
    async def _processing_loop(self) -> None:
        """Main processing loop for handling torches."""
        while self.is_running:
            try:
                # Wait for a torch with timeout
                torch = await asyncio.wait_for(
                    self.torch_queue.get(),
                    timeout=1.0
                )
                
                # Process the torch
                await self.process_torch(torch)
                
                # Mark task as done
                self.torch_queue.task_done()
                
            except asyncio.TimeoutError:
                # No torch available, continue loop
                continue
            except Exception as e:
                logger.error(f"Error in processing loop: {e}")
                await asyncio.sleep(1)
    
    async def _cleanup_loop(self) -> None:
        """Cleanup loop for removing old torches."""
        while self.is_running:
            try:
                await asyncio.sleep(3600)  # Run every hour
                await self._cleanup_old_torches()
            except Exception as e:
                logger.error(f"Error in cleanup loop: {e}")
    
    async def _cleanup_old_torches(self) -> None:
        """Remove torches older than TTL."""
        cutoff_time = datetime.utcnow() - timedelta(hours=self.torch_ttl)
        
        to_remove = []
        for torch_id, torch in self.processed_torches.items():
            if torch.timestamp < cutoff_time:
                to_remove.append(torch_id)
        
        for torch_id in to_remove:
            del self.processed_torches[torch_id]
        
        if to_remove:
            logger.info(f"Cleaned up {len(to_remove)} old torches")
    
    async def _setup_mcp_subscriptions(self) -> None:
        """Setup MCP channel subscriptions."""
        if not self.mcp_protocol:
            return
        
        # Subscribe to channels this campfire should listen to
        channels = self.config.get('subscribe_channels', [])
        for channel in channels:
            await self.mcp_protocol.subscribe(channel, self._handle_mcp_torch)
    
    async def _handle_mcp_torch(self, message: Dict[str, Any]) -> None:
        """
        Handle incoming torch from MCP.
        
        Args:
            message: MCP message containing torch data
        """
        try:
            torch = Torch.from_mcp_message(message)
            await self.add_torch(torch)
        except Exception as e:
            logger.error(f"Error handling MCP torch: {e}")
    
    async def _send_torch_via_mcp(self, torch: Torch) -> None:
        """
        Send a torch via MCP.
        
        Args:
            torch: Torch to send
        """
        if not self.mcp_protocol:
            return
        
        try:
            message = torch.to_mcp_message()
            await self.mcp_protocol.send_message(torch.channel, message)
        except Exception as e:
            logger.error(f"Error sending torch via MCP: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get campfire statistics.
        
        Returns:
            Dictionary with campfire stats
        """
        return {
            'name': self.name,
            'is_running': self.is_running,
            'campers_count': len(self.campers),
            'processed_torches_count': len(self.processed_torches),
            'queue_size': self.torch_queue.qsize(),
            'campers': [camper.__class__.__name__ for camper in self.campers]
        }
    
    def __str__(self) -> str:
        """String representation of the campfire."""
        return f"Campfire({self.name}, {len(self.campers)} campers)"
    
    def __repr__(self) -> str:
        """Detailed string representation."""
        return (f"Campfire(name='{self.name}', "
                f"campers={[c.__class__.__name__ for c in self.campers]}, "
                f"running={self.is_running})")


class CampfireManager:
    """
    Manager for multiple campfires, handling coordination and routing.
    """
    
    def __init__(self, mcp_protocol: Optional[MCPProtocol] = None):
        """
        Initialize the campfire manager.
        
        Args:
            mcp_protocol: Shared MCP protocol for all campfires
        """
        self.campfires: Dict[str, Campfire] = {}
        self.mcp_protocol = mcp_protocol
        self.is_running = False
    
    def add_campfire(self, campfire: Campfire) -> None:
        """
        Add a campfire to the manager.
        
        Args:
            campfire: Campfire to add
        """
        self.campfires[campfire.name] = campfire
        
        # Set MCP protocol if not already set
        if self.mcp_protocol and not campfire.mcp_protocol:
            campfire.mcp_protocol = self.mcp_protocol
    
    async def start_all(self) -> None:
        """Start all campfires."""
        self.is_running = True
        
        tasks = []
        for campfire in self.campfires.values():
            task = asyncio.create_task(campfire.start())
            tasks.append(task)
        
        await asyncio.gather(*tasks)
    
    async def stop_all(self) -> None:
        """Stop all campfires."""
        self.is_running = False
        
        tasks = []
        for campfire in self.campfires.values():
            task = asyncio.create_task(campfire.stop())
            tasks.append(task)
        
        await asyncio.gather(*tasks)
    
    def get_campfire(self, name: str) -> Optional[Campfire]:
        """Get a campfire by name."""
        return self.campfires.get(name)
    
    def get_all_stats(self) -> Dict[str, Any]:
        """Get stats for all campfires."""
        return {
            'manager_running': self.is_running,
            'campfires_count': len(self.campfires),
            'campfires': {name: cf.get_stats() for name, cf in self.campfires.items()}
        }