import os
import numpy as np
from flask import Flask, request, jsonify, send_from_directory
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# 지식 베이스 (day1에서 쓴 것과 동일)
documents = [
    "Python은 강력한 프로그래밍 언어입니다.",
    "OpenAI는 AI 연구를 선도하는 기업입니다.",
    "FAISS는 벡터 검색을 위한 라이브러리입니다.",
]


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)


# 서버 시작 시 문서들을 미리 임베딩해 둔다
doc_embeddings = np.array([get_embedding(doc) for doc in documents], dtype=np.float32)


def rag_answer(question):
    # ① 질문 임베딩 → 코사인 유사도가 가장 큰 문서 검색
    q = get_embedding(question)
    sims = doc_embeddings @ q / (np.linalg.norm(doc_embeddings, axis=1) * np.linalg.norm(q))
    retrieved_doc = documents[int(np.argmax(sims))]

    # ② 검색된 문서를 근거로 LLM이 답변
    prompt = f"참고 정보: {retrieved_doc}\n\n질문: {question}\n답변:"
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "당신은 제공된 참고 정보를 바탕으로 답합니다."},
            {"role": "user", "content": prompt},
        ],
    )
    return retrieved_doc, response.choices[0].message.content


app = Flask(__name__)  # 정적 파일은 ./static 폴더에서 자동 서빙 (/static/...)


@app.route("/")
def index():
    return send_from_directory(app.static_folder, "rag.html")


@app.route("/api/rag", methods=["POST"])
def rag():
    question = request.json["message"]
    retrieved_doc, answer = rag_answer(question)
    return jsonify({"retrieved": retrieved_doc, "answer": answer})


if __name__ == "__main__":
    app.run(debug=True, port=5000)
