import os
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

# 1. 문서 로드 및 벡터 DB 구축 (서버 구동 시 1회 고정)
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)

# 미션 3: 도메인 맞춤 시스템 프롬프트
SYSTEM_PROMPT = """당신은 스포츠 영양학 및 근비대(Muscle Hypertrophy) 훈련 분야의 최고 권위자이자 논문 리뷰 전문가입니다.
제공된 [논문 컨텍스트]를 바탕으로 사용자의 질문에 전문적이고 명확하며 객관적인 어조로 답변하세요.

[제약 조건]
1. 반드시 제공된 컨텍스트 내의 정보만을 사용하여 답변하세요.
2. 논문에 없는 내용을 질문하면 "해당 내용은 제공된 논문에서 확인할 수 없습니다."라고 명확히 선을 그으세요.
3. 근육 성장 기전이나 훈련법을 설명할 때는 일반인도 이해하기 쉽게 단계별로 구조화하여 설명하세요.
4. 모든 답변은 한국어로 작성하며, 주요 의학 및 운동학 전문 용어는 괄호 안에 영문을 병기하세요 (예: 단백질 합성(Protein Synthesis)).

[논문 컨텍스트]
{context}"""

qa_prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("placeholder", "{chat_history}"),
    ("human", "{input}"),
])

# Gradio의 대화 기록 형식을 LangChain이 이해할 수 있게 변환하는 헬퍼 함수
def format_history(history):
    formatted = []
    for user_msg, ai_msg in history:
        formatted.append(("human", user_msg))
        formatted.append(("ai", ai_msg))
    return formatted

# 미션 1, 2, 5 통합: 스트리밍, 동적 설정, 출처 파싱
def chat_response(message, history, temperature, k, model_name):
    # 미션 2: UI에서 넘겨받은 k 값으로 검색 범위 동적 조절
    docs = vectorstore.similarity_search(message, k=k)
    context = "\n\n".join(doc.page_content for doc in docs)
    
    # 미션 2: UI에서 넘겨받은 모델과 온도로 LLM 동적 생성
    llm = ChatGoogleGenerativeAI(model=model_name, temperature=temperature)
    
    # 프롬프트 조립
    prompt_value = qa_prompt.invoke({
        "context": context,
        "chat_history": format_history(history),
        "input": message
    })
    
    partial_message = ""
    # 미션 5: llm.stream()을 활용한 실시간 스트리밍 출력
    for chunk in llm.stream(prompt_value):
        partial_message += chunk.content
        yield partial_message # 글자가 생성될 때마다 UI로 밀어냄
        
    # 미션 1: PyPDFLoader 메타데이터에서 출처 및 페이지 추출 (page는 0부터 시작하므로 +1)
    sources = []
    for doc in docs:
        source_file = os.path.basename(doc.metadata.get('source', 'Unknown'))
        page_num = doc.metadata.get('page', 0) + 1
        sources.append(f"{source_file} (p.{page_num})")
        
    # 리스트 중복 제거 후 최종 텍스트 조립
    unique_sources = list(dict.fromkeys(sources))
    source_str = "\n\n📎 **출처:** " + ", ".join(unique_sources)
    
    # 최종적으로 답변 끝에 출처를 덧붙여서 전송
    yield partial_message + source_str

# 미션 4: 대화 내역 다운로드 파일 생성 함수
def download_chat_history(history):
    file_path = "chat_history.txt"
    with open(file_path, "w", encoding="utf-8") as f:
        for user_msg, ai_msg in history:
            f.write(f"🧑‍💻 사용자: {user_msg}\n")
            f.write(f"🤖 AI: {ai_msg}\n")
            f.write("-" * 50 + "\n")
    return file_path

# UI 레이아웃 구성
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 💪 근비대 극대화 논문 Q&A 봇 (Pro Version)")
    
    # 미션 2: 접을 수 있는 설정 패널
    with gr.Accordion("⚙️ 챗봇 상세 설정", open=False):
        with gr.Row():
            model_dd = gr.Dropdown(choices=["gemini-2.0-flash", "gemini-1.5-pro", "gemini-1.5-flash"], value="gemini-2.0-flash", label="🤖 모델 선택")
            temp_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, label="🌡️ Temperature (창의성/환각 조절)")
            k_slider = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="📚 참고할 문서 조각 수 (k)")

    # 핵심 챗봇 인터페이스 (설정 패널의 값들을 additional_inputs로 연결)
    chat_interface = gr.ChatInterface(
        fn=chat_response,
        additional_inputs=[temp_slider, k_slider, model_dd],
        chatbot=gr.Chatbot(height=500),
        title="",
        description="'Maximizing Muscle Hypertrophy' 논문 내용을 바탕으로 근성장 메커니즘을 질문해 보세요."
    )
    
    # 미션 4: 대화 내역 다운로드 영역
    with gr.Row():
        download_btn = gr.Button("💾 현재 대화 내역 저장 및 다운로드", variant="primary")
        download_file = gr.File(label="다운로드 준비 완료 (버튼을 누르세요)")
        
    # 버튼 클릭 이벤트 (채팅창의 히스토리를 가져와 파일로 변환)
    download_btn.click(
        fn=download_chat_history, 
        inputs=[chat_interface.chatbot], 
        outputs=[download_file]
    )

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