from flask import Flask, request, jsonify
from langchain_groq import ChatGroq
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
import subprocess
import os
import logging
import time
from dotenv import load_dotenv


load_dotenv()
app = Flask(__name__)
logging.basicConfig(level=logging.INFO)

if os.getenv("GROQ_API_KEY") is None:
    app.logger.error("Error: GROQ_API_KEY tidak diatur dalam environment.")

CHROMA_PATH = "chromadb"

# Variabel Global
rag_chain = None
suggestion_chain = None
vector_db_global = None  
retriever = None
rewrite_chain = None
embeddings_global = None


PROMPT_TEMPLATE = """
Anda adalah Asisten Virtual BKKBN yang profesional, ramah, dan edukatif.
Tugas utama dan SATU-SATUNYA adalah menjawab pertanyaan pengguna HANYA berdasarkan [Konteks] yang diberikan.

BATASAN MUTLAK (ANTI-HALUSINASI):
1. Anda DILARANG KERAS menggunakan pengetahuan umum, memori bawaan, atau informasi dari internet untuk menjawab.
2. JIKA pertanyaan TIDAK RELEVAN dengan program BKKBN/Keluarga Berencana (misalnya: ramalan, cuaca, politik, hari baik, dll), ATAU informasi untuk menjawab BENAR-BENAR TIDAK ADA di dalam [Konteks], Anda WAJIB membalas dengan kalimat persis berikut tanpa tambahan apapun:
"Maaf, saya belum memiliki informasi mengenai hal tersebut. Namun, Anda bisa mencoba memilih salah satu topik rekomendasi di bawah ini."

Instruksi Format Jawaban:
1. Jawablah dengan struktur yang rapi dan mudah dibaca.
2. Gunakan poin-poin (bullet points) atau penomoran jika menjelaskan langkah-langkah atau daftar.
3. Gunakan **huruf tebal** untuk kata kunci penting.
4. Hindari paragraf yang terlalu panjang (tembok teks). Pisahkan antar paragraf.
5. PENALARAN SINONIM: Jika pertanyaan menggunakan kata tertentu, namun konteks menggunakan kata lain yang bermakna sama, Anda HARUS menggunakan informasi tersebut untuk menjawab. Jangan kaku pada satu kata yang sama persis!
6. SANGAT PENTING: Jika informasi BENAR-BENAR tidak ada di konteks, Anda WAJIB menjawab HANYA dengan: "Maaf, saya belum memiliki informasi mengenai hal tersebut. Namun, Anda bisa mencoba memilih salah satu topik rekomendasi di bawah ini."
7. Penutup Hangat: Akhiri dengan kalimat seperti:"Kalau masih bingung, tanya aja ya!" "Mau aku bantu jelasin lebih detail juga bisa 👍"
8. Nada Komunikasi : Ramah, santai, tapi tetap sopan Tidak terlalu formal kaku.

Konteks:
{context}

Pertanyaan:
{question}

Jawaban:
"""

REWRITE_PROMPT = """Tugas Anda adalah memperbaiki pertanyaan pengguna agar lebih deskriptif dan formal dalam konteks program BKKBN (Keluarga Berencana, Stunting, Pernikahan, Alat Kontrasepsi).
Jangan menjawab pertanyaannya, cukup ubah kalimatnya menjadi pertanyaan yang lengkap.

Contoh:
Input: "Stunting"
Output: "Apa yang dimaksud dengan stunting dan bagaimana pencegahannya?"

Input: "Syarat nikah"
Output: "Apa saja persyaratan dokumen dan kesehatan untuk pendaftaran pernikahan?"

Input: {user_query}
Output:"""

SUGGESTION_PROMPT_TEMPLATE = """
Anda adalah asisten yang bertugas membuat saran pertanyaan lanjutan.
Berdasarkan [Konteks Dokumen] berikut, buatlah tepat 3 rekomendasi pertanyaan lanjutan singkat yang relevan dengan sisa informasi yang ada di teks.

[Konteks Dokumen]:
{context}

Pertanyaan Pengguna Sebelumnya:
{question}

Jawaban Bot Saat Ini:
{answer}

BATASAN MUTLAK (ANTI-HALUSINASI):
1. LOGIKA PENYELAMAT: Jika "Jawaban Bot Saat Ini" berisi kalimat penolakan seperti "Maaf, saya belum memiliki informasi", maka ABAIKAN [Konteks Dokumen]. Sebagai gantinya, keluarkan 3 pertanyaan dasar BKKBN ini:
1. Apa saja syarat pendaftaran nikah?
2. Bagaimana alur konseling calon pengantin?
3. Apa itu aplikasi Elsimil?
2. Jika bot berhasil menjawab, buatlah 3 pertanyaan lanjutan yang 100% jawabannya PASTI ADA di dalam sisa informasi [Konteks Dokumen] di atas.
3. DILARANG KERAS menggunakan pengetahuan umum Anda di luar teks!
4. Jika teks hanya membahas A dan B, JANGAN PERNAH membuat pertanyaan tentang C.
5. SANGAT PENTING: JANGAN menyertakan frasa seperti "berdasarkan teks" atau "menurut dokumen". Buatlah kalimat pertanyaannya senatural mungkin.

Keluarkan HANYA 3 pertanyaan lanjutan dengan format list angka (1. 2. 3.) tanpa kalimat pengantar atau penutup apapun.
"""

def initialize_rag_chain():
    global vector_db_global, embeddings_global 
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

    embeddings_global = embeddings  
    try:
        if not os.path.exists(CHROMA_PATH) or not os.listdir(CHROMA_PATH):
            app.logger.warning("Database ChromaDB tidak ditemukan atau kosong.")
            return None, None, None, None

        # Embeddings
        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        
        # Vector DB 
        vector_db_global = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
        retriever = vector_db_global.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={'k': 4, 'score_threshold': 0.70}
        )
        
        llm_rewrite = ChatGroq(model="llama-3.1-8b-instant", temperature=0.0)
        rewrite_prompt = ChatPromptTemplate.from_template(
            "Ubah kalimat ini menjadi satu pertanyaan formal tanpa pembukaan dan tanpa menjawabnya: {user_query}"
        )
        rewrite_chain_init = rewrite_prompt | llm_rewrite | StrOutputParser()
        
        llm_main = ChatGroq(model="llama-3.3-70b-versatile", temperature=0.3)
        main_prompt = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
        main_chain_init = main_prompt | llm_main | StrOutputParser()

        llm_suggestion = ChatGroq(model="llama-3.1-8b-instant", temperature=0.0)
        suggestion_prompt = ChatPromptTemplate.from_template(SUGGESTION_PROMPT_TEMPLATE)
        suggestion_chain_init = suggestion_prompt | llm_suggestion | StrOutputParser()

        app.logger.info("LangChain (Main, Suggestion, Rewrite) berhasil diinisialisasi.")
        return main_chain_init, suggestion_chain_init, retriever, rewrite_chain_init

    except Exception as e:
        app.logger.error(f"Gagal menginisialisasi LangChain: {e}")
        return None, None, None, None

# Inisialisasi awal
rag_chain, suggestion_chain, retriever, rewrite_chain = initialize_rag_chain()


small_talk_responses = {
    "cantik": "Terima kasih! Fokus saya adalah memberikan informasi yang akurat untuk Anda.",
    "ganteng": "Terima kasih! Fokus saya adalah memberikan informasi yang akurat untuk Anda.",
    "pintar": "Terima kasih, saya belajar dari dokumen-dokumen yang diberikan.",
    "cerdas": "Terima kasih, saya belajar dari dokumen-dokumen yang diberikan.",
    "siapa kamu": "Saya adalah Asisten Virtual BKKBN, siap membantu Anda dengan informasi seputar program kami.",
    "terima kasih": "Sama-sama! Senang bisa membantu.",
    "makasih": "Sama-sama! Senang bisa membantu.",
    "halo": "Halo juga! Ada yang bisa saya bantu?",
    "hai": "Hai juga! Ada yang bisa saya bantu?",
    "selamat pagi": "Selamat pagi! Ada yang bisa saya bantu?",
    "selamat siang": "Selamat siang! Ada yang bisa saya bantu?",
    "selamat sore": "Selamat sore! Ada yang bisa saya bantu?",
    "selamat malam": "Selamat malam! Ada yang bisa saya bantu?",
    "kontak": "Silakan hubungi nomor 082141236706 untuk informasi lebih lanjut.",
}

@app.route('/chat', methods=['POST'])
def chat():
    global rag_chain, suggestion_chain, vector_db_global, retriever, rewrite_chain

    try:
        data = request.get_json()
        user_message = data.get("message", "").strip() 
        session_id = data.get("session_id", "Tanpa ID")
        
        if not user_message:
            return jsonify({"error": "Pesan tidak boleh kosong."}), 400

        lower_message = user_message.lower()
        if lower_message in small_talk_responses:
            return jsonify({
                "reply": small_talk_responses[lower_message],
                "suggestions": [],
                "score": 1.0 
            })

        if not rag_chain or not retriever or not rewrite_chain:
            return jsonify({"reply": "Maaf, saat ini saya belum memiliki dokumen pengetahuan..."})
        # start_time = time.time()

        # # proses retrieval + generate jawaban
        # docs = retriever.invoke(user_message)

        # bot_response = rag_chain.invoke({
        #     "context": context_text, 
        #     "question": user_message 
        # })

        # end_time = time.time()
        # response_time = round(end_time - start_time, 2)
        # app.logger.info(f"Menerima pesan dari User: '{user_message}'")
        # app.logger.info(f"Waktu respon: {response_time} detik")

        refined_query = rewrite_chain.invoke({"user_query": user_message}).strip()
        hybrid_query = f"{user_message} {refined_query}"
        app.logger.info(f"Hybrid Query: '{hybrid_query}'")

        docs_for_score = vector_db_global.similarity_search_with_score(hybrid_query, k=1)
        top_score = round(max(0, 1 - (docs_for_score[0][1] / 2)), 4) if docs_for_score else 0

        docs_for_score = vector_db_global.similarity_search_with_relevance_scores(hybrid_query, k=4)

        docs_for_score = docs_for_score[:2]

        doc_scores = []

        for doc, score in docs_for_score:
            doc_scores.append({
            "score": float(score)
        })
        docs = retriever.invoke(hybrid_query)
        context_text = "\n\n".join([doc.page_content for doc in docs])
        
        bot_response = rag_chain.invoke({
            "context": context_text, 
            "question": user_message 
        })
        try:
            query_vector = embeddings_global.embed_query(hybrid_query)
            query_vector = list(query_vector)  
        except Exception as e:
            app.logger.error(f"Vector error: {e}")
            query_vector = []

        suggestions_list = []
        if "Maaf, saya belum memiliki informasi" in bot_response:
            suggestions_list = ["Apa itu stunting?", "Syarat nikah apa saja?", "Apa itu Elsimil?"]
        else:
            suggestions_text = suggestion_chain.invoke({
                "context": context_text,
                "question": user_message, 
                "answer": bot_response
            })
            suggestions_list = [s.lstrip('1234567890. ') for s in suggestions_text.split('\n') if s.strip()]
            suggestions_list = suggestions_list[:3] 

        return jsonify({
            "reply": bot_response,
            "suggestions": suggestions_list,
            "score": top_score, 
            "vector": query_vector,
            "doc_scores": doc_scores
        })

    except Exception as e:
        app.logger.error(f"Error pada endpoint /chat: {e}", exc_info=True)
        return jsonify({"error": "Terjadi kesalahan pada server."}), 500


@app.route('/re-ingest', methods=['POST'])
def reingest_documents():
    global rag_chain, suggestion_chain, vector_db_global, retriever, rewrite_chain

    try:
        app.logger.info("Memulai proses ingestion ulang...")
        
        if vector_db_global is not None:
            try:
                vector_db_global.delete_collection()
                app.logger.info("Tabel memori database berhasil dikosongkan secara aman.")
            except Exception as e:
                app.logger.warning(f"Gagal mengosongkan tabel memori: {e}")
        
        rag_chain = None
        suggestion_chain = None
        vector_db_global = None
        retriever = None
        rewrite_chain = None

        python_executable = os.sys.executable
        
        result = subprocess.run(
            [python_executable, 'ingest.py'], 
            capture_output=True, 
            text=True, 
            encoding='utf-8' 
        )
        
        if result.returncode != 0:
            error_msg = result.stderr if result.stderr else result.stdout
            app.logger.error(f"Ingest Gagal: {error_msg}")
            return jsonify({
                "error": "Proses ingestion gagal dijalankan.", 
                "details": error_msg 
            }), 500

        app.logger.info("Proses ingestion fisik selesai.")
        
        rag_chain, suggestion_chain, retriever, rewrite_chain = initialize_rag_chain()
        app.logger.info("LangChain berhasil diperbarui dari awal.")
        
        return jsonify({
            "message": "Proses ingestion berhasil dan database diperbarui.", 
            "output": result.stdout
        }), 200

    except Exception as e:
        app.logger.error(f"Error pada endpoint /re-ingest: {e}", exc_info=True)
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False, threaded=True)