import os
import zipfile
import requests
import json
from typing import List
from pyserini.search.lucene import LuceneSearcher
from rankify.dataset.dataset import Document, Context
from tqdm import tqdm
from pyserini.eval.evaluate_dpr_retrieval import has_answers, SimpleTokenizer
from rankify.utils.pre_defind_models import INDEX_TYPE

class BM25Retriever:
    """
    Implements **BM25** `[2]`_ , a **probabilistic ranking model** for retrieving documents from large-scale corpora.

    .. _[2]: https://www.nowpublishers.com/article/Details/INR-019

    BM25, based on the **Probabilistic Relevance Framework**, ranks documents based on their relevance to a given query. 
    This retriever utilizes **Pyserini's LuceneSearcher** for efficient batch-based retrieval.

    References
    ----------
    .. [2] Robertson, S. & Zaragoza, H. (2009). *The Probabilistic Relevance Framework: BM25 and Beyond*.
       Foundations and Trends® in Information Retrieval, 3(4), 333-389. 
       DOI: [https://www.nowpublishers.com/article/Details/INR-019](https://www.nowpublishers.com/article/Details/INR-019)

    Attributes
    ----------
    n_docs : int
        Number of top documents to retrieve for each query.
    batch_size : int
        Number of queries processed in a single batch for efficiency.
    threads : int
        Number of parallel threads used for retrieval.
    index_type : str
        The type of index (`wiki` or `msmarco`) used for retrieval.
    index_url : str
        URL to download the prebuilt BM25 index.
    index_folder : str
        Path where the BM25 index is stored.
    index_path : str
        Path to the Lucene index used by Pyserini.
    title_map_path : str
        Path to a JSON file mapping passage IDs to titles.
    searcher : LuceneSearcher
        Pyserini-based search engine for BM25 retrieval.
    pid2title : dict
        Dictionary mapping document IDs to their titles.
    
    Examples
    --------
    Basic Usage:

    >>> from rankify.dataset.dataset import Document, Question
    >>> from rankify.retrievers.retriever import Retriever
    >>> retriever = Retriever(model="bm25", n_docs=5, index_type="wiki")
    >>> documents = [Document(question=Question("Who discovered gravity?"))]
    >>> retrieved_documents = retriever.retrieve(documents)
    >>> print(retrieved_documents[0].contexts[0].text)
    """
    def __init__(self,model="bm25", n_docs: int = 10, batch_size: int = 36, threads: int = 30 , index_type: str = 'wiki') -> None:
        """
        Initializes the BM25 retriever.

        Parameters
        ----------
        n_docs : int, optional
            Number of top documents to retrieve per query (default is 10).
        batch_size : int, optional
            Number of queries processed in a single batch (default is 36).
        threads : int, optional
            Number of parallel threads used for retrieval (default is 30).
        index_type : str, optional
            Type of index to use (`wiki` or `msmarco`) (default is "wiki").
        
        Raises
        ------
        ValueError
            If the specified index type is not supported.
        """
        self.n_docs = n_docs
        self.batch_size = batch_size
        self.threads = threads
        self.tokenizer = SimpleTokenizer()

        if 'bm25' not in INDEX_TYPE or index_type not in INDEX_TYPE['bm25']:
            raise ValueError(f"Index type {index_type} is not supported.")
        
        self.index_type = index_type
        self.index_url = INDEX_TYPE['bm25'][index_type]['url']
        self.index_folder = os.path.join(os.environ.get("RERANKING_CACHE_DIR", "./cache"), 'index',  f"bm25_index_{index_type}")
        if index_type =="wiki":
            self.index_path =  os.path.join(self.index_folder, f"bm25_index")
        else:
            self.index_path =  os.path.join(self.index_folder, f"bm25_index_{index_type}")
            
        self.title_map_path = os.path.join(self.index_path, "corpus.json")


        self._ensure_index_downloaded()
        self.searcher = LuceneSearcher(self.index_path)
        with open(self.title_map_path, "r", encoding="utf-8") as f:
            self.pid2title = json.load(f)

    def _ensure_index_downloaded(self) -> None:
        """
        Ensures that the BM25 index and associated files are downloaded and extracted.

        If the index does not exist locally, it downloads and extracts the index from the 
        specified URL.

        Raises
        ------
        RuntimeError
            If the index download fails.
        """
        if not os.path.exists(self.index_path):
            os.makedirs(self.index_folder, exist_ok=True)
            zip_path = os.path.join(self.index_folder, f"bm25_index_{self.index_type}.zip")
            print("Downloading BM25 index...")
            response = requests.get(self.index_url, stream=True)
            with open(zip_path, "wb") as f:
                for chunk in tqdm(response.iter_content(chunk_size=1024), desc="Downloading BM25 Index"):
                    f.write(chunk)
            print("Extracting BM25 index...")
            with zipfile.ZipFile(zip_path, "r") as zip_ref:
                zip_ref.extractall(self.index_folder)
            os.remove(zip_path)
    def retrieve__(self, documents: List[Document]) -> List[Document]:
        """
        Retrieves relevant contexts for each document in the input list using BM25.

        Parameters
        ----------
        documents : List[Document]
            A list of `Document` objects containing queries and expected answers.

        Returns
        -------
        List[Document]
            The input `Document` objects updated with retrieved contexts.
        """
        print(f"Retrieving {len(documents)} documents one at a time...")

        for document in tqdm(documents, desc="Retrieving documents", unit="doc"):
            query = document.question.question  # Extract the query string
            try:
                # Perform search for the single query
                hits = self.searcher.search(query, self.n_docs)

                contexts = []
                for hit in hits:
                    try:
                        lucene_doc = self.searcher.doc(hit.docid)
                        print(lucene_doc.raw())
                        raw_content = json.loads(lucene_doc.raw())  # Parse the raw JSON
                        text = raw_content.get("contents", "")
                        title = self.pid2title.get(hit.docid, "No Title")

                        context = Context(
                            id=int(hit.docid),
                            title=title,
                            text=text,
                            score=hit.score,
                            has_answer=has_answers(text, document.answers.answers, self.tokenizer)
                        )
                        contexts.append(context)
                    except Exception as e:
                        print(f"Error processing document ID {hit.docid}: {e}")

                # Assign the retrieved contexts to the document
                document.contexts = contexts

            except Exception as e:
                print(f"Error retrieving contexts for query '{query}': {e}")

        return documents
    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Performs batch search using Pyserini's LuceneSearcher.

        Parameters
        ----------
        queries : List[str]
            List of queries to search.
        qids : List[str]
            List of unique query IDs.
        k : int
            Number of documents to retrieve per query.
        batch_size : int
            Number of queries processed per batch.
        threads : int
            Number of threads to use for query processing.

        Returns
        -------
        dict
            Dictionary mapping query IDs to search results.
        
        Raises
        ------
        RuntimeError
            If batch search fails for any query.
        """
        # Extract the question text and unique IDs
        question_texts = [doc.question.question for doc in documents]  # Extract plain string questions
        batch_qids = [str(i) for i in range(len(question_texts))]  # Generate unique IDs for each query

        print(f"Retrieving {len(documents)} documents in batches of {self.batch_size}...")

        # Perform batch search for all questions
        batch_results = self._batch_search(question_texts, batch_qids, self.n_docs, self.batch_size, self.threads)

        for i, qid in enumerate(tqdm(batch_qids, desc="Processing documents", unit="doc")):
            document = documents[i]
            hits = batch_results.get(qid, [])

            contexts = []
            for hit in hits:
                try:
                    lucene_doc = self.searcher.doc(hit.docid)
                    raw_content = json.loads(lucene_doc.raw())  # Parse the raw JSON
                    #print(raw_content)
                    text = raw_content.get("contents", "")
                    title = self.pid2title.get(hit.docid, "No Title")

                    context = Context(
                        id=int(hit.docid),
                        title=title,
                        text=text,
                        score=hit.score,
                        has_answer=has_answers(text, document.answers.answers, self.tokenizer)
                    )
                    contexts.append(context)
                except Exception as e:
                    print(f"Error processing document ID {hit.docid}: {e}")

            document.contexts = contexts

        return documents

    def _batch_search(self, queries: List[str], qids: List[str], k: int, batch_size: int, threads: int):
        """
        Performs batch search using Pyserini's batch_search method.
        """
        batch_results = {}
        for start in tqdm(range(0, len(queries), batch_size), desc="Batch search"):
            end = min(start + batch_size, len(queries))
            batch_queries = queries[start:end]
            batch_qids = qids[start:end]

            try:
                batch_hits = self.searcher.batch_search(batch_queries, batch_qids, k=k, threads=threads)
                batch_results.update(batch_hits)
            except Exception as e:
                print(f"Batch search failed for queries {batch_queries} with error {e}")

        return batch_results
