import os
import sys
import numpy as np
import pandas as pd
from tqdm import tqdm
import torch
from sentence_transformers import SentenceTransformer
import streamlit as st
from dotenv import load_dotenv

load_dotenv()

def load_and_parse_data(data_path, emb_path):
    # FIX: Removed memory_map=True to prevent file locking
    docs = pd.read_parquet(data_path) 
    chunks = list(docs["data"])

    vecs = pd.read_parquet(emb_path)
    embs = vecs = np.stack(vecs['embedding'].values)
    return chunks, embs

@st.cache_resource(show_spinner=False)
def get_model():
    # 1. Determine where we are (Local vs Frozen/Exe)
    if getattr(sys, 'frozen', False):
        base_dir = sys._MEIPASS
    else:
        base_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 2. Check for the local folder
    local_model_path = os.path.join(base_dir, "models", "multilingual-e5-large")
    
    # 3. Decision Logic: Use Local if it exists, otherwise use Cloud
    if os.path.exists(local_model_path):
        print(f"📂 [System] Found local model at: {local_model_path}")
        model_source = local_model_path
    else:
        print(f"☁️ [System] Local model not found. Downloading from Hugging Face...")
        model_source = "intfloat/multilingual-e5-large" 

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"🔄 [System] Loading Embedding Model ({model_source}) on {device}...")
    
    try:
        model = SentenceTransformer(model_source, device=device)
        print("✅ [System] Model loaded.")
        return model
    except Exception as e:
        print(f"❌ [System] Failed to load model: {e}")
        raise e

def bge_embed(texts, batch_size=256):
    """
    Standardizes text for E5 models. 
    E5 requires 'query: ' prefix for search and 'passage: ' for storage.
    """
    model = get_model()
    embeddings = []
    
    if isinstance(texts, str):
        texts = [texts]

    # --- E5 SPECIFIC UPDATE ---
    # If the input is a single string (a query), we add the 'query: ' prefix.
    # If you are using this function to embed your database (passages), 
    # you should change this prefix to 'passage: '.
    processed_texts = []
    for t in texts:
        if not t.startswith("query: ") and not t.startswith("passage: "):
            # Defaulting to query prefix for the search function
            processed_texts.append(f"query: {t}")
        else:
            processed_texts.append(t)
    # ---------------------------
        
    for i in range(0, len(processed_texts), batch_size):
        batch = processed_texts[i:i + batch_size]
        batch_emb = model.encode(
            batch, 
            batch_size=len(batch), 
            show_progress_bar=False, 
            normalize_embeddings=True 
        )
        embeddings.extend(batch_emb)
        
    return embeddings