import gradio as gr
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
import warnings
from transformers import pipeline
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import os
from dotenv import load_dotenv

warnings.filterwarnings("ignore")
load_dotenv()

# Constants and configurations
APP_TITLE = "💊 Asisten Kesehatan Feminacare"
INITIAL_MESSAGE = """Halo! 👋 Saya adalah asisten kesehatan feminacare yang siap membantu Anda dengan informasi seputar kesehatan wanita. 
Silakan ajukan pertanyaan apa saja dan saya akan membantu Anda dengan informasi yang akurat."""

# Model configurations
MODEL_NAME = "SeaLLMs/SeaLLMs-v3-7B-Chat"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
TOP_K_DOCS = 5

def initialize_models():
    """Initialize the embedding model and vector store"""
    data_directory = os.path.join(os.path.dirname(__file__), "vector_db_dir")
    embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    vector_store = Chroma(
        embedding_function=embedding_model, 
        persist_directory=data_directory
    )
    return vector_store

def create_llm():
    """Initialize the language model with optimized parameters"""
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        torch_dtype=torch.float16,
        quantization_config=bnb_config
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
    
    text_generation_pipeline = pipeline(
        model=model,
        tokenizer=tokenizer,
        task="text-generation",
        temperature=0.2,
        do_sample=True,
        repetition_penalty=1.1,
        return_full_text=False,
        max_new_tokens=200,
        eos_token_id=terminators,
        device_map="auto"
    )
    
    return HuggingFacePipeline(pipeline=text_generation_pipeline)

PROMPT_TEMPLATE = """
Anda adalah asisten kesehatan profesional dengan nama Feminacare. 
Berikan informasi yang akurat, jelas, dan bermanfaat berdasarkan konteks yang tersedia.
Context yang tersedia:
{context}
Chat history:
{chat_history}
Question: {question}
Instruksi untuk menjawab:
1. Berikan jawaban yang LENGKAP dan TERSTRUKTUR
2. Selalu sertakan SUMBER informasi dari konteks yang diberikan
3. Jika informasi tidak tersedia dalam konteks, katakan: "Maaf, saya tidak memiliki informasi yang cukup untuk menjawab pertanyaan tersebut secara akurat. Silakan konsultasi dengan tenaga kesehatan untuk informasi lebih lanjut."
4. Gunakan bahasa yang mudah dipahami
5. Jika relevan, berikan poin-poin penting menggunakan format yang rapi
6. Akhiri dengan anjuran untuk konsultasi dengan tenaga kesehatan jika diperlukan
Answer:
"""

class HealthAssistant:
    def __init__(self):
        self.vector_store = initialize_models()
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            output_key='answer'
        )
        self.qa_chain = self.setup_qa_chain()

    def setup_qa_chain(self):
        """Set up the QA chain with improved configuration"""
        custom_prompt = PromptTemplate(
            template=PROMPT_TEMPLATE,
            input_variables=["context", "question", "chat_history"]
        )
        
        return ConversationalRetrievalChain.from_llm(
            llm=create_llm(),
            retriever=self.vector_store.as_retriever(),
            memory=self.memory,
            return_source_documents=True,
        )

    def respond(self, message, history):
        """Generate response for user input"""
        if not message:
            return ""
        
        response = self.qa_chain({"question": message})
        return response["answer"]

    def clear_history(self):
        """Clear conversation history"""
        self.memory.clear()
        return []

def create_demo():
    # Initialize the assistant
    assistant = HealthAssistant()
    
    # Create the Gradio interface
    with gr.Blocks(css="footer {visibility: hidden}") as demo:
        gr.Markdown(f"# {APP_TITLE}")
        gr.Markdown("""
        Asisten digital ini dirancang untuk membantu Anda berkonsultasi tentang kesehatan wanita.
        
        _Catatan: Informasi yang diberikan bersifat umum. Selalu konsultasikan dengan tenaga kesehatan untuk saran yang lebih spesifik._
        """)
        
        chatbot = gr.Chatbot(
            value=[[None, INITIAL_MESSAGE]],
            height=400
        )
        
        with gr.Row():
            msg = gr.Textbox(
                placeholder="Ketik pertanyaan Anda di sini...",
                show_label=False,
                scale=9
            )
            submit = gr.Button("Kirim", scale=1)
            clear = gr.Button("🗑️ Hapus Riwayat", scale=1)
        
        # Set up event handlers
        submit_click = submit.click(
            assistant.respond,
            inputs=[msg, chatbot],
            outputs=[chatbot],
            queue=True
        )
        submit_click.then(lambda: "", None, msg)
        
        msg.submit(
            assistant.respond,
            inputs=[msg, chatbot],
            outputs=[chatbot],
            queue=True
        ).then(lambda: "", None, msg)
        
        clear.click(
            assistant.clear_history,
            outputs=[chatbot],
            queue=False
        )
    
    return demo

# Create and launch the demo
demo = create_demo()