from langchain_community.vectorstores import FAISS
from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from google import genai
import os

# Make sure your environment variable GOOGLE_API_KEY is set
API_KEY = os.getenv("GOOGLE_API_KEY")
if not API_KEY:
    raise ValueError("Missing GOOGLE_API_KEY environment variable!")

# Initialize client with API key
client = genai.Client(api_key=API_KEY)


class GeminiEmbeddings(Embeddings):
    """LangChain wrapper for Google Gemini embeddings"""
    
    def embed_documents(self, texts):
        if not texts:
            return []
        response = client.models.embed_content(
            model="gemini-embedding-001",
            contents=texts
        )
        # Each response.embeddings[i].values is a list of floats
        return [e.values for e in response.embeddings]

    def embed_query(self, text):
        response = client.models.embed_content(
            model="gemini-embedding-001",
            contents=[text]
        )
        return response.embeddings[0].values


def create_vector_store(texts):
    docs = [Document(page_content=t) for t in texts if t.strip()]
    if not docs:
        return None

    embeddings = GeminiEmbeddings()
    vectorstore = FAISS.from_texts(
        texts=[d.page_content for d in docs],
        embedding=embeddings
    )
    return vectorstore