import gradio as gr
import os
from langchain_community.vectorstores import FAISS
from langchain_together import TogetherEmbeddings, Together
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.memory import ConversationBufferMemory
from typing import List, Tuple
import re

TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')


class ChatBot:
    def __init__(self):
        self.embeddings = TogetherEmbeddings(
            model="togethercomputer/m2-bert-80M-32k-retrieval",
            together_api_key=TOGETHER_API_KEY
        )













        
        self.vectorstore = FAISS.load_local(
            ".",
            embeddings=self.embeddings,
            allow_dangerous_deserialization=True  
        )
        self.retriever = self.vectorstore.as_retriever()
        
        self.model = Together(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
            temperature=0.4,
            max_tokens=256,
            top_k=30,
            together_api_key=TOGETHER_API_KEY
        )
        
        self.memory = ConversationBufferMemory(
            return_messages=True,
            memory_key="chat_history",
            output_key="answer"
        )
        
        self.template = """Quyidagi ko'rsatmalarga qat'iy rioya qiling:



1. Faqat o'zbek tilida javob bering
2. Faqat berilgan ma'lumotlar asosida javob bering
3. Agar savol tushunarsiz bo'lsa yoki ma'lumot bo'lmasa, "Kechirasiz, bu haqida ma'lumotga ega emasman" deb javob bering
4. O'zingizdan savol bermang
5. Javobni takrorlamang
6. Salomlashish uchun "Assalomu alaykum" yoki "Vaalaykum assalom" dan foydalaning

Kontekst: {context}
Suhbat Tarixi: {chat_history}
Savol: {question}

Javob:"""



























        self.prompt = ChatPromptTemplate.from_template(self.template)
        
        self.chain = (
            {
                "context": self.retriever,
                "chat_history": lambda x: self.get_chat_history(),
                "question": RunnablePassthrough()
            }
            | self.prompt
            | self.model
            | StrOutputParser()





















        )

    def get_chat_history(self) -> str:
        messages = self.memory.load_memory_variables({})["chat_history"]
        return "\n".join([f"{m.type}: {m.content}" for m in messages])
    
    def process_response(self, response: str) -> str:
        unwanted_tags = ["[INST]", "[/INST]", "<s>", "</s>"]
        for tag in unwanted_tags:
            response = response.replace(tag, "")
        
        response = re.sub(r"```.*?```", "", response, flags=re.DOTALL)
        response = re.sub(r"print\(.*?\)", "", response)
        response = re.sub(r'\s+', ' ', response)
    
        return response.strip()

    def chat(self, message: str, history: List[Tuple[str, str]]) -> str:


























        try:


            
            if message == "__init__":
                return "Assalomu alaykum. Sizga qanday yordam bera olaman?"
                
            self.memory.chat_memory.add_user_message(message)
            response = self.chain.invoke(message)
            clean_response = self.process_response(response)
            
            if not clean_response or len(clean_response.split()) < 3:
                clean_response = "Kechirasiz, savolingizni tushunolmadim. Iltimos, batafsilroq savol bering."
            
            self.memory.chat_memory.add_ai_message(clean_response)
            return clean_response
        except Exception as e:
            return f"Xatolik yuz berdi: {str(e)}"

    def reset_chat(self) -> List[Tuple[str, str]]:
        self.memory.clear()
        return []

def create_demo() -> gr.Interface:
    chatbot = ChatBot()






    
    with gr.Blocks() as demo:
        gr.Markdown("""# RAG Chatbot
        Beeline Uzbekistanning jismoniy shaxslar uchun tariflari haqida ma'lumotlar beruvchi bot""")
        
        
        chatbot_interface = gr.Chatbot(
            height=600,
            show_copy_button=True,

        )
        
        with gr.Row():
            msg = gr.Textbox(
                show_label=False,
                placeholder="Xabaringizni shu yerda yozing",
                container=False
            )
            submit = gr.Button("Xabarni yuborish", variant="primary")
        
        clear = gr.Button("Yangi suhbat")
        
        def respond(message, chat_history):
            message = message.strip()
            if not message:
                return "", chat_history
            
            bot_message = chatbot.chat(message, chat_history)
            chat_history.append((message, bot_message))
            return "", chat_history
        
        def init_chat():
            
            initial_greeting = chatbot.chat("__init__", [])
            return [("", initial_greeting)]
        

        submit.click(respond, [msg, chatbot_interface], [msg, chatbot_interface])
        msg.submit(respond, [msg, chatbot_interface], [msg, chatbot_interface])
        clear.click(init_chat, None, chatbot_interface)
        
        
        demo.load(init_chat, None, chatbot_interface)
    
    return demo


demo = create_demo()

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