from typing import List, Dict, Any, Optional, Protocol, Union, Tuple, runtime_checkable
from dataclasses import dataclass
import logging
import asyncio
import time
from abc import ABC, abstractmethod
from sentence_transformers import CrossEncoder
import numpy as np

# Configure logging
logger = logging.getLogger(__name__)


@dataclass
class RetrievalResult:
    """Represents a single retrieval result."""
    document: str
    metadata: Dict[str, Any]
    score: float
    rank: Optional[int] = None
    source: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary."""
        return {
            "document": self.document,
            "metadata": self.metadata,
            "score": self.score,
            "rank": self.rank,
            "source": self.source
        }


@runtime_checkable
class RetrieverSource(Protocol):
    """Protocol for a retrieval source that can be queried asynchronously."""
    
    async def aget_relevant_documents(
        self, 
        query: str, 
        top_k: int = 5, 
        **kwargs
    ) -> List[RetrievalResult]:
        """Retrieve relevant documents from this source."""
        ...


class VectorDatabaseProtocol(Protocol):
    """Protocol defining the interface for vector databases."""
    
    def query(self, *args, **kwargs) -> Dict[str, Any]:
        """Query the vector database and return results."""
        ...


class VectorSourceAdapter:
    """Standardized source for vector databases."""
    
    def __init__(self, vector_db: Any, embedder: Any, collection_name: Optional[str] = None):
        self.vector_db = VectorDatabaseAdapter(vector_db)
        self.embedder = embedder
        self.collection_name = collection_name

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        from .tracing_helper import trace_span
        async with trace_span("vector_source", {"query": query, "top_k": top_k}):
            try:
                # Embedding is usually fast but we could make it async if needed
                # For now assume embedder.get_embedding is sync or we wrap it
                if asyncio.iscoroutinefunction(self.embedder.get_embedding):
                    embedding = await self.embedder.get_embedding(query)
                else:
                    embedding = self.embedder.get_embedding(query)
                
                raw_response = self.vector_db.query(
                    query_embedding=embedding,
                    top_k=top_k,
                    collection_name=kwargs.get("collection_name", self.collection_name)
                )
                
                docs = raw_response.get("documents", [])
                metas = raw_response.get("metadatas", [])
                scores = raw_response.get("scores", [])
                
                results = []
                for doc, meta, score in zip(docs, metas, scores):
                    results.append(RetrievalResult(
                        document=doc,
                        metadata=meta,
                        score=float(score),
                        source="vector"
                    ))
                return results
            except Exception as e:
                logger.error(f"Vector Source error: {e}")
                return []


class GraphSourceAdapter:
    """Source for Knowledge Graph (Neo4j, Memgraph, MongoDB)."""
    
    def __init__(self, graph_helper: Any, index_name: Optional[str] = None):
        self.graph_helper = graph_helper
        self.index_name = index_name

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        from .tracing_helper import trace_span
        async with trace_span("graph_source", {"query": query, "top_k": top_k}):
            try:
                # Graph retrieval often involves cypher queries or vector search within the graph DB
                # If helper has search_similar_by_vector
                if hasattr(self.graph_helper, "search_similar_by_vector_async"):
                    embedding = kwargs.get("embedding")
                    if not embedding and "embedder" in kwargs:
                        embedder = kwargs["embedder"]
                        if asyncio.iscoroutinefunction(embedder.get_embedding):
                            embedding = await embedder.get_embedding(query)
                        else:
                            embedding = embedder.get_embedding(query)
                    
                    if embedding:
                        hits = await self.graph_helper.search_similar_by_vector_async(
                            embedding=embedding,
                            top_k=top_k,
                            index_name=kwargs.get("index_name", self.index_name)
                        )
                        return [RetrievalResult(
                            document=h.get("content") or h.get("text") or str(h),
                            metadata=h.get("metadata", h),
                            score=h.get("score", 0.5),
                            source="graph"
                        ) for h in hits]
                
                elif hasattr(self.graph_helper, "search_similar_by_vector"):
                    # Fallback to sync
                    embedding = kwargs.get("embedding")
                    # ... same logic for sync ...
                    if embedding:
                        hits = self.graph_helper.search_similar_by_vector(
                            embedding=embedding,
                            top_k=top_k,
                            index_name=kwargs.get("index_name", self.index_name)
                        )
                        return [RetrievalResult(
                            document=h.get("content") or h.get("text") or str(h),
                            metadata=h.get("metadata", h),
                            score=h.get("score", 0.5),
                            source="graph"
                        ) for h in hits]
                
                return []
            except Exception as e:
                logger.error(f"Graph Source error: {e}")
                return []


class MCPSourceAdapter:
    """Source for MCP Tools."""
    
    def __init__(self, mcp_manager: Any, server_name: str, tool_name: str):
        self.mcp_manager = mcp_manager
        self.server_name = server_name
        self.tool_name = tool_name

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        from .tracing_helper import trace_span
        async with trace_span("mcp_source", {"query": query, "server": self.server_name, "tool": self.tool_name}):
            try:
                response = await self.mcp_manager.call_tool(
                    server_name=self.server_name,
                    tool_name=self.tool_name,
                    arguments={"query": query, "top_k": top_k}
                )
                
                # MCP response format varies, but usually it's text or a list of items
                if isinstance(response, list):
                    return [RetrievalResult(
                        document=str(item.get("content", item)),
                        metadata=item.get("metadata", {}),
                        score=1.0, # MCP rarely returns standardized scores
                        source=f"mcp:{self.server_name}"
                    ) for item in response[:top_k]]
                
                return [RetrievalResult(document=str(response), metadata={}, score=1.0, source=f"mcp:{self.server_name}")]
            except Exception as e:
                logger.error(f"MCP Source error ({self.server_name}): {e}")
                return []


class GCSSourceAdapter:
    """Source for Google Cloud Storage."""
    
    def __init__(self, gcs_helper: Any, bucket_name: str):
        self.gcs_helper = gcs_helper
        self.bucket_name = bucket_name

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        try:
            # GCS doesn't have native content search, but we can search by prefix or metadata
            # This is a placeholder for metadata-based retrieval
            blobs = self.gcs_helper.list_blobs(self.bucket_name, prefix=kwargs.get("prefix", ""))
            hits = []
            for blob in blobs:
                if query.lower() in blob.name.lower():
                    hits.append(RetrievalResult(
                        document=f"File: {blob.name} in bucket {self.bucket_name}",
                        metadata={"name": blob.name, "bucket": self.bucket_name, "type": "gcs_blob"},
                        score=0.8,
                        source="gcs"
                    ))
                if len(hits) >= top_k:
                    break
            return hits
        except Exception as e:
            logger.error(f"GCS Source error: {e}")
            return []


class DriveSourceAdapter:
    """Source for Google Drive."""
    
    def __init__(self, drive_helper: Any):
        self.drive_helper = drive_helper

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        try:
            # Google Drive API supports 'q' parameter for search
            files = self.drive_helper.search_files(query=f"name contains '{query}'", max_results=top_k)
            return [RetrievalResult(
                document=f"Drive File: {f.get('name')} (ID: {f.get('id')})",
                metadata=f,
                score=0.9,
                source="google_drive"
            ) for f in files]
        except Exception as e:
            logger.error(f"Drive Source error: {e}")
            return []


class LocalFileSourceAdapter:
    """Source for local files using DocumentLoaderHelper."""
    
    def __init__(self, loader_helper: Any, base_path: str):
        self.loader_helper = loader_helper
        self.base_path = base_path

    async def aget_relevant_documents(self, query: str, top_k: int = 5, **kwargs) -> List[RetrievalResult]:
        try:
            # For local files, we might search filenames or content if small enough
            # For now, let's treat this as a keyword search over file contents in a directory
            import os
            hits = []
            for root, _, files in os.walk(self.base_path):
                for file in files:
                    if any(file.endswith(ext) for ext in ['.txt', '.md', '.pdf', '.docx']):
                        full_path = os.path.join(root, file)
                        try:
                            # Use loader to load and check content
                            # This is expensive for many files, but works for focused directories
                            docs = self.loader_helper.load_file(full_path)
                            for doc in docs:
                                if query.lower() in doc.get("content", "").lower():
                                    hits.append(RetrievalResult(
                                        document=doc.get("content")[:1000],
                                        metadata={"path": full_path, "type": "local_file"},
                                        score=0.7,
                                        source="local_files"
                                    ))
                                if len(hits) >= top_k:
                                    break
                        except:
                            continue
                    if len(hits) >= top_k:
                        break
            return hits
        except Exception as e:
            logger.error(f"Local File Source error: {e}")
            return []


class ParallelRetrieverX:
    """
    Advanced Parallel Retriever that orchestrates multiple sources.
    Uses Reciprocal Rank Fusion (RRF) to merge results.
    """
    
    def __init__(
        self, 
        sources: List[RetrieverSource] = None,
        reranker_model: Optional[str] = None,
        use_rerank: bool = False,
        rrf_k: int = 60
    ):
        self.sources = sources or []
        self.rrf_k = rrf_k
        self.use_rerank = use_rerank
        self.reranker = CrossEncoder(reranker_model) if reranker_model and use_rerank else None

    def add_source(self, source: RetrieverSource):
        self.sources.append(source)

    async def retrieve(
        self, 
        query: str, 
        top_k: int = 10, 
        **kwargs
    ) -> List[RetrievalResult]:
        """Parallel retrieval with RRF fusion."""
        if not self.sources:
            logger.warning("No sources configured for ParallelRetrieverX")
            return []

        # 1. Dispatch parallel requests
        from .tracing_helper import trace_span
        async with trace_span("parallel_retrieve", {"query": query, "sources_count": len(self.sources)}):
            tasks = [source.aget_relevant_documents(query, top_k=top_k*2, **kwargs) for source in self.sources]
            source_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 2. Reciprocal Rank Fusion (RRF)
        # RRF Score(d) = sum( 1 / (k + rank(d, s)) ) for each source s
        all_docs = {}
        for results in source_results:
            if isinstance(results, Exception):
                logger.error(f"Source task failed: {results}")
                continue
                
            for rank, res in enumerate(results, 1):
                doc_key = res.document # Simple heuristic, ideally use a hash or ID
                if doc_key not in all_docs:
                    all_docs[doc_key] = {"res": res, "rrf_score": 0.0}
                
                all_docs[doc_key]["rrf_score"] += 1.0 / (self.rrf_k + rank)

        # 3. Sort by RRF score
        fused_results = []
        for key, data in all_docs.items():
            res = data["res"]
            res.score = data["rrf_score"] # Overwrite score with fused score
            fused_results.append(res)
        
        fused_results.sort(key=lambda x: x.score, reverse=True)
        final_results = fused_results[:top_k * 2] # Keep buffer for reranking

        # 4. Optional Reranking
        if self.reranker and final_results:
            pairs = [[query, r.document] for r in final_results]
            rerank_scores = self.reranker.predict(pairs)
            for res, score in zip(final_results, rerank_scores):
                res.score = float(score)
            final_results.sort(key=lambda x: x.score, reverse=True)

        return final_results[:top_k]


# Keep legacy classes for backward compatibility but encourage EnhancedRetrieverX/ParallelRetrieverX
class EnhancedRetrieverX:
    """Legacy wrapper for backward compatibility or simple two-stage usage."""
    def __init__(self, vector_db, embedder, **kwargs):
        self.source = VectorSourceAdapter(vector_db, embedder)
        self.orchestrator = ParallelRetrieverX(sources=[self.source], **kwargs)

    def retrieve(self, query: str, top_k: int = 10, **kwargs) -> List[RetrievalResult]:
        # Sync wrapper for async retrieve
        return asyncio.run(self.orchestrator.retrieve(query, top_k, **kwargs))


class RetrieverX:
    """Ultra-legacy RetrieverX class."""
    def __init__(self, vector_db, embedder, reranker_model=None, use_rerank=True):
        self.retriever = EnhancedRetrieverX(vector_db, embedder, reranker_model=reranker_model, use_rerank=use_rerank)

    def retrieve(self, query: str, collection_name: str, top_k: int = 10) -> List[Dict[str, Any]]:
        results = self.retriever.retrieve(query, top_k, collection_name=collection_name)
        return [{"document": r.document, "metadata": r.metadata, "score": r.score} for r in results]
