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)

# --- 체인: 01 과 같되 프롬프트에 [문서] 자리를 추가 ---
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"))


@app.get("/")
def index():
    return send_from_directory(app.static_folder, "02.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")
    chunks = RecursiveCharacterTextSplitter(
        chunk_size=500, chunk_overlap=50).split_documents(loader.load())
    store.add_documents(chunks)
    return jsonify({"message": f"'{file.filename}' 업로드 완료 ({len(chunks)}개 청크)"})


# --- 질의: 검색 → {context} 로 chain 호출 ---
@app.post("/ask")
def ask():
    question = request.json["question"]
    hits = store.similarity_search(question, k=3)
    if not hits:
        return jsonify({"answer": "아직 업로드된 문서가 없습니다. 먼저 문서를 올려 주세요."})
    context = "\n\n".join(d.page_content for d in hits)
    answer = chain.invoke({"context": context, "question": question})
    return jsonify({"answer": answer})


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