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

warnings.filterwarnings("ignore")
load_dotenv()

# Constants and configurations
TITLE = "💊 Asisten Kesehatan Feminacare"
DESCRIPTION = """
# 💊 Asisten Kesehatan Feminacare
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.*
"""

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 auto device mapping"""
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    terminators = [tokenizer.eos_token_id]
    if hasattr(tokenizer, 'convert_tokens_to_ids'):
        try:
            terminators.append(tokenizer.convert_tokens_to_ids("<|eot_id|>"))
        except:
            pass
    
    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,
    )
    
    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 historyt:
{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):
        vector_store = initialize_models()
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            output_key='answer'
        )
        
        custom_prompt = PromptTemplate(
            template=PROMPT_TEMPLATE,
            input_variables=["context", "question", "chat_history"]
        )
        
        self.qa_chain = ConversationalRetrievalChain.from_llm(
            llm=create_llm(),
            retriever=vector_store.as_retriever(),
            memory=self.memory,
            combine_docs_chain_kwargs={"prompt": custom_prompt},
            return_source_documents=True,
        )

    def respond(self, message, history):
        """Process the message and return a response"""
        response = self.qa_chain({"question": message})
        return response["answer"]
    
    def clear_history(self):
        """Clear the conversation memory"""
        self.memory.clear()
        return None

def create_demo():
    assistant = HealthAssistant()
    
    # Define the interface
    with gr.Blocks(title=TITLE) as demo:
        gr.Markdown(DESCRIPTION)
        
        chatbot = gr.Chatbot(
            label="Chat History",
            height=600,
            show_copy_button=True,
        )
        
        with gr.Row():
            msg = gr.Textbox(
                label="Ketik pertanyaan Anda di sini...",
                placeholder="Contoh: Apa itu PCOS?",
                scale=9
            )
            submit = gr.Button("Kirim", scale=1)
            
        clear = gr.Button("🗑️ Hapus Riwayat Chat")
        
        # Set up event handlers
        submit_click = submit.click(
            assistant.respond,
            inputs=[msg, chatbot],
            outputs=[chatbot],
            show_progress="full"
        )
        submit_click.then(lambda: "", None, msg)  # Clear input after sending
        
        msg.submit(
            assistant.respond,
            inputs=[msg, chatbot],
            outputs=[chatbot],
            show_progress="full"
        ).then(lambda: "", None, msg)  # Clear input after sending
        
        clear.click(
            assistant.clear_history,
            outputs=[chatbot],
            show_progress=True
        )
        
        # Add some CSS styling
        gr.Markdown("""
        <style>
        .gradio-container {
            max-width: 1200px !important;
            margin: auto;
        }
        </style>
        """)
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        enable_queue=True
    )