# Impor library yang dibutuhkan
from flask import Flask, render_template, request, jsonify, session
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from dotenv import load_dotenv
import os

# Memuat variabel lingkungan dari file .env
load_dotenv()

# Inisialisasi aplikasi Flask
app = Flask(__name__)
# Menetapkan secret key untuk keamanan sesi
app.secret_key = os.urandom(24)

# --- KONFIGURASI AWAL ---

# Memuat vectorstore dari direktori "data"
try:
    vectorstore = Chroma(
        persist_directory="data",
        embedding_function=HuggingFaceEmbeddings(model_name="firqaaa/indo-sentence-bert-large")
    )
    print("Vectorstore berhasil dimuat.")
except Exception as e:
    print(f"Kesalahan saat memuat vectorstore: {e}")

# Menyiapkan retriever untuk mengambil dokumen dari vectorstore
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 20})

# Menyiapkan model bahasa (LLM) dari Google
llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)

# --- MANAJEMEN MEMORI PERCAKAPAN ---

# Variabel global untuk menyimpan objek memori di sisi server
# Kunci: session_id, Nilai: objek ConversationBufferMemory
session_memories = {}

def get_session_memory(session_id: str) -> ConversationBufferMemory:
    """
    Mengambil atau membuat objek memori untuk session_id tertentu.
    Objek memori disimpan di variabel global `session_memories`.
    """
    if session_id not in session_memories:
        print(f"Membuat instance memori baru untuk sesi: {session_id}")
        session_memories[session_id] = ConversationBufferMemory(
            memory_key="chat_history", return_messages=True
        )
    return session_memories[session_id]

# --- PEMBUATAN RAG CHAIN (RANTAI RAG) ---

# Prompt untuk memformulasikan ulang pertanyaan berdasarkan riwayat percakapan
contextualize_q_system_prompt = (
    "Berdasarkan riwayat percakapan dan pertanyaan terbaru dari pengguna, "
    "buatlah pertanyaan baru yang berdiri sendiri dan dapat dipahami tanpa melihat riwayat percakapan. "
    "JANGAN menjawab pertanyaan tersebut, cukup formulasikan ulang jika diperlukan, jika tidak, kembalikan pertanyaan asli."
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

# Rantai yang bertugas memformulasikan ulang pertanyaan
history_aware_retriever = create_history_aware_retriever(
    llm, retriever, contextualize_q_prompt
)

# Prompt utama untuk menjawab pertanyaan berdasarkan konteks yang diberikan retriever
qa_system_prompt = (
    "Anda adalah asisten AI yang teliti untuk BPVP Sorong. "
    "Gunakan HANYA informasi dari Konteks yang diberikan untuk menjawab pertanyaan. "
    "JANGAN gunakan pengetahuan eksternal. Jawaban Anda harus berdasarkan fakta dari konteks. "
    "Jika pertanyaan tidak bisa dijawab dari konteks, katakan: 'Maaf, saya tidak menemukan informasi tersebut dalam dokumen saya.'\n\n"
    "Konteks:\n{context}"
)
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

# Rantai yang menggabungkan dokumen-dokumen konteks menjadi satu string
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

# Rantai RAG utama yang mengorkestrasi seluruh alur
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)


# --- ROUTE (ENDPOINT) UNTUK APLIKASI WEB FLASK ---

@app.route('/')
def home():
    """Menampilkan halaman utama dan membuat session_id jika belum ada."""
    if 'session_id' not in session:
        session['session_id'] = os.urandom(16).hex()
    return render_template('index.html')

@app.route('/get', methods=['GET'])
def get_response():
    """Menerima pesan dari pengguna, memprosesnya dengan RAG chain, dan mengembalikan jawaban."""
    user_message = request.args.get('msg')
    session_id = session.get('session_id')
    
    if not session_id:
        return jsonify({"error": "Sesi tidak ditemukan. Silakan muat ulang halaman."}), 400
        
    # Mengambil memori yang sesuai untuk sesi ini
    memory = get_session_memory(session_id)
    chat_history = memory.load_memory_variables({})["chat_history"]

    # Menjalankan RAG chain dengan input yang diperlukan
    response = rag_chain.invoke({"input": user_message, "chat_history": chat_history})
    
    # Menyimpan percakapan baru ke dalam memori
    memory.save_context({"input": user_message}, {"answer": response["answer"]})
    
    # Mengembalikan hanya teks jawaban ke frontend
    return jsonify(response["answer"])

@app.route('/load_history', methods=['GET'])
def load_history():
    """Mengembalikan riwayat percakapan untuk sesi saat ini."""
    session_id = session.get('session_id')
    if not session_id:
        return jsonify([])
        
    memory = get_session_memory(session_id)
    history = memory.load_memory_variables({})["chat_history"]
    
    # Memformat riwayat agar mudah dibaca oleh frontend
    formatted_history = []
    for i in range(0, len(history), 2):
        user_msg = history[i].content
        bot_msg = history[i+1].content if i+1 < len(history) else ""
        formatted_history.append({"sender": "user", "message": user_msg})
        formatted_history.append({"sender": "bot", "message": bot_msg})
        
    return jsonify(formatted_history)

@app.route('/clear_history', methods=['POST'])
def clear_history():
    """Menghapus riwayat percakapan untuk sesi saat ini dari memori server."""
    session_id = session.get('session_id')
    
    if session_id and session_id in session_memories:
        del session_memories[session_id]
        print(f"Riwayat untuk sesi {session_id} telah dihapus dari memori server.")
    
    return jsonify({"status": "success", "message": "Riwayat percakapan telah dihapus"})

# Menjalankan aplikasi
if __name__ == '__main__':
    app.run(debug=True)