import os

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

from backend.ai.embedding import watson_embedding
from backend.ai.llm import watson_llm

UPLOAD_PATH = "uploads"

def upload_document(file):
    # file 저장
    file_path = os.path.join(UPLOAD_PATH, file.filename)

    with open(file_path, "wb") as f:
        f.write(file.file.read())

    # pdf 업로드 => 분할 => 인덱스 생성
    # pdf 로드
    loader = PyPDFLoader(file_path)
    docs = loader.load()

    splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=30)
    chunks = splitter.split_documents(docs)

    faiss_store = FAISS.from_documents(documents=chunks, embedding=watson_embedding)
    faiss_store.save_local("./db/vectorstore")

    return {"message": "업로드 성공"}

# 질문 => 유사도 검색 => 문서 => llm 답변 생성
def rag_chat(question: str):

    faiss_store = FAISS.load_local("./db/vectorstore", watson_embedding, allow_dangerous_deserialization=True)

    retriever = faiss_store.as_retriever(search_kwargs={"k": 3})

    ### LLM
    # 1. prompt
    message = """\
    당신은 PDF 기반 RAG AI 입니다.
    다음 문서를 참고해서 질문에 답변하세요.

    문서:
    {context}

    질문:
    {question}
    """

    rag_prompt = ChatPromptTemplate.from_template(message)
    # 2. chain
    chain = {"context": retriever, "question": RunnablePassthrough()} | rag_prompt | watson_llm | StrOutputParser()

    # 3. answer
    answer = chain.invoke(question)
    return answer

async def rag_chat_stream(question: str):

    faiss_store = FAISS.load_local("./db/vectorstore", watson_embedding, allow_dangerous_deserialization=True)

    retriever = faiss_store.as_retriever(search_kwargs={"k": 3})

    ### LLM
    # 1. prompt
    message = """\
    당신은 PDF 기반 RAG AI 입니다.
    다음 문서를 참고해서 질문에 답변하세요.

    문서:
    {context}

    질문:
    {question}
    """

    rag_prompt = ChatPromptTemplate.from_template(message)
    # 2. chain
    chain = {"context": retriever, "question": RunnablePassthrough()} | rag_prompt | watson_llm | StrOutputParser()

    # 3. answer
    # chain.stream() : 동기방식(요청 => 응답을 할 떄까지 기다리는 방식)
    # chain.astream() : 비동기방식(다른 일 처리 가능)
    async for chunk in chain.astream(question):
        yield chunk