"""
Context Aggregator for PraisonAI Agents.

Aggregates context from multiple memory sources concurrently.
Inspired by CrewAI's ContextualMemory pattern.
"""

import asyncio
import logging
from praisonaiagents._logging import get_logger
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

logger = get_logger(__name__)

@dataclass
class AggregatedContext:
    """Result of context aggregation."""
    context: str = ""
    sources_used: List[str] = field(default_factory=list)
    tokens_used: int = 0
    fetch_times: Dict[str, float] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "context": self.context[:500] + "..." if len(self.context) > 500 else self.context,
            "sources_used": self.sources_used,
            "tokens_used": self.tokens_used,
            "fetch_times": self.fetch_times,
        }

class ContextAggregator:
    """
    Aggregates context from multiple memory sources concurrently.
    
    Implements the CrewAI ContextualMemory pattern for PraisonAI.
    Fetches from memory, knowledge, and RAG in parallel and merges results.
    
    Example:
        aggregator = ContextAggregator()
        
        # Register sources
        aggregator.register_source("memory", memory.search)
        aggregator.register_source("knowledge", knowledge.search)
        aggregator.register_source("rag", rag.retrieve)
        
        # Aggregate context
        result = await aggregator.aggregate("user query", max_tokens=4000)
        print(result.context)
    """
    
    def __init__(
        self,
        max_tokens: int = 4000,
        separator: str = "\n\n---\n\n",
        include_source_labels: bool = True,
    ):
        """
        Initialize context aggregator.
        
        Args:
            max_tokens: Maximum tokens for aggregated context
            separator: Separator between source contexts
            include_source_labels: Whether to label each source section
        """
        self.max_tokens = max_tokens
        self.separator = separator
        self.include_source_labels = include_source_labels
        
        # Registered context sources: name -> (fetch_fn, priority)
        self._sources: Dict[str, Tuple[Callable, int]] = {}
    
    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count (4 chars per token heuristic)."""
        return len(text) // 4 + 1
    
    def register_source(
        self,
        name: str,
        fetch_fn: Callable[[str], Union[str, List[Dict[str, Any]]]],
        priority: int = 50,
    ) -> None:
        """
        Register a context source.
        
        Args:
            name: Source name (e.g., "memory", "knowledge", "rag")
            fetch_fn: Function that takes query and returns context
            priority: Priority (lower = first, 0-100)
        """
        self._sources[name] = (fetch_fn, priority)
    
    def unregister_source(self, name: str) -> None:
        """Remove a registered source."""
        self._sources.pop(name, None)
    
    @property
    def sources(self) -> List[str]:
        """Get list of registered source names."""
        return list(self._sources.keys())
    
    async def _fetch_source(
        self,
        name: str,
        query: str,
    ) -> Tuple[str, Optional[str], float]:
        """
        Fetch context from a single source.
        
        Returns: (context_str, source_name, fetch_time)
        """
        import time
        
        fetch_fn, _ = self._sources.get(name, (None, 0))
        if not fetch_fn:
            return "", None, 0.0
        
        start = time.time()
        try:
            # Call the fetch function
            if asyncio.iscoroutinefunction(fetch_fn):
                result = await fetch_fn(query)
            else:
                # Run sync function in executor
                loop = asyncio.get_event_loop()
                result = await loop.run_in_executor(None, fetch_fn, query)
            
            # Convert result to string
            if isinstance(result, str):
                context = result
            elif isinstance(result, list):
                # Handle list of results (e.g., from memory search)
                parts = []
                for item in result:
                    if isinstance(item, dict):
                        text = item.get("text") or item.get("memory") or item.get("content", "")
                        if text:
                            parts.append(str(text))
                    elif isinstance(item, str):
                        parts.append(item)
                context = "\n".join(parts)
            else:
                context = str(result) if result else ""
            
            elapsed = time.time() - start
            return context, name, elapsed
            
        except Exception as e:
            logger.warning(f"Failed to fetch from {name}: {e}")
            return "", None, time.time() - start
    
    async def aggregate(
        self,
        query: str,
        sources: Optional[List[str]] = None,
        max_tokens: Optional[int] = None,
    ) -> AggregatedContext:
        """
        Aggregate context from multiple sources concurrently.
        
        Args:
            query: Search query
            sources: Specific sources to use (default: all registered)
            max_tokens: Override max tokens
            
        Returns:
            AggregatedContext with merged results
        """
        max_tokens = max_tokens or self.max_tokens
        sources_to_use = sources or list(self._sources.keys())
        
        if not sources_to_use:
            return AggregatedContext()
        
        # Filter to registered sources only
        sources_to_use = [s for s in sources_to_use if s in self._sources]
        
        if not sources_to_use:
            return AggregatedContext()
        
        # Fetch from all sources concurrently
        tasks = [
            self._fetch_source(name, query)
            for name in sources_to_use
        ]
        
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Process results
        contexts: List[Tuple[str, str, float, int]] = []  # (context, name, time, priority)
        fetch_times: Dict[str, float] = {}
        
        for result in results:
            if isinstance(result, Exception):
                logger.warning(f"Source fetch failed: {result}")
                continue
            
            context, name, elapsed = result
            if context and name:
                _, priority = self._sources.get(name, (None, 50))
                contexts.append((context, name, elapsed, priority))
                fetch_times[name] = elapsed
        
        # Sort by priority (lower = first)
        contexts.sort(key=lambda x: x[3])
        
        # Merge and truncate to fit token budget
        merged_parts = []
        total_tokens = 0
        sources_used = []
        separator_tokens = self._estimate_tokens(self.separator)
        
        for context, name, _, _ in contexts:
            context_tokens = self._estimate_tokens(context)
            
            # Check if we can fit this context
            if total_tokens + context_tokens + separator_tokens > max_tokens:
                # Try to fit partial
                remaining = max_tokens - total_tokens - separator_tokens
                if remaining > 100:
                    # Truncate context
                    chars_to_keep = int(remaining * 4 * 0.9)
                    truncated = context[:chars_to_keep] + "..."
                    
                    if self.include_source_labels:
                        merged_parts.append(f"**{name}**:\n{truncated}")
                    else:
                        merged_parts.append(truncated)
                    
                    sources_used.append(name)
                    total_tokens += self._estimate_tokens(truncated)
                break
            
            if self.include_source_labels:
                merged_parts.append(f"**{name}**:\n{context}")
            else:
                merged_parts.append(context)
            
            sources_used.append(name)
            total_tokens += context_tokens + separator_tokens
        
        merged_context = self.separator.join(merged_parts)
        
        return AggregatedContext(
            context=merged_context,
            sources_used=sources_used,
            tokens_used=total_tokens,
            fetch_times=fetch_times,
        )
    
    def aggregate_sync(
        self,
        query: str,
        sources: Optional[List[str]] = None,
        max_tokens: Optional[int] = None,
    ) -> AggregatedContext:
        """
        Synchronous version of aggregate.
        
        Args:
            query: Search query
            sources: Specific sources to use
            max_tokens: Override max tokens
            
        Returns:
            AggregatedContext with merged results
        """
        try:
            # Check if we're in an async context
            try:
                asyncio.get_running_loop()
                # We're in an async context - need to run in thread
                import concurrent.futures
                
                def _run_async_in_thread():
                    """Run async method in new event loop within thread."""
                    return asyncio.run(self.aggregate(query, sources, max_tokens))
                
                with concurrent.futures.ThreadPoolExecutor() as pool:
                    future = pool.submit(_run_async_in_thread)
                    return future.result()
            except RuntimeError:
                # No running event loop - safe to use asyncio.run
                return asyncio.run(self.aggregate(query, sources, max_tokens))
        except Exception:
            # Fallback - create new event loop
            return asyncio.run(self.aggregate(query, sources, max_tokens))

def create_aggregator_from_config(
    memory = None,
    knowledge = None,
    rag = None,
    max_tokens: int = 4000,
) -> ContextAggregator:
    """
    Create a ContextAggregator with standard sources.
    
    Args:
        memory: Memory instance with search() method
        knowledge: Knowledge instance with search() method
        rag: RAG instance with retrieve() method
        max_tokens: Maximum tokens for aggregated context
        
    Returns:
        Configured ContextAggregator
    """
    aggregator = ContextAggregator(max_tokens=max_tokens)
    
    if memory and hasattr(memory, 'search_short_term'):
        aggregator.register_source("memory", memory.search_short_term, priority=10)
    
    if knowledge and hasattr(knowledge, 'search'):
        aggregator.register_source("knowledge", knowledge.search, priority=20)
    
    if rag and hasattr(rag, 'retrieve'):
        aggregator.register_source("rag", rag.retrieve, priority=30)
    
    return aggregator
