import os
import json
from typing import List, Dict
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import linear_kernel
from rank_bm25 import BM25Okapi
import re
import numpy as np

CORPUS_DIR = "../corpus"


def _tokenize(text: str):
    return re.findall(r"\b\w+\b", text.lower())


class SimpleSearchEngine:
    def __init__(self, corpus_dir: str = CORPUS_DIR):
        self.corpus_dir = corpus_dir
        self.documents: List[Dict] = []

        # TF-IDF
        self.vectorizer = None
        self.doc_tfidf = None

        # BM25
        self.bm25 = None
        self._bm25_tokens = []

    def _load_documents(self):
        docs = []
        doc_id = 0

        if not os.path.exists(self.corpus_dir):
            print(f"Warning: Corpus directory '{self.corpus_dir}' does not exist.")
            return

        for root, _, files in os.walk(self.corpus_dir):
            for fname in files:
                if not fname.endswith(".json"):
                    continue

                path = os.path.join(root, fname)
                try:
                    with open(path, "r", encoding="utf-8") as f:
                        data = json.load(f)

                    text = data.get("tf_idf_text", "").strip()
                    if not text:
                        continue

                    docs.append({
                        "id": doc_id,
                        "title": data.get("title", os.path.splitext(fname)[0]),
                        "url": data.get("url", ""),
                        "path": path,
                        "text": text,
                    })
                    doc_id += 1

                except Exception as e:
                    print(f"Error reading {path}: {e}")

        self.documents = docs
        print(f"Loaded {len(self.documents)} documents from {self.corpus_dir}")

    def build_index(self):
        self._load_documents()
        texts = [doc["text"] for doc in self.documents]

        if not texts:
            print("No documents found to index.")
            return

        # TF-IDF
        self.vectorizer = TfidfVectorizer(token_pattern=r"\b\w+\b")
        self.doc_tfidf = self.vectorizer.fit_transform(texts)
        print("TF-IDF index built.")

        # BM25
        self._bm25_tokens = [_tokenize(t) for t in texts]
        self.bm25 = BM25Okapi(self._bm25_tokens)
        print("BM25 index built.")

    def _minmax_normalize(self, scores: np.ndarray):
        min_s = scores.min()
        max_s = scores.max()
        if max_s == min_s:
            return np.zeros_like(scores)
        return (scores - min_s) / (max_s - min_s)

    def search(self, query: str, top_k: int = 5):
        if self.vectorizer is None or self.doc_tfidf is None:
            raise RuntimeError("Index not built. Call build_index() first.")

        query_vec = self.vectorizer.transform([query])
        scores = linear_kernel(query_vec, self.doc_tfidf).flatten()
        ranked_idx = scores.argsort()[::-1][:top_k]

        results = []
        for idx in ranked_idx:
            doc = self.documents[idx]
            results.append({
                "title": doc["title"],
                "url": doc["url"],
                "path": doc["path"],
                "text": doc["text"],
                "score": float(scores[idx])
            })

        return results

    def search_bm25(self, query: str, top_k: int = 5):
        if self.bm25 is None:
            raise RuntimeError("BM25 index not built. Call build_index() first.")

        scores = np.array(self.bm25.get_scores(_tokenize(query)))
        ranked_idx = scores.argsort()[::-1][:top_k]

        results = []
        for idx in ranked_idx:
            doc = self.documents[idx]
            results.append({
                "title": doc["title"],
                "url": doc["url"],
                "path": doc["path"],
                "text": doc["text"],
                "score": float(scores[idx])
            })

        return results

    def search_hybrid(self, query: str, top_k: int = 5, alpha: float = 0.5):
        if self.vectorizer is None or self.doc_tfidf is None or self.bm25 is None:
            raise RuntimeError("Indexes not built. Call build_index() first.")

        # TF-IDF scores
        query_vec = self.vectorizer.transform([query])
        tfidf_scores = linear_kernel(query_vec, self.doc_tfidf).flatten()

        # BM25 scores
        bm25_scores = np.array(
            self.bm25.get_scores(_tokenize(query)),
            dtype=float
        )

        # Normalize
        tfidf_norm = self._minmax_normalize(tfidf_scores)
        bm25_norm = self._minmax_normalize(bm25_scores)

        # Combine
        hybrid_scores = alpha * tfidf_norm + (1 - alpha) * bm25_norm
        ranked_idx = hybrid_scores.argsort()[::-1][:top_k]

        results = []
        for idx in ranked_idx:
            doc = self.documents[idx]
            results.append({
                "title": doc["title"],
                "url": doc["url"],
                "path": doc["path"],
                "text": doc["text"],
                "score": float(hybrid_scores[idx])
            })

        return results