# ============================================================
# [실습] 4단계: 문서 삭제  (= 03 + 삭제)
#
# INDEXED(이름만) 대신 DOC_IDS(이름→청크 id) 로 바꿔, 삭제 시 그 문서의 청크 id 를 지운다.
#
# 실행:  python 04_delete_실습.py   →   http://localhost:5000
#
# 구현 항목:
#  1) /upload : store.add_documents 가 돌려준 id 를 DOC_IDS[파일명] 에 모은다
#  2) DELETE /files/<name> : DOC_IDS 에서 꺼낸 id 들을 store.delete(ids=...) 로 제거
#  3) /ask    : 03 과 동일 (검색 → context → chain → 답 + 출처)
# ============================================================

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"))
DOC_IDS = {}   # {파일명: [청크 id]}


@app.get("/")
def index():
    return send_from_directory(app.static_folder, "04.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()
    for d in docs:
        d.metadata["source"] = file.filename
    chunks = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50).split_documents(docs)
    # 1) 여기에 코드를 작성하세요 (ids = store.add_documents(chunks); DOC_IDS[파일명] 에 누적)
    return jsonify({"message": f"'{file.filename}' 업로드 완료", "files": list(DOC_IDS)})


@app.delete("/files/<name>")
def delete_file(name):
    # 2) 여기에 코드를 작성하세요
    #    DOC_IDS.pop(name) → 없으면 404 → store.delete(ids=...) → 남은 목록 반환
    pass


@app.post("/ask")
def ask():
    question = request.json["question"]
    hits = store.similarity_search(question, k=2)
    if not hits:
        return jsonify({"answer": "검색할 문서가 없습니다.", "sources": []})
    # 3) 여기에 코드를 작성하세요 (context → chain.invoke → 답 + sources)
    pass


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