import chromadb
import pandas as pd
import logging

# --- 설정 (필요시 수정) ---
CHROMA_HOST = "localhost"  # ChromaDB가 실행 중인 서버 주소
CHROMA_PORT = 8001         # ChromaDB 포트
COLLECTION_NAME = "segue_saved_qna13" # 데이터를 추출할 컬렉션 이름
OUTPUT_EXCEL_FILE = "segue_saved_qna_master.xlsx" # 저장할 엑셀 파일 이름
# -----------------------------

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

def export_collection_to_excel():
    """
    ChromaDB 컬렉션의 모든 데이터를 읽어 Excel 파일로 저장합니다.
    """
    try:
        # 1. ChromaDB 클라이언트에 연결합니다.
        logging.info(f"ChromaDB 서버({CHROMA_HOST}:{CHROMA_PORT})에 연결을 시도합니다...")
        client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)

        # 2. 지정된 컬렉션을 가져옵니다.
        logging.info(f"'{COLLECTION_NAME}' 컬렉션을 불러옵니다...")
        collection = client.get_collection(name=COLLECTION_NAME)

        # 3. 컬렉션의 모든 데이터를 가져옵니다.
        # .get()에 아무 인자도 주지 않으면 컬렉션의 모든 데이터를 반환합니다.
        logging.info("컬렉션에서 모든 데이터를 추출합니다...")
        data = collection.get(
            include=["documents", "metadatas"] # 질문(documents)과 답변(metadatas)을 가져옵니다.
        )

        # 4. 추출한 데이터를 엑셀에 저장할 형태로 가공합니다.
        qna_list = []
        documents = data.get('documents', [])
        metadatas = data.get('metadatas', [])

        if not documents:
            logging.warning("컬렉션에 데이터가 없습니다. 빈 엑셀 파일이 생성됩니다.")

        for i, question in enumerate(documents):
            # metadata가 있고, 그 안에 'answer' 키가 있는지 확인합니다.
            answer = metadatas[i].get('answer', 'N/A') if metadatas and len(metadatas) > i else 'N/A'
            qna_list.append({
                "question": question,
                "answer": answer
            })

        logging.info(f"총 {len(qna_list)}개의 Q&A 데이터를 성공적으로 추출했습니다.")

        # 5. Pandas DataFrame으로 변환 후 Excel 파일로 저장합니다.
        if qna_list:
            df = pd.DataFrame(qna_list)

            logging.info(f"'{OUTPUT_EXCEL_FILE}' 파일로 저장합니다...")
            df.to_excel(OUTPUT_EXCEL_FILE, index=False, engine='openpyxl')
            logging.info("🎉 작업 완료! 엑셀 파일이 성공적으로 생성되었습니다.")
        else:
            # 데이터가 없을 경우 빈 엑셀 파일을 만들어줍니다.
            pd.DataFrame(columns=["question", "answer"]).to_excel(OUTPUT_EXCEL_FILE, index=False)
            logging.info("데이터가 없어 빈 엑셀 파일을 생성했습니다.")

    except Exception as e:
        logging.error(f"오류가 발생했습니다: {e}")
        logging.error("ChromaDB 서버가 실행 중인지, 컬렉션 이름이 정확한지 확인해주세요.")

if __name__ == "__main__":
    export_collection_to_excel()