import os
import gradio as gr
from model.controller import Controller
import zipfile

os.chdir("/home/user/app")

os.system('wget -O processed_cases.csv "https://drive.usercontent.google.com/download?id=1jMuQtywo0mbj7ZHCCsyE8xurbSyVVCst&export=download&confirm=t&uuid=2f681c98-86f8-4159-9e03-673cdcbc7cb51"')
os.system('wget -O chromadb_collection.zip "https://drive.usercontent.google.com/download?id=1gz5-gxSlySEtPTzL_VPQ9e8jxHFuL0ZJ&export=download&confirm=t&uuid=de946efb-47b3-435d-b432-3bd5c01c73fb"')

with zipfile.ZipFile("chromadb_collection.zip", 'r') as zip_ref:
    zip_ref.extractall()

os.system('mv content/chromadb_collections chromadb_collections')
os.system('rm -r content')

bot = Controller()

def chatbot_interface(user_input, chat_id=2311):
    return bot.handle_message(chat_id, user_input)

def validate_input(user_input):
    if not user_input or user_input.strip() == "":
        return False, "🚫 Please enter a valid legal question. It cannot be empty."
    if len(user_input) < 5:
        return False, "⚠️ Your question is too short. Please provide more details."
    return True, None

custom_css = """
@font-face {
    font-family: 'Vazir';
    src: url('https://cdn.jsdelivr.net/gh/rastikerdar/vazir-font/vf/Vazir.woff2') format('woff2'),
         url('https://cdn.jsdelivr.net/gh/rastikerdar/vazir-font/vf/Vazir.woff') format('woff');
}
.gradio-container {
    background-color: #f9f9f9;
}
.chatbox, .inputbox {
    font-family: 'Vazir', sans-serif;
    font-size: 16px;
}
"""

with gr.Blocks(css=custom_css) as interface:
    
    gr.Markdown("""
    <div style="text-align: center; font-family: 'Vazir';">
    <h1 style="color: #4a90e2;">⚖️ RAG Law Chatbot ⚖️</h1>
    <p style="font-size: 18px; color: #333;">Welcome to the legal chatbot! 👨‍⚖️👩‍⚖️<br>Ask any legal question, and our assistant will help you! 📜🏛️</p>
    </div>
    """)
    
    # Organize the chatbot area in a column for vertical stacking
    with gr.Column():
        chatbot = gr.Chatbot(label="🧑‍⚖️ Legal Chatbot Assistant 🧑‍⚖️", height=400, elem_classes=["chatbox"])
    
    # Use Row to align input and button horizontally
    with gr.Row():
        user_input = gr.Textbox(show_label=False, placeholder="Enter your law question here... ⚖️", container=True)
        send_button = gr.Button("📤 Send")
    
    # Chat update function to append new messages to the chatbot
    def chat_update(user_message, history):
        history = history or []
        
        is_valid, validation_message = validate_input(user_message)
        if not is_valid:
            history.append((user_message, validation_message))
            return history, ""

        bot_reply = chatbot_interface(user_message)
        history.append((user_message, bot_reply))
        return history, ""

    # Connect the button click to the chat update function
    send_button.click(chat_update, [user_input, chatbot], [chatbot, user_input])

interface.launch()