"""Retrieval orchestrator for TCMGM."""
import logging
from typing import List, Dict, Optional
from datetime import datetime

from outhad_contextkit.memory.temporal.types import TimeWindow

logger = logging.getLogger(__name__)


class RetrievalOrchestrator:
    """Orchestrates retrieval across vector, graph, and timeline."""
    
    def __init__(self, vector_store, graph_store, timeline_builder, embedding_model=None):
        """
        Initialize the retrieval orchestrator.
        
        Args:
            vector_store: Vector store instance for semantic search
            graph_store: Graph store instance for entity relationships
            timeline_builder: TimelineBuilder instance for temporal queries
            embedding_model: Embedding model for generating query embeddings
        """
        self.vector_store = vector_store
        self.graph_store = graph_store
        self.timeline = timeline_builder
        self.embedding_model = embedding_model
        logger.info("RetrievalOrchestrator initialized successfully")
    
    def fused_search(
        self,
        query: str,
        user_id: str,
        time_window: Optional[TimeWindow] = None,
        include_causal: bool = True,
        include_multimodal: bool = True,
        top_k: int = 10
    ) -> Dict:
        """
        Fused retrieval combining vector, graph, and timeline.
        
        Args:
            query: Search query
            user_id: User identifier
            time_window: Optional time window filter
            include_causal: Include causal chain exploration
            include_multimodal: Include cross-modal results
            top_k: Number of results
        
        Returns:
            Dict with merged results from all retrieval methods
        """
        logger.info(f"Starting fused search for query: '{query}' (user: {user_id})")
        
        results = {
            "vector_results": [],
            "graph_results": [],
            "timeline_results": [],
            "causal_chains": [],
            "multimodal_results": [],
            "fused_ranking": []
        }
        
        # 1. Vector search (semantic similarity)
        logger.debug("Performing vector search...")
        vector_results = self._vector_search(query, user_id, time_window, top_k)
        results["vector_results"] = vector_results
        logger.info(f"Vector search returned {len(vector_results)} results")
        
        # 2. Graph search (entity relationships)
        logger.debug("Performing graph search...")
        graph_results = self._graph_search(query, user_id, time_window)
        results["graph_results"] = graph_results
        logger.info(f"Graph search returned {len(graph_results)} results")
        
        # 3. Timeline search (temporal proximity)
        if time_window:
            logger.debug("Performing timeline search...")
            timeline_results = self._timeline_search(user_id, time_window)
            results["timeline_results"] = timeline_results
            logger.info(f"Timeline search returned {len(timeline_results)} results")
        
        # 4. Causal chain exploration
        if include_causal and vector_results:
            logger.debug("Exploring causal chains...")
            causal_chains = self._explore_causal_chains(vector_results, user_id)
            results["causal_chains"] = causal_chains
            logger.info(f"Found {len(causal_chains)} causal chains")
        
        # 5. Cross-modal search
        if include_multimodal:
            logger.debug("Performing cross-modal search...")
            multimodal_results = self._cross_modal_search(query, user_id)
            results["multimodal_results"] = multimodal_results
            logger.info(f"Cross-modal search returned {len(multimodal_results)} results")
        
        # 6. Fuse and rank all results
        logger.debug("Fusing and ranking results...")
        fused = self._fuse_results(results, query)
        results["fused_ranking"] = fused[:top_k]
        logger.info(f"Fused ranking complete: {len(results['fused_ranking'])} results")
        
        return results
    
    def _vector_search(
        self,
        query: str,
        user_id: str,
        time_window: Optional[TimeWindow],
        top_k: int
    ) -> List[Dict]:
        """Perform vector search."""
        try:
            # Generate query embeddings using the embedding model
            if not self.embedding_model:
                logger.warning("No embedding model available for vector search")
                return []
            
            query_embeddings = self.embedding_model.embed(query, "search")
            
            # Use the vector store's search method
            filters = {"user_id": user_id}
            
            # CRITICAL FIX: Do NOT filter vector search by time window
            # Vector search should return semantically relevant results regardless of time
            # Temporal filtering is handled by timeline search and in result fusion
            # Filtering by timestamp here causes 0 results because timestamps aren't properly indexed
            
            # Perform search with embeddings (no temporal filtering)
            search_results = self.vector_store.search(
                query=query,
                vectors=query_embeddings,
                filters=filters,
                limit=top_k
            )
            
            # Format results
            formatted_results = []
            for result in search_results:
                if hasattr(result, 'payload'):
                    # Qdrant/Chroma format
                    formatted_results.append({
                        "id": result.id,
                        "content": result.payload.get("data", ""),
                        "score": getattr(result, 'score', 0.0),
                        "metadata": result.payload,
                        "source": "vector"
                    })
                elif isinstance(result, dict):
                    # Dict format (already formatted from Memory.search)
                    formatted_results.append({
                        "id": result.get("id", ""),
                        "content": result.get("memory", ""),
                        "score": result.get("score", 0.0),
                        "metadata": result,
                        "source": "vector"
                    })
            
            return formatted_results
        except Exception as e:
            logger.error(f"Vector search failed: {e}")
            return []
    
    def _graph_search(
        self,
        query: str,
        user_id: str,
        time_window: Optional[TimeWindow]
    ) -> List[Dict]:
        """Perform graph search."""
        try:
            filters = {"user_id": user_id}
            
            # Use graph's search method
            graph_results = self.graph_store.search(query, filters)
            
            # Format results
            formatted_results = []
            for result in graph_results:
                if isinstance(result, dict):
                    formatted_results.append({
                        "id": result.get("id", ""),
                        "content": result.get("name", ""),
                        "score": 0.8,  # Default relevance score
                        "metadata": result,
                        "source": "graph"
                    })
                else:
                    formatted_results.append({
                        "id": str(result),
                        "content": str(result),
                        "score": 0.8,
                        "metadata": {},
                        "source": "graph"
                    })
            
            return formatted_results
        except Exception as e:
            logger.error(f"Graph search failed: {e}")
            return []
    
    def _timeline_search(
        self,
        user_id: str,
        time_window: TimeWindow
    ) -> List[Dict]:
        """Search timeline within time window."""
        try:
            timeline_events = self.timeline.get_timeline(user_id, time_window)
            
            # Format results
            formatted_results = []
            for event in timeline_events:
                formatted_results.append({
                    "id": event.get("id", ""),
                    "content": event.get("content", ""),
                    "score": 0.7,  # Default temporal relevance
                    "timestamp": event.get("timestamp", ""),
                    "metadata": event,
                    "source": "timeline"
                })
            
            return formatted_results
        except Exception as e:
            logger.error(f"Timeline search failed: {e}")
            return []
    
    def _explore_causal_chains(
        self,
        seed_events: List[Dict],
        user_id: str,
        max_depth: int = 3
    ) -> List[Dict]:
        """Explore causal chains from seed events."""
        from outhad_contextkit.memory.temporal.causal_queries import get_causal_chain
        
        chains = []
        for event in seed_events[:3]:  # Top 3 events
            event_id = event.get('id')
            if not event_id:
                continue
            
            try:
                # Get forward chain (effects)
                forward_chain = get_causal_chain(
                    self.graph_store.graph,
                    event_id,
                    {"user_id": user_id},
                    max_depth=max_depth,
                    direction="forward"
                )
                
                # Get backward chain (causes)
                backward_chain = get_causal_chain(
                    self.graph_store.graph,
                    event_id,
                    {"user_id": user_id},
                    max_depth=max_depth,
                    direction="backward"
                )
                
                chains.append({
                    "seed_event": event,
                    "forward_chain": forward_chain,
                    "backward_chain": backward_chain
                })
            except Exception as e:
                logger.error(f"Failed to explore causal chain for event {event_id}: {e}")
        
        return chains
    
    def _cross_modal_search(
        self,
        query: str,
        user_id: str
    ) -> List[Dict]:
        """Perform cross-modal search."""
        try:
            from outhad_contextkit.memory.temporal.cross_modal import cross_modal_search
            
            # Get query embedding
            # This would use the multimodal embedder if available
            # For now, return empty as this requires multimodal content in the database
            
            # TODO: Implement full cross-modal search when multimodal content is available
            # This would:
            # 1. Get query embedding
            # 2. Query for nodes with modality != 'text'
            # 3. Perform cross-modal similarity search
            
            logger.debug("Cross-modal search placeholder - no multimodal content indexed")
            return []
        except Exception as e:
            logger.error(f"Cross-modal search failed: {e}")
            return []
    
    def _fuse_results(
        self,
        all_results: Dict,
        query: str
    ) -> List[Dict]:
        """Fuse and rank results from all sources."""
        fused = []
        
        # Add vector results with weight
        for result in all_results["vector_results"]:
            fused.append({
                "content": result.get("content", ""),
                "id": result.get("id", ""),
                "source": "vector",
                "score": result.get("score", 0.0),
                "weight": 0.4,  # Vector weight
                "metadata": result.get("metadata", {})
            })
        
        # Add graph results with weight
        for result in all_results["graph_results"]:
            fused.append({
                "content": result.get("content", ""),
                "id": result.get("id", ""),
                "source": "graph",
                "score": 0.7,  # Default score for graph
                "weight": 0.3,  # Graph weight
                "metadata": result.get("metadata", {})
            })
        
        # Add timeline results with weight
        for result in all_results["timeline_results"]:
            fused.append({
                "content": result.get("content", ""),
                "id": result.get("id", ""),
                "source": "timeline",
                "score": 0.6,  # Default temporal relevance
                "weight": 0.3,  # Timeline weight
                "timestamp": result.get("timestamp", ""),
                "metadata": result.get("metadata", {})
            })
        
        # Calculate final scores
        for item in fused:
            item["final_score"] = item["score"] * item["weight"]
        
        # Sort by final score
        fused.sort(key=lambda x: x["final_score"], reverse=True)
        
        # Deduplicate by content while preserving order
        seen_content = set()
        deduplicated = []
        for item in fused:
            content = item.get("content", "")
            if content and content not in seen_content:
                seen_content.add(content)
                deduplicated.append(item)
        
        return deduplicated

