import chromadb
from chromadb.utils import embedding_functions

# 2. 벡터 DB 설정
persist_directory = "./chroma_db"
chroma_client = chromadb.PersistentClient(path=persist_directory)

# ✅ Chroma 전용 임베딩 함수 사용 (오류 방지)
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# 컬렉션 생성
collection = chroma_client.get_or_create_collection(
    name="manuals",
    embedding_function=sentence_transformer_ef  # ← 여기가 핵심!
)


def init(job: str):
    """
    직무별 매뉴얼 데이터를 벡터 DB에 초기화합니다.

    Args:
        job (str): 직무 코드 (예: 'FI', 'HR')
    """
    print(f'{job} init start')
    # 1. 문서 준비 (실제로는 PDF/Word 등에서 추출)
    manuals = [
        "지출 결의서는 사용 목적, 금액, 일자, 증빙 서류를 반드시 첨부하여 전자 결재 시스템에 등록해야 합니다.",
        "월말 마감은 매월 25일부터 시작되며, 모든 부서는 28일까지 비용 집행 내역을 최종 확정해야 합니다.",
        "외화 송금은 반드시 외환관리부의 사전 승인을 받은 후 금융팀을 통해 진행되어야 하며, 계약서 사본을 첨부해야 합니다.",
        "세금계산서는 발행일로부터 10일 이내에 ERP에 등록되지 않으면 비용 처리가 불가합니다.",
        "장기자산(차량, 사무기기 등)은 매년 1월에 정기 감가상각 점검을 받아야 하며, 자산관리부서가 이를 주관합니다.",
        "현금 보관은 원칙적으로 금지되며, 불가피한 경우는 금고 보관 후 당일 중 재무팀에 입금 처리해야 합니다.",
        "연말 정산 대상 직원은 매년 12월 10일까지 개인 소득공제 자료를 인사 시스템에 제출해야 합니다.",
        "예산 초과 지출은 사전에 재무부와 협의 후 예산 조정 승인을 받아야 하며, 미승인 시 결재가 거부됩니다.",
        "재무 제표 초안은 분기 마감 후 5영업일 이내에 감사법인에 제출되어야 하며, 최종 승인은 CFO가 담당합니다.",
    ]

    # 문서 ID 생성 및 추가
    doc_ids = [f"DOC_{job}_{i}" for i in range(len(manuals))]
    
    collection.add(
        documents=manuals,
        ids=doc_ids,
        metadatas=[{"source": "fi_manual_v1.pdf", "version": 1.0, "dept": job} for _ in doc_ids]
    )

    print(f'{job} init end')


if __name__ == "__main__":
    # FI : 재무 HR : 인사
    init(job="FI")