import os
import uuid
from flask import Flask, request, jsonify, session, send_from_directory
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from dotenv import load_dotenv

load_dotenv()
app = Flask(__name__)
app.secret_key = "day3_history_demo"

UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
# 04 의 프롬프트에 history 자리만 추가
prompt = ChatPromptTemplate.from_messages([
    ("system", "아래 [문서]와 이전 대화를 근거로 한국어로 간결히 답하세요. 문서에 없으면 '문서에 없습니다'라고만 답하세요."),
    MessagesPlaceholder("history"),
    ("human", "[문서]\n{context}\n\n[질문] {question}"),
])
chain = prompt | llm | StrOutputParser()

# 세션별 대화 저장소 (배운 그대로)
SESSIONS = {}   # {session_id: InMemoryChatMessageHistory}


def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    if session_id not in SESSIONS:
        SESSIONS[session_id] = InMemoryChatMessageHistory()
    return SESSIONS[session_id]


# 체인을 '메모리 달린 체인'으로 감싼다
conversational = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="question",      # 사용자 입력으로 기록할 키
    history_messages_key="history",     # 과거 대화를 끼워 넣을 placeholder 이름
)

store = InMemoryVectorStore(OpenAIEmbeddings(model="text-embedding-3-small"))
DOC_IDS = {}


@app.get("/")
def index():
    return send_from_directory(app.static_folder, "05.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)
    ids = store.add_documents(chunks)
    DOC_IDS.setdefault(file.filename, []).extend(ids)
    return jsonify({"message": f"'{file.filename}' 업로드 완료 ({len(chunks)}개 청크)", "files": list(DOC_IDS)})


@app.delete("/files/<name>")
def delete_file(name):
    ids = DOC_IDS.pop(name, None)
    if not ids:
        return jsonify({"error": f"'{name}' 은(는) 없는 문서입니다"}), 404
    store.delete(ids=ids)
    return jsonify({"message": f"'{name}' 삭제됨 ({len(ids)}개 청크 제거)", "files": list(DOC_IDS)})


@app.post("/ask")
def ask():
    if "sid" not in session:
        session["sid"] = uuid.uuid4().hex          # 세션마다 고유 id (대화 격리)
    question = request.json["question"]
    hits = store.similarity_search(question, k=2)
    if not hits:
        return jsonify({"answer": "검색할 문서가 없습니다.", "sources": []})
    context = "\n\n".join(f"[출처: {d.metadata.get('source')}] {d.page_content}" for d in hits)

    answer = conversational.invoke(
        {"context": context, "question": question},
        config={"configurable": {"session_id": session["sid"]}},   # 이 세션의 대화에 누적
    )

    sources = []
    for d in hits:
        s = d.metadata.get("source")
        if s not in sources:
            sources.append(s)
    return jsonify({"answer": answer, "sources": sources})


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