"""
The high level BM25 search API. It wraps bm25s into a simple to use search interface,
enabling 1-line indexing and 1-line searching. By default, it will require:
- numba compilation for speed up
- stemming for better search quality
- stopword removal for better search quality
"""

import json
import csv
from pathlib import Path
from bm25s import BM25
from bm25s.tokenization import Tokenizer
import Stemmer
from typing import List



class BM25Search:
    def __init__(
        self,
        corpus: List[str],
        language: str = "english",
        bm25_kwargs: dict = None,
        tokenizer_kwargs: dict = None,
        tokenizer_cls: Tokenizer = Tokenizer,
    ):
        if bm25_kwargs is not None and "corpus" in bm25_kwargs:
            raise ValueError(
                "The 'corpus' argument in bm25_kwargs is reserved and cannot be set manually."
            )
        if language != "english":
            raise NotImplementedError("Currently only English language is supported.")
        
        self.leave_progress = leave_progress = False
        self.show_progress = show_progress = True
        self.corpus = corpus

        stemmer = Stemmer.Stemmer("english")
        bm25_kwargs_default = dict(
            backend="numba", csc_backend="numpy", auto_compile=False
        )
        tokenizer_kwargs_default = dict(
            stemmer=stemmer, stopwords="english", lower=True
        )

        if isinstance(bm25_kwargs, dict):
            bm25_kwargs_default.update(bm25_kwargs)
        elif bm25_kwargs is not None:
            raise ValueError("bm25_kwargs must be a dict or None.")

        if isinstance(tokenizer_kwargs, dict):
            tokenizer_kwargs_default.update(tokenizer_kwargs)
        elif tokenizer_kwargs is not None:
            raise ValueError("tokenizer_kwargs must be a dict or None.")

        if "corpus" in bm25_kwargs_default:
            raise ValueError(
                "The 'corpus' argument in bm25_kwargs is reserved and cannot be set manually."
            )

        # note: we do not pass corpus here, as we will keep it separately. This means the BM25
        # object will return document ids instead of texts when retrieving.
        self.retriever = BM25(**bm25_kwargs_default)
        self.tokenizer: Tokenizer = tokenizer_cls(**tokenizer_kwargs_default)

        # tokenize the corpus
        tokenized = self.tokenizer.tokenize(
            corpus,
            leave_progress=leave_progress,
            show_progress=show_progress,
            update_vocab=True,
            return_as="tuple",
        )

        # compile and index (warmup=False to avoid segfaults in some CI environments)
        self.retriever.compile(activate_numba=True, warmup=False)
        
        create_empty_token = True
        # If the corpus is empty or has no tokens, we can't create an empty token
        # as it relies on vocab dict having some content or logic that fails if empty
        if len(tokenized.vocab) == 0:
            create_empty_token = False
            
        self.retriever.index(
            tokenized,
            leave_progress=leave_progress,
            show_progress=show_progress,
            create_empty_token=create_empty_token,
        )

    def search(self, queries: List[str], k: int = 10, n_jobs: int = 1):
        # Ensure k is not larger than the corpus size
        num_docs = len(self.corpus)
        if k > num_docs:
            k = num_docs

        tokenized_queries = self.tokenizer.tokenize(
            queries,
            update_vocab=False,
            show_progress=self.show_progress,
            leave_progress=self.leave_progress,
            return_as="tuple",
            allow_empty=False,
        )
        
        # Handle empty queries explicitly to avoid issues with Numba backend
        # We filter out queries that result in no tokens
        non_empty_indices = []
        empty_indices = []
        
        # We access internal ids list from Tokenized namedtuple
        # Convert to string list first as retrieve expects that or Tokenized object
        # But checking emptiness is easier on ids
        
        for i, q_ids in enumerate(tokenized_queries.ids):
            if len(q_ids) > 0:
                non_empty_indices.append(i)
            else:
                empty_indices.append(i)
        
        # Prepare results structure
        results = [None] * len(queries)
        
        # Process empty queries immediately
        for i in empty_indices:
            # Empty query results in empty list of documents
            # Or do we want to return empty docs with 0 score? 
            # Standard BM25 usually returns nothing for empty query, or 0 score for all docs.
            # Let's return empty list as top-k results.
            results[i] = []

        if len(non_empty_indices) > 0:
            # Create a new Tokenized object for non-empty queries
            non_empty_ids = [tokenized_queries.ids[i] for i in non_empty_indices]
            non_empty_tokenized = self.tokenizer.to_tokenized_tuple(non_empty_ids)
            
            # Retrieve for non-empty queries
            # note: because we did not pass `corpus` when initializing BM25,
            # the retrieve() will return document ids instead of texts.
            doc_ids, scores = self.retriever.retrieve(
                query_tokens=non_empty_tokenized,
                k=k,
                sorted=True,
                return_as="tuple",
                show_progress=self.show_progress,
                leave_progress=self.leave_progress,
                n_threads=n_jobs,
                chunksize=50,
                backend_selection="auto",
            )
            
            # Map back to original indices
            num_docs = doc_ids.shape[1]
            
            for idx, original_idx in enumerate(non_empty_indices):
                query_results = []
                for di in range(num_docs):
                    doc_id = doc_ids[idx, di]
                    doc_text = self.corpus[doc_id]
                    query_results.append(
                        {
                            "id": int(doc_id),
                            "score": float(scores[idx, di]),
                            "document": doc_text,
                        }
                    )
                results[original_idx] = query_results

        return results


def index(documents, language: str = "english"):
    return BM25Search(corpus=documents, language=language)


def load(path, document_column=None):
    """
    Loads a csv, json, jsonl, or txt file from the given path. For json, we expect a list of dicts.
    Returns a list of strings (corpus) that can be passed to `bm25s.index`.

    Parameters
    ----------
    path : str
        The file path to load the documents from.

    document_column : str, optional
        The column name to use as the document text when loading from csv, or the key when loading from json/jsonl.
        If None, the first column will be used by default.
    
    Returns
    -------
    List[str]
        A list of strings representing the documents.
    """

    path = Path(path)
    documents = []

    if path.suffix == ".txt":
        with open(path, "r", encoding="utf-8") as f:
            documents = [line.strip() for line in f if line.strip()]

    elif path.suffix == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                if len(data) > 0 and isinstance(data[0], str):
                    documents = data
                elif len(data) > 0 and isinstance(data[0], dict):
                    if document_column is None:
                        # Use the first key available in the first element
                        document_column = list(data[0].keys())[0]
                    documents = [d[document_column] for d in data]
            else:
                raise ValueError("JSON file must contain a list of strings or dicts.")

    elif path.suffix == ".jsonl":
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                d = json.loads(line)
                if document_column is None and not documents:
                    # Infer column from first line
                    document_column = list(d.keys())[0]
                
                if document_column in d:
                    documents.append(d[document_column])
                else:
                    # skip or error? Let's skip if key missing, or error. 
                    # raising error is safer for "load"
                    raise ValueError(f"Key '{document_column}' not found in JSONL line: {line[:50]}...")

    elif path.suffix == ".csv":
        with open(path, "r", newline="", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            if document_column is None:
                if reader.fieldnames:
                    document_column = reader.fieldnames[0]
                else:
                    return BM25Search(corpus=[])

            for row in reader:
                if document_column in row:
                    documents.append(row[document_column])
                else:
                     raise ValueError(f"Column '{document_column}' not found in CSV.")
    else:
        raise ValueError(f"Unsupported file extension: {path.suffix}")

    return documents
