import os

import chromadb
import pandas as pd
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from dotenv import load_dotenv

load_dotenv()

# === CONFIG ===
CSV_PATH = "master_jtbd_comments.csv"
COLLECTION_NAME = "test"
API_KEY = os.environ["CHROMA_API_KEY"]
CHROMA_CLOUD_URL = "https://api.trychroma.com"  # default endpoint

# === 1. Load CSV ===
df = pd.read_csv(CSV_PATH)
print(f"✅ Loaded {len(df)} rows from {CSV_PATH}")

# Combine text columns into one field
TEXT_COLUMNS = ["comment", "situation", "struggle", "outcome"]

def combine_text(row):
    parts = [str(row[c]) for c in TEXT_COLUMNS if c in row and pd.notna(row[c])]
    return " | ".join(parts)

df["combined_text"] = df.apply(combine_text, axis=1)

# === 2. Initialize Chroma Cloud client ===
client = chromadb.CloudClient(
  api_key=API_KEY,
  tenant=os.environ["CHROMA_TENANT"],
  database='test'
)

# === 3. Create / get collection ===
collection = client.get_or_create_collection(
    name=COLLECTION_NAME,
    metadata={"description": "JTBD comments + situations/struggles/outcomes"}
)

# === 4. Embed text with OpenAI ===
embedder = OpenAIEmbeddingFunction(api_key=os.getenv("OPENAI_API_KEY"))

# === 5. Upload data ===
batch_size = 100
for i in range(0, len(df), batch_size):
    batch = df.iloc[i:i+batch_size]
    collection.add(
        documents=batch["combined_text"].tolist(),
        metadatas=batch[["jtbd", "source_file"]].to_dict(orient="records"),
        ids=[str(idx) for idx in batch.index],
        embeddings=embedder(batch["combined_text"].tolist())
    )
    print(f"→ Uploaded {i + len(batch)}/{len(df)} rows")

print("🎉 All data uploaded to Chroma Cloud collection:", COLLECTION_NAME)