"""
개선된 벡터 스토어 모듈 - Milvus 설정 최적화
"""
from typing import List, Dict, Any, Optional
import uuid
from langchain.schema import Document

# 벡터 스토어 임포트
try:
    # 최신 버전 임포트
    from langchain_milvus import Milvus
    from langchain_community.vectorstores import FAISS
    from langchain_huggingface import HuggingFaceEmbeddings
    MODERN_IMPORTS = True
    print("최신 langchain 패키지 임포트 성공")
except ImportError:
    # 이전 버전 임포트
    from langchain_community.vectorstores import Milvus, FAISS
    from langchain_community.embeddings import HuggingFaceEmbeddings
    MODERN_IMPORTS = False
    print("레거시 langchain_community 패키지 사용")

from config import MILVUS_HOST, MILVUS_PORT, MILVUS_COLLECTION, EMBEDDING_MODEL

class VectorStore:
    def __init__(self, use_milvus: bool = True):
        """
        벡터 스토어 초기화

        Args:
            use_milvus: Milvus 사용 여부 (False이면 FAISS 사용)
        """
        self.use_milvus = use_milvus

        # 임베딩 모델 설정
        print(f"임베딩 모델 로드 중: {EMBEDDING_MODEL}")
        model_kwargs = {
            "device": "cpu",
            "trust_remote_code": True  # 원격 코드 실행 허용 (필수)
        }
        encode_kwargs = {"normalize_embeddings": True}

        self.embeddings = HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs
        )
        self.vector_store = None

        print(f"임베딩 모델 초기화 완료: {EMBEDDING_MODEL}")

    def init_milvus(self) -> Milvus:
        """
        Milvus 벡터 스토어 초기화

        Returns:
            Milvus 벡터 스토어 인스턴스
        """
        connection_args = {
            "host": MILVUS_HOST,
            "port": MILVUS_PORT,
        }

        # 벡터 검색 인덱스 파라미터 (FLAT 인덱스 및 코사인 유사도 메트릭)
        index_params = {
            "index_type": "FLAT",  # 정확도 우선 FLAT 인덱스
            "metric_type": "COSINE",  # 코사인 유사도 (정규화된 벡터에 적합)
            "params": {}  # FLAT 인덱스에는 추가 파라미터 없음
        }

        return Milvus(
            embedding_function=self.embeddings,
            collection_name=MILVUS_COLLECTION,
            connection_args=connection_args,
            index_params=index_params
        )

    def init_faiss(self) -> FAISS:
        """
        FAISS 벡터 스토어 초기화 (로컬 대체용)

        Returns:
            FAISS 벡터 스토어 인스턴스
        """
        return FAISS.from_documents([], self.embeddings)

    def create_or_load(self, documents: Optional[List[Document]] = None) -> Any:
        """
        벡터 스토어 생성 또는 로드

        Args:
            documents: 저장할 문서 리스트 (None이면 빈 스토어 생성)

        Returns:
            벡터 스토어 인스턴스
        """
        if self.use_milvus:
            if documents:
                # 문서가 제공된 경우 새 컬렉션 생성
                try:
                    # 연결 설정
                    connection_args = {
                        "host": MILVUS_HOST,
                        "port": MILVUS_PORT,
                    }

                    # 검색 인덱스 설정
                    index_params = {
                        "index_type": "FLAT",  # 정확도 우선
                        "metric_type": "COSINE",  # 코사인 유사도
                        "params": {}
                    }

                    print(f"Milvus 컬렉션 생성: {MILVUS_COLLECTION} (기존 컬렉션 삭제)")

                    # 문서로부터 Milvus 컬렉션 생성
                    self.vector_store = Milvus.from_documents(
                        documents=documents,
                        embedding=self.embeddings,
                        collection_name=MILVUS_COLLECTION,
                        connection_args=connection_args,
                        index_params=index_params,
                        drop_old=True  # 기존 컬렉션 삭제 (재구축)
                    )

                    print(f"Milvus 컬렉션 생성 완료: {len(documents)}개 문서 인덱싱됨")

                except Exception as e:
                    print(f"Milvus 컬렉션 생성 실패: {e}")
                    # 대체 방안으로 FAISS 사용
                    print("대체 방안으로 FAISS 사용")
                    self.use_milvus = False
                    self.vector_store = FAISS.from_documents(documents, self.embeddings)
            else:
                # 기존 컬렉션 로드
                try:
                    self.vector_store = self.init_milvus()
                except Exception as e:
                    print(f"Milvus 컬렉션 로드 실패: {e}")
                    # 대체 방안으로 FAISS 사용
                    print("대체 방안으로 FAISS 사용")
                    self.use_milvus = False
                    self.vector_store = self.init_faiss()
        else:
            # FAISS 사용
            if documents:
                print(f"FAISS 인덱스 생성: {len(documents)}개 문서")
                self.vector_store = FAISS.from_documents(documents, self.embeddings)
                print("FAISS 인덱스 생성 완료")
            else:
                self.vector_store = self.init_faiss()
                print("빈 FAISS 인덱스 초기화 완료")

        return self.vector_store

    def add_documents(self, documents: List[Document]) -> None:
        """
        벡터 스토어에 문서 추가

        Args:
            documents: 추가할 문서 리스트
        """
        if self.vector_store is None:
            self.create_or_load(documents)
        else:
            if self.use_milvus:
                self.vector_store.add_documents(documents)
            else:
                self.vector_store.add_documents(documents)

    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        """
        벡터 유사도 검색 수행

        Args:
            query: 검색 쿼리
            k: 반환할 결과 수

        Returns:
            유사도가 높은 문서 리스트
        """
        if self.vector_store is None:
            raise ValueError("벡터 스토어가 초기화되지 않았습니다.")

        print(f"검색 쿼리: '{query}', 상위 {k}개 결과 요청")
        results = self.vector_store.similarity_search(query, k=k)
        print(f"검색 완료: {len(results)}개 결과 찾음")

        return results

    def save_local(self, path: str = "faiss_index") -> None:
        """
        FAISS 인덱스 로컬 저장 (Milvus 사용 안 할 경우)

        Args:
            path: 저장 경로
        """
        if not self.use_milvus and self.vector_store is not None:
            self.vector_store.save_local(path)
            print(f"FAISS 인덱스 로컬 저장 완료: {path}")

    """
    FAISS 역직렬화 허용 설정이, 포함된 벡터 스토어 코드
    """

    # vector_store.py 파일에서 load_local 메서드 수정

    def load_local(self, path: str = "faiss_index") -> None:
        """
        FAISS 인덱스 로컬 로드 (Milvus 사용 안 할 경우)

        Args:
            path: 로드할 인덱스 경로
        """
        if not self.use_milvus:
            try:
                print(f"FAISS 인덱스 로드 중: {path}")

                # 역직렬화 허용 옵션 추가 (보안 경고 확인 필요)
                self.vector_store = FAISS.load_local(
                    path,
                    self.embeddings,
                    allow_dangerous_deserialization=True  # 역직렬화 허용
                )
                print(f"FAISS 인덱스 로드 완료: {path}")
            except Exception as e:
                print(f"FAISS 인덱스 로드 실패: {e}")

                # 오류 세부 정보 출력
                import traceback
                traceback.print_exc()

                # 새 인덱스 초기화
                self.vector_store = self.init_faiss()
                print("새 FAISS 인덱스 초기화됨")