from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_community.document_loaders import TextLoader
from huggingface_hub import InferenceClient 
import transformers
from sentence_transformers import SentenceTransformer
from datasets import Dataset, Features, Value, Sequence
import pandas as pd
import faiss
import os
import torch
import gradio as gr

# LOGIN HF Auth
from huggingface_hub import login

# Ambil token API dari environment variable (jika disimpan di secrets)
import os
hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Autentikasi secara manual menggunakan token
login(token=hf_token)

ST_MODEL = "LazarusNLP/all-indo-e5-small-v4"
BASE_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
# BASE_MODEL = "meta-llama/Llama-3.1-405B-Instruct"
# BASE_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
# BASE_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
DOMAIN_DATA_DIR = "./data"

# SYS_MSG = """
# Kamu adalah asisten dalam sebuah perusahaan penyedia listrik (PLN) yang membantu menjawab pertanyaan seputar 'sexual harassment' dalam Bahasa Indonesia.
# Jawab dengan singkat menggunakan konteks untuk menjawab pertanyaan dalam Bahasa Indonesia.
# """

SYS_MSG = """
Kamu adalah Bestie, asisten yang dirancang untuk membantu karyawan PLN Group dalam memahami, melaporkan, dan menangani insiden pelecehan seksual di tempat kerja. Peran Anda adalah memberikan informasi yang akurat, membantu pengguna berbagi pengalaman mereka, menilai apakah pengalaman tersebut termasuk pelecehan seksual, serta memandu mereka dalam proses pelaporan dengan penuh empati dan kerahasiaan.

Panduan:

 - Informasi dan Edukasi:
Berikan definisi yang jelas, tidak menghakimi, dan akurat mengenai pelecehan seksual berdasarkan kebijakan dan pedoman internal PLN Group.
Tawarkan contoh-contoh konkret tentang apa yang termasuk dalam pelecehan seksual, termasuk namun tidak terbatas pada pelecehan verbal, fisik, visual, atau non-verbal.
Beri tahu pengguna tentang hak-hak mereka dan sumber daya dukungan yang tersedia di dalam PLN Group.
- Bantuan dalam Berbagi Pengalaman:
Dorong pengguna untuk membagikan cerita mereka dengan memastikan ruang yang aman dan bersifat rahasia.
Yakinkan pengguna bahwa informasi mereka akan diperlakukan dengan sangat hati-hati dan mereka tidak wajib mengungkapkan detail jika tidak merasa nyaman.
Berikan pertanyaan yang sensitif untuk memahami konteks pengalaman mereka (misalnya, lokasi, waktu, orang yang terlibat) tanpa memaksa atau menghakimi.

-Penilaian Pengalaman:
Analisis cerita pengguna sesuai dengan kriteria yang tercantum dalam kebijakan pelecehan seksual PLN Group.
Jika kasus tersebut memenuhi definisi pelecehan seksual, sampaikan dengan tenang kepada pengguna dan berikan panduan langkah selanjutnya.
Jika kasus tersebut tidak memenuhi definisi, jelaskan dengan lembut alasan mengapa dan tawarkan sumber daya atau bantuan lain, seperti layanan kesehatan mental atau bantuan internal lebih lanjut.

-Bantuan dalam Pelaporan:
Berikan instruksi kepada pengguna tentang cara melaporkan insiden ke tim atau departemen yang tepat di PLN Group.
Jika perlu, berikan informasi tentang dukungan hukum atau konseling sesuai dengan program dukungan karyawan PLN Group.
Tawarkan opsi kepada pengguna untuk melaporkan pengalaman mereka secara anonim jika diinginkan, sambil menjelaskan keterbatasan pelaporan anonim jika ada.

-Nada dan Gaya:
Selalu jaga nada yang empatik, tidak menghakimi, dan mendukung.
Gunakan bahasa yang sederhana dan jelas agar mudah dipahami oleh semua karyawan.
Hindari bahasa yang menghakimi atau meremehkan.
Responlah dengan sensitif terhadap emosi pengguna, baik itu kebingungan, kesedihan, atau kemarahan.
"""

# --------------------------------------------------------------
# RAG

TOP_K = 1

domain_data = [os.path.join(DOMAIN_DATA_DIR, f) for f in os.listdir(DOMAIN_DATA_DIR) if f.endswith('.txt')]
pages = []

for file in domain_data:
    text_loader = TextLoader(file)
    file_pages = text_loader.load()
    pages.extend(file_pages)

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=64,
    separators=["\n\n"]
)

documents = splitter.split_documents(pages)
content = [doc.page_content.strip() for doc in documents]

ST = SentenceTransformer(ST_MODEL)
embeddings = ST.encode(content)

features = Features({
    'text': Value('string'),
    'embeddings': Sequence(Value('float32'))
})

data = {'text': content, 'embeddings': [embedding.tolist() for embedding in embeddings]}
dataset = Dataset.from_dict(data, features=features)

dataset.add_faiss_index(column='embeddings')

def retrieve(query, top_k=1):
    query_embedding = ST.encode([query])
    scores, retrieved_examples = dataset.get_nearest_examples('embeddings', query_embedding, k=top_k)

    return scores, retrieved_examples['text']

# --------------------------------------------------------------

client = InferenceClient(BASE_MODEL)

examples = [
    ["Saya mengalami pelecehan seksual."],
    ["Bantu saya melaporkan pelecehan seksual."],
    ["Bagaimana prosedur pelaporan pelecehan seksual yang saya alami?"],
    ["Bagaiamana melaporkan pelecahan seksual yang saya lihat?"],
]

def respond(
    message,
    history: list[tuple[str, str]],
    max_tokens=512,
    temperature=0.5,
    top_p=0.9,
):
    # Retrieve top 3 relevant documents based on the user's query
    score, retrieved_docs = retrieve(message, top_k=TOP_K)

    # print(f"Score: {score}")
    if score <= 13:
        # Prepare the retrieved context
        context = "\n".join([f"{doc}" for i, doc in enumerate(retrieved_docs)])
    else:
        context = ""
    
    print(f"Feed:\n{context}")

    messages = [{"role": "system", "content": SYS_MSG}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # messages.append({"role": "user", "content": message})
    
    # Append the current user message along with the retrieved context
    user_context = f"{message}\nKonteks:\n{context}"
    messages.append({"role": "user", "content": user_context})

    # print(f"\nMessages:\n{messages}\n")
    
    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    textbox=gr.Textbox(placeholder="Enter message here", container=False, scale = 7),
    chatbot=gr.Chatbot(value=[[None, "Halo namaku Bestie, apa yang perlu kamu ceritakan hari ini?"]])
)

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