import os
from pinecone import Pinecone, ServerlessSpec
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])

if 'your-index' not in pc.list_indexes().names():
    pc.create_index(
        name='your-index',
        dimension=1536,
        metric='cosine',
        spec=ServerlessSpec(
            cloud='aws',
            region='us-east-1'
        )
    )

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

def upsert_texts(texts, ids):
    embeddings = []
    for text in texts:
        response = client.embeddings.create(
            model="text-embedding-ada-002",
            input=text
        )
        embeddings.append(response.data[0].embedding)

    pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])

    index = pc.Index("your-index")

    vectors = []
    for id, emb, text in zip(ids, embeddings, texts):
        vectors.append((id, emb, {"text": text}))

    index.upsert(vectors)

def query_text(query, top_k=5):
    pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
    index = pc.Index('your-index')

    embedding = client.embeddings.create(
        model="text-embedding-ada-002",
        input=query
    ).data[0].embedding

    results = index.query(vector=embedding, top_k=top_k, include_metadata=True)
    return results