import numpy as np
import faiss
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
client = OpenAI()


def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-ada-002")
    return np.array(response.data[0].embedding, dtype=np.float32)


documents = [
    "Python은 강력한 프로그래밍 언어입니다.",
    "OpenAI는 AI 연구를 선도하는 기업입니다.",
    "FAISS는 벡터 검색을 위한 라이브러리입니다.",
]
doc_embeddings = np.array([get_embedding(doc) for doc in documents], dtype=np.float32)

# 코사인 유사도 = '정규화한 벡터의 내적'
faiss.normalize_L2(doc_embeddings)                  # ① 벡터를 길이 1로 정규화
index = faiss.IndexFlatIP(doc_embeddings.shape[1])  # ② 내적(IP) 기반 인덱스
index.add(doc_embeddings)                           # ③ 문서 벡터 저장

# 질문도 같은 모델로 임베딩 + 똑같이 정규화 → 검색 결과값이 곧 코사인 유사도
query = "검색 라이브러리가 뭐야?"
query_vec = np.array([get_embedding(query)], dtype=np.float32)
faiss.normalize_L2(query_vec)

scores, indices = index.search(query_vec, k=2)      # 상위 2개 검색
for score, idx in zip(scores[0], indices[0]):
    print(f"문서 {idx}: 유사도 {score:.4f} -> {documents[idx]}")
