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=8001,
    settings=Settings(chroma_api_impl="rest")
)

# ② 기존 컬렉션 가져오기 (excel_honbun과 같은 DB를 사용)
collection_name = "excel_honbun"  # 기존 collection에 합쳐서 등록
collection = client.get_collection(collection_name)

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

# ④ 데이터 전처리 (모든 컬럼 정보 통합)
texts = []
for idx, row in df.iterrows():
    注文者 = str(row.get('注文者', '')).strip()
    元請下請 = str(row.get('元請又は下請', '')).strip()
    JV区分 = str(row.get('JVの別', '')).strip()
    工事名 = str(row.get('工事名', '')).strip()
    工事現場 = str(row.get('工事現場の名称・都市等', '')).strip()
    請負代金 = str(row.get('請負代金の額', '')).strip()
    請負単位 = str(row.get('請負代金の額.1', '')).strip()  # 千円 등 단위
    着工年月 = str(row.get('着工年月', '')).strip()
    完成年月 = str(row.get('完成又は完成予定年月', '')).strip()

    combined_text = (
        f"[施工実績]\n"
        f"注文者: {注文者}\n"
        f"元請又は下請: {元請下請}\n"
        f"JVの別: {JV区分}\n"
        f"工事名: {工事名}\n"
        f"工事現場: {工事現場}\n"
        f"請負代金: {請負代金} {請負単位}\n"
        f"着工年月: {着工年月}\n"
        f"完成又は完成予定年月: {完成年月}"
    )

    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=[f"construction_{idx}"],  # id 충돌 방지 위해 접두어 사용
            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}")