from sentence_transformers import SentenceTransformer
from langchain_core.embeddings import Embeddings
import torch
import re
import emoji
import unicodedata
import nltk
from nltk.corpus import stopwords

nltk.download('stopwords')
stop_words = set(stopwords.words('english'))

device = "cuda" if torch.cuda.is_available() else "cpu"
text_model = SentenceTransformer("intfloat/e5-base", device=device)


def preprocess_text(text):
    text = emoji.replace_emoji(text, replace='')
    text = ''.join(c for c in text if unicodedata.category(c)[0] != 'C')
    text = text.lower()
    words = re.findall(r'\b[a-z]+\b', text)
    filtered = [w for w in words if w not in stop_words]

    return " ".join(filtered)


def get_text_embedding(text):
    clean_text = preprocess_text(text)
    emb = text_model.encode(clean_text, convert_to_numpy=True, normalize_embeddings=True)
    return emb.tolist()


class TextEmbeddings(Embeddings):
    def embed_documents(self, texts):
        return [get_text_embedding(text) for text in texts]

    def embed_query(self, text):
        return get_text_embedding(text)