import json
from typing import List, Union

import chromadb
from chromadb.utils import embedding_functions

# === 경로 설정 (모두 로컬) ===
EMBEDDING_MODEL_PATH = "./models/ko-sroberta-multitask"

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

# 임베딩 함수 설정
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=EMBEDDING_MODEL_PATH,  # 로컬 폴더 경로 가능
    device="cpu",
    normalize_embeddings=True
)

# 컬렉션 생성 또는 가져오기
collection = chroma_client.get_or_create_collection(
    name="orgchart",
    embedding_function=embedding_fn,
    metadata={"hnsw:space": "cosine"}
)


def init(sessionId: str, data: List[Union[str, dict]]):
    """
    데이터를 벡터 DB에 초기화(저장)합니다.

    Args:
        sessionId (str): 세션 ID
        data (List[Union[str, dict]]): 저장할 데이터 리스트 (문자열 또는 딕셔너리)
    """
    print(f'{sessionId} init start')

    # 문서 ID 생성
    doc_ids = [f"{sessionId}_{i}" for i in range(len(data))]

    # 데이터 처리: 문자열이면 그대로, 객체면 JSON 문자열로 변환
    documents = []

    doc_list = []
    for q in data:
        # 각 필드를 안전하게 추출 (None 방어)
        name = q.get('name') or ""
        dept = q.get('deptNm') or ""
        grade = q.get('gradeNm') or ""
        position = q.get('ptsnNm') or ""
        office_phone = q.get('ofcePhn') or ""
        mobile_phone = q.get('mblPhn') or ""
        chief_name = q.get('chiefNm') or ""
        state_code = q.get('state') or ""

        # 상태 코드 한글화
        state_map = {'C': '재직', 'T': '퇴사', 'H': '휴직'}
        status = state_map.get(state_code, "정보없음")

        # [핵심] 검색 엔진이 좋아할만한 서술형 문장 생성
        # 부서명과 이름을 앞부분에 배치하여 가중치 유도
        if name == '':
            doc = (
                f"부서: {dept}. "
                f"해당 {dept}의 부서장은 {chief_name}입니다."
            )
        else:
            doc = (
                f"부서: {dept}. 이름: {name}. {dept} 소속의 {name} {grade}입니다. "
                f"직위는 {position}이며 현재 {status} 중입니다. "
                f"사내 전화번호(사선)는 {office_phone}입니다."
            )
        doc_list.append(doc)

    # 벡터 DB에 추가
    collection.add(
        documents=doc_list,
        ids=doc_ids,
        metadatas=[{"sessionId": sessionId, 'deptCd': d.get('deptCd') or "", 'gradeCd': d.get('gradeCd') or ""} for d in data]
    )

    print(f'{sessionId} init end')