import os
import uuid
import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 1. LLM 초기화
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)

# 2. 문서 로드 및 벡터 DB 구축
loader = PyPDFLoader("Maximizing Muscle Hypertrophy.pdf")
pages = loader.load_and_split()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(pages)

embeddings = GoogleGenerativeAIEmbeddings(model="gemini-embedding-001")
vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
retriever = vectorstore.as_retriever()

# 검색된 문서를 하나의 문자열로 결합하는 헬퍼 함수
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 3. 프롬프트 정의
qa_prompt = ChatPromptTemplate.from_messages([
    ("system", """논문 리뷰 전문가입니다. 제공된 문서를 바탕으로 한국어로 답변하세요.
문서에 없는 내용은 모른다고 답하세요.

{context}"""),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
])

# 4. 에러가 나던 chains 모듈을 버리고 LCEL(파이프라인) 문법으로 RAG 체인 구축
rag_chain = (
    RunnablePassthrough.assign(context=(lambda x: format_docs(retriever.invoke(x["input"]))))
    | qa_prompt
    | llm
    | StrOutputParser()
)

# 5. 메모리(대화 기록) 연동
store = {}
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

# 6. Gradio 연동 함수
def chat_response(message, history, session_id):
    # LCEL 체인은 딕셔너리가 아닌 문자열을 바로 반환하므로 ["answer"] 추출이 필요 없음
    response = conversational_rag_chain.invoke(
        {"input": message},
        config={"configurable": {"session_id": session_id}}
    )
    return response

# 7. 다중 사용자 환경 UI 실행
with gr.Blocks() as demo:
    session_state = gr.State(lambda: str(uuid.uuid4()))
    
    gr.ChatInterface(
        fn=chat_response,
        additional_inputs=[session_state],
        title="💪 근비대 극대화 논문 Q&A 봇",
        description="'Maximizing Muscle Hypertrophy' 논문에 대해 궁금한 점을 물어보세요!"
    )

if __name__ == "__main__":
    demo.launch()