import os
import logging
import numpy as np
from sentence_transformers import SentenceTransformer
from openai import OpenAI
from scipy.sparse import save_npz
import pickle
from datasets import load_from_disk
from sklearn.feature_extraction.text import TfidfVectorizer

# === Logging setup ===
logger = logging.getLogger("precompute")
logging.basicConfig(level=logging.INFO)

# === API key handling ===
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    OPENAI_API_KEY = input("Please enter your OpenAI API Key (set OPENAI_API_KEY environment variable for future runs): ")
    if not OPENAI_API_KEY:
        raise EnvironmentError("OPENAI_API_KEY must be provided either as an environment variable or input.")
openai_client = OpenAI(api_key=OPENAI_API_KEY)

# === Load CAP dataset ===
LOCAL_PATH = "./cap_dataset"  # Local path for testing
if not os.path.exists(LOCAL_PATH):
    raise FileNotFoundError(f"CAP dataset not found at {LOCAL_PATH}. Download it first.")
cap_dataset = load_from_disk(LOCAL_PATH)
cap_texts = [doc['text'] for doc in cap_dataset]
logger.info(f"Loaded {len(cap_texts)} CAP texts.")

# === TF-IDF Precomputation ===
if not (os.path.exists("./data/cap_tfidf.pkl") and os.path.exists("./data/cap_tfidf_matrix.npz")):
    logger.info("Precomputing TF-IDF...")
    tfidf = TfidfVectorizer(stop_words='english', max_features=100_000)
    tfidf_matrix = tfidf.fit_transform(cap_texts)
    os.makedirs("./data", exist_ok=True)
    with open("./data/cap_tfidf.pkl", 'wb') as f:
        pickle.dump(tfidf, f)
    save_npz("./data/cap_tfidf_matrix.npz", tfidf_matrix)
    logger.info("✅ Saved TF-IDF cache files.")
else:
    logger.info("TF-IDF cache files already exist, skipping.")

# === GTE Embeddings Precomputation ===
if not os.path.exists("./data/cap_gte.npy"):
    logger.info("Precomputing GTE embeddings...")
    encoder_gte = SentenceTransformer("Alibaba-NLP/gte-Qwen2-1.5B-instruct")
    embeddings_gte = encoder_gte.encode(cap_texts, normalize_embeddings=True)
    os.makedirs("./data", exist_ok=True)
    np.save("./data/cap_gte.npy", embeddings_gte)
    logger.info("✅ Saved GTE embeddings.")
else:
    logger.info("GTE embeddings cache file already exists, skipping.")

# === OpenAI Embeddings Precomputation ===
if not os.path.exists("./data/cap_openai.npy"):
    logger.info("Precomputing OpenAI embeddings...")
    def get_openai_embeddings(texts):
        chunk_size = 100  # Adjust based on average text length and token limit
        embeddings = []
        for i in range(0, len(texts), chunk_size):
            chunk = texts[i:i + chunk_size]
            response = openai_client.embeddings.create(
                model="text-embedding-3-large",
                input=chunk
            )
            embeddings.extend([item.embedding for item in response.data])
            logger.info(f"Processed chunk {i//chunk_size + 1} of {len(texts)//chunk_size + 1}")
            time.sleep(1)  # Rate limit buffer for Tier 5
        return np.array(embeddings)
    embeddings_openai = get_openai_embeddings(cap_texts)
    os.makedirs("./data", exist_ok=True)
    np.save("./data/cap_openai.npy", embeddings_openai)
    logger.info("✅ Saved OpenAI embeddings.")
else:
    logger.info("OpenAI embeddings cache file already exists, skipping.")

logger.info("✅ Precomputation completed. Cache files are ready for use.")