import json
from elasticsearch import Elasticsearch
from tqdm import tqdm

class BM25Retriever:
    def __init__(self, es_host="http://localhost:9200", index_name="bm25_index"):
        """
        Initialize the BM25 Retriever.
        
        Args:
            es_host (str): Elasticsearch host URL
            index_name (str): Name of the Elasticsearch index
        """
        self.es = Elasticsearch(es_host)
        self.index_name = index_name
    
    def create_index(self):
        """
        Create or recreate the Elasticsearch index with appropriate mappings.
        """
        if self.es.indices.exists(index=self.index_name):
            self.es.indices.delete(index=self.index_name)
        
        self.es.indices.create(
            index=self.index_name,
            body={
                "mappings": {
                    "properties": {
                        "text": {
                            "type": "text",
                            "analyzer": "standard"
                        },
                        "title": {
                            "type": "keyword"
                        },
                        "chunk_id": {
                            "type": "integer"
                        }
                    }
                }
            }
        )
    
    def index_corpus(self, corpus_file=None, corpus_data=None):
        """
        Index the corpus documents into Elasticsearch.
        
        Args:
            corpus_file (str): Path to the corpus JSONL file
        """
        actions = []
        assert corpus_file or corpus_data, "Either corpus_file or corpus_data must be provided."
        if corpus_data:
            for _, data in enumerate(corpus_data):
                chunk_id = data['chunk_id']
                text = data['text']
                title = data.get('title', '')
                actions.append({"index": {"_index": self.index_name, "_id": chunk_id}})
                actions.append({"text": text, "title": title, "chunk_id": chunk_id})
                if len(actions) >= 2000:
                    self.es.bulk(body=actions, refresh=True)
                    actions = []
        else:
            with open(corpus_file, 'r') as f:
                for line in tqdm(f, desc="Indexing corpus"):
                    doc = json.loads(line)
                    action = {"index": {"_index": self.index_name, "_id": doc["chunk_id"]}}
                    actions.append(action)
                    actions.append(doc)
                    if len(actions) >= 2000:
                        self.es.bulk(body=actions, refresh=True)
                        actions = []
        if actions:
            self.es.bulk(body=actions, refresh=True)
    
    def search_queries(self, queries_file=None, query_data=None, top_k=20):
        """
        Search queries against the indexed corpus.
        
        Args:
            queries_file (str): Path to the queries JSONL file
            top_k (int): Number of top documents to retrieve
            
        Returns:
            list: List of search results with query info and top-k document IDs
        """
        results = []
        assert queries_file or query_data, "Either queries_file or query_data must be provided."
        if query_data:
            for query in query_data:
                query_data = query['question']
                response = self.es.search(
                    index=self.index_name,
                    body={
                        "query": {
                            "match": {
                                "text": query_data
                            }
                        },
                        "size": top_k
                    }
                )
                
                top_k_doc_ids = [int(hit["_id"]) for hit in response["hits"]["hits"]]
                
                results.append({
                    "qid": query['qid'],
                    "text": query_data,
                    "answer": query['answers'],
                    "top_k_doc_id": top_k_doc_ids
                })
        else:
            with open(queries_file, 'r') as f:
                for line in tqdm(f, desc="Searching queries"):
                    query_data = json.loads(line)
                    query = query_data["question"]
                    
                    response = self.es.search(
                        index=self.index_name,
                        body={
                            "query": {
                                "match": {
                                    "text": query
                                }
                            },
                            "size": top_k
                        }
                    )
                    
                    top_k_doc_ids = [int(hit["_id"]) for hit in response["hits"]["hits"]]
                    
                    results.append({
                        "qid": query_data["id"],
                        "text": query_data["question"],
                        "answer": query_data["answers"],
                        "top_k_doc_id": top_k_doc_ids
                    })
        return results
    
    def save_results(self, results, output_file):
        """
        Save search results to a JSONL file.
        
        Args:
            results (list): List of search results
            output_file (str): Path to the output JSONL file
        """
        with open(output_file, 'w') as f:
            for result in results:
                f.write(json.dumps(result) + "\n")
    
    def run_retrieval(self, corpus_file, queries_file, output_file, top_k=20):
        """
        Run the complete retrieval pipeline.
        
        Args:
            corpus_file (str): Path to the corpus JSONL file
            queries_file (str): Path to the queries JSONL file
            output_file (str): Path to the output JSONL file
            top_k (int): Number of top documents to retrieve
        """
        print("Creating index...")
        self.create_index()
        
        print("Indexing corpus...")
        self.index_corpus(corpus_file)
        
        print("Searching queries...")
        results = self.search_queries(queries_file, top_k)
        
        print(f"Writing results to {output_file}...")
        self.save_results(results, output_file)
        
        print("Done.")