# ============================================================
# [실습] 3단계: 문서 여러 개 + 출처 표시  (= 02 + source)
#
# 02 에서 바뀌는 건 두 가지뿐 (chain 은 그대로):
#   · 업로드할 때 각 청크에 source(파일명) 를 심고, 인덱싱된 파일 목록을 반환
#   · 질의 결과에서 source 를 모아 답과 함께 반환 (가장 관련 높은 문서가 맨 앞)
#
# 실행:  python 03_multi_docs_실습.py   →   http://localhost:5000
#
# 구현 항목:
#  1) /upload : 각 문서에 metadata["source"]=파일명 심기, INDEXED 목록 갱신·반환
#  2) /ask    : k=2 검색 → context → chain.invoke, sources(관련도 순·중복 제거) 반환
# ============================================================

import os
from flask import Flask, request, jsonify, send_from_directory
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from dotenv import load_dotenv

load_dotenv()
app = Flask(__name__)
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
prompt = ChatPromptTemplate.from_messages([
    ("system", "아래 [문서]에 있는 내용만 근거로 한국어로 간결히 답하세요. 문서에 없으면 모른다고 답하세요."),
    ("human", "[문서]\n{context}\n\n[질문] {question}"),
])
chain = prompt | llm | StrOutputParser()

store = InMemoryVectorStore(OpenAIEmbeddings(model="text-embedding-3-small"))
INDEXED = []


@app.get("/")
def index():
    return send_from_directory(app.static_folder, "03.html")


@app.post("/upload")
def upload():
    file = request.files.get("file")
    if not file or not file.filename:
        return jsonify({"error": "파일이 없습니다"}), 400
    path = os.path.join(UPLOAD_DIR, file.filename)
    file.save(path)
    ext = file.filename.rsplit(".", 1)[-1].lower()
    loader = PyPDFLoader(path) if ext == "pdf" else TextLoader(path, encoding="utf-8")
    docs = loader.load()
    # 1) 여기에 코드를 작성하세요 (각 doc.metadata["source"]=파일명 → 청킹 → add → INDEXED 갱신)
    return jsonify({"message": f"'{file.filename}' 업로드 완료", "files": INDEXED})


@app.post("/ask")
def ask():
    question = request.json["question"]
    hits = store.similarity_search(question, k=2)
    if not hits:
        return jsonify({"answer": "아직 업로드된 문서가 없습니다.", "sources": []})
    # 2) 여기에 코드를 작성하세요
    #    context(출처 포함) → answer = chain.invoke(...) → sources(관련도 순, 중복 제거)
    pass


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