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 서버로 연결 (1.x 서버 반드시 실행중이어야 함)
client = chromadb.HttpClient(
    host="localhost",
    port=8000,
    settings=Settings(chroma_api_impl="rest")
)

# ② 컬렉션 생성 또는 가져오기 (컬렉션명은 영문으로)
collection_name = "excel_honbun"
if collection_name not in [c.name for c in client.list_collections()]:
    collection = client.create_collection(collection_name)
else:
    collection = client.get_collection(collection_name)

# ③ 엑셀 파일 읽기
df = pd.read_excel('excel_honbun.xlsx', sheet_name='honbun')

# ④ 데이터 전처리 (章タイトル + 章概要 + 章本文를 묶어서 하나의 문장으로 변환)
texts = []
for idx, row in df.iterrows():
    chapter = str(row.get('章', '')).strip()
    title = str(row.get('章タイトル', '')).strip()
    summary = str(row.get('章概要', '')).strip()
    content = str(row.get('章本文', '')).strip()

    combined_text = f"章: {chapter}\n章タイトル: {title}\n章概要: {summary}\n章本文: {content}"
    texts.append(combined_text)

# ⑤ 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=[str(idx)],
            documents=[text],
            embeddings=[embedding]
        )

        print(f"[{idx+1}/{len(texts)}] 등록 완료")

        time.sleep(0.5)  # (속도조절: API rate limit 방지)

    except Exception as e:
        print(f"[{idx}] 에러 발생: {e}")