import pandas as pd
import openai
from openai import OpenAI
import chromadb
import time
from config import logger, openai_api_key
from chromadb.config import Settings

# ① ChromaDB HTTP 서버 연결
client = chromadb.HttpClient(
    host="localhost",
    port=8001,
    settings=Settings(chroma_api_impl="rest")
)

# ② 컬렉션명 정의 (신규 생성)
collection_name = "segue_profile_split6"
try:
    collection = client.get_collection(collection_name)
except:
    collection = client.create_collection(collection_name)

# ③ 엑셀 데이터 읽기
df = pd.read_excel("segue_profile_split6.xlsx")

# ④ 데이터 전처리 및 임베딩 준비
texts = []
metadatas = []
ids = []

for idx, row in df.iterrows():
    content = str(row.get("content", "")).strip()
    page_type = str(row.get("page_type", "")).strip()
    url = str(row.get("url", "")).strip()
    year = row.get("year", None)

    # year가 NaN이면 제외
    year_meta = int(year) if pd.notna(year) else None

    texts.append(content)
    metadatas.append({
        "page_type": page_type,
        "url": url,
        **({"year": year_meta} if year_meta is not None else {})
    })
    ids.append(f"page_{idx}")

# ⑤ OpenAI API 준비
openai.api_key = openai_api_key
client_openai = OpenAI(api_key=openai_api_key)

# ⑥ 임베딩 및 ChromaDB 등록
for idx, text in enumerate(texts):
    try:
        response = client_openai.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        embedding = response.data[0].embedding

        collection.add(
            ids=[ids[idx]],
            documents=[text],
            embeddings=[embedding],
            metadatas=[metadatas[idx]]
        )

        print(f"[{idx + 1}/{len(texts)}] 登録完了")
        time.sleep(0.5)

    except Exception as e:
        print(f"[{idx}] エラー発生: {e}")