import os
from typing import List
import numpy as np
from transformers import AutoTokenizer, AutoModel
from dotenv import load_dotenv
import torch
import ast  # <-- FIXED: Added missing import for sanitization

load_dotenv()

# --- SANITIZATION LAYER START ---
# Force-overwrite the environment HF_TOKEN to remove any byte-literal wrapping or literal newlines
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    try:
        if hf_token.startswith("b'") or hf_token.startswith('b"'):
            hf_token = ast.literal_eval(hf_token).decode("utf-8")
        else:
            hf_token = hf_token.strip()
        
        # Clean up escaped literal newlines (\n) causing the Illegal Header crash
        hf_token = hf_token.replace("\\n", "").replace("\n", "").strip()
        
        # Overwrite the active environment memory mapping globally
        os.environ["HF_TOKEN"] = hf_token
    except Exception:
        # Fallback to popping it out of memory if it can't be parsed safely
        os.environ.pop("HF_TOKEN", None)
# --- SANITIZATION LAYER END ---

# Set cache directory
os.environ["HF_HOME"] = os.getenv("HF_HOME", "/tmp/huggingface_cache")


class EmbeddingGenerator:
    def __init__(self, model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"):
        """Initialize embedding generator for multilingual text"""
        self.model_name = model_name
        
        # Pull the clean token string safely out of the cleaned environment memory
        token_str = os.getenv("HF_TOKEN")
        
        # Pass token explicitly even for local checks to override residual library bugs
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name, 
            local_files_only=False,
            token=token_str
        )
        self.model = AutoModel.from_pretrained(
            model_name,
            torch_dtype=torch.float32, # Use torch_dtype instead of dtype
            device_map="cpu",
            low_cpu_mem_usage=False,
            local_files_only=False,
            token=token_str
        )
        self.model = self.model.to("cpu")
        self.model.eval()

    def get_embedding(self, text: str) -> List[float]:
        """Generate embedding for text"""
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512,
            add_special_tokens=True
        )

        with torch.no_grad():
            outputs = self.model(**inputs)

        embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
        return embeddings.tolist()

    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for a batch of texts"""
        return [self.get_embedding(text) for text in texts]


def get_embedding_function():
    """Return the embedding function"""
    embedder = EmbeddingGenerator()
    return embedder.get_embedding


if __name__ == "__main__":
    embed_gen = EmbeddingGenerator()
    test_text = "Hello world, नमस्ते दुनिया"
    embedding = embed_gen.get_embedding(test_text)
    print(f"Embedding length: {len(embedding)}")
    print(f"First 10 values: {embedding[:10]}")