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

# --- 📜 설정: 이 부분을 자신의 환경에 맞게 수정하세요 ---
EXCEL_FILE_PATH = "segue_saved_qna_master.xlsx"      # 마스터 엑셀 파일 경로
COLLECTION_NAME = "segue_saved_qna13"          # 사용할 ChromaDB 컬렉션 이름
EMBEDDING_MODEL = "text-embedding-3-small"     # 임베딩 모델

# ChromaDB 서버 정보
CHROMA_HOST = "localhost"
CHROMA_PORT = 8001

# OpenAI API 클라이언트 준비
# config.py에서 직접 불러온 유효한 API 키를 명시적으로 사용합니다.
client_openai = OpenAI(api_key=openai_api_key)
# ----------------------------------------------------

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def sync_excel_to_chromadb():
    """
    엑셀 파일을 기준으로 ChromaDB 컬렉션을 완전히 동기화(재생성)합니다.
    """
    try:
        # 1. 엑셀 파일 읽기
        logging.info(f"'{EXCEL_FILE_PATH}' 파일을 읽습니다...")
        df = pd.read_excel(EXCEL_FILE_PATH)
        # question이 비어있는 행은 무시
        df.dropna(subset=['question'], inplace=True)
        df['answer'] = df['answer'].fillna('') # answer가 비어있으면 빈 문자열로 처리

        questions = df['question'].astype(str).str.strip().tolist()
        answers = df['answer'].astype(str).str.strip().tolist()

        if not questions:
            logging.warning("엑셀에 데이터가 없습니다. DB를 비웁니다.")
        else:
            logging.info(f"엑셀에서 {len(questions)}개의 Q&A 데이터를 찾았습니다.")

        # 2. 모든 질문을 '한 번에' 임베딩 (매우 효율적)
        logging.info(f"총 {len(questions)}개의 질문을 한번에 임베딩합니다. 잠시 기다려주세요...")
        response = client_openai.embeddings.create(
            model=EMBEDDING_MODEL,
            input=questions
        )
        embeddings = [data.embedding for data in response.data]
        logging.info("임베딩 생성 완료.")

        # 3. ChromaDB 클라이언트 연결
        logging.info(f"ChromaDB 서버({CHROMA_HOST}:{CHROMA_PORT})에 연결합니다...")
        client_chroma = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)

        # 4. 기존 컬렉션을 '삭제'하고 새로 생성 (가장 안전한 방법)
        logging.info(f"기존 '{COLLECTION_NAME}' 컬렉션을 삭제하고 새로 생성합니다...")
        try:
            client_chroma.delete_collection(name=COLLECTION_NAME)
        except Exception:
            logging.info(f"기존 '{COLLECTION_NAME}' 컬렉션이 없어 새로 생성만 진행합니다.")

        collection = client_chroma.create_collection(name=COLLECTION_NAME)

        # 5. 모든 데이터를 '한 번에' DB에 저장
        logging.info(f"새로운 데이터 {len(questions)}개를 DB에 저장합니다...")
        collection.add(
            ids=[f"qna_{i}" for i in range(len(questions))], # ID는 행 순서대로 부여
            documents=questions,
            embeddings=embeddings,
            metadatas=[{'answer': ans} for ans in answers]
        )

        logging.info(f"🎉 작업 완료! '{COLLECTION_NAME}' 컬렉션이 엑셀 파일과 완벽하게 동기화되었습니다.")

    except FileNotFoundError:
        logging.error(f"🚨 파일을 찾을 수 없습니다: '{EXCEL_FILE_PATH}'")
    except Exception as e:
        logging.error(f"🚨 작업 중 오류가 발생했습니다: {e}", exc_info=True)

# 이 스크립트를 직접 실행할 때 sync_excel_to_chromadb 함수를 호출합니다.
if __name__ == "__main__":
    sync_excel_to_chromadb()