from utills import *
import streamlit as st
from streamlit_chat import message
from streamlit_lottie import st_lottie
import json
from Functions import RFPProcessor
from Prompts_and_Chains import LLMChains
function = RFPProcessor()
chains_obj = LLMChains()


if "is_category_selected" not in st.session_state:
    st.session_state["is_category_selected"] = False
if "user_input" not in st.session_state:
    st.session_state["user_input"] = ""


def local_css(file_name):
    with open(file_name, "r") as f:
        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)


def load_lottiefile(filepath: str):
    '''Load lottie animation file'''
    with open(filepath, "r") as f:
        return json.load(f)


def main():
    st.set_page_config(page_title="Justice League Chatbot",
                       page_icon="⚖️", layout="wide")
    local_css("style.css")
    st_lottie(load_lottiefile("images/welcome.json"), speed=1,
              reverse=False, loop=True, quality="high", height=300)
    st.title("🦸‍♂️ Justice League Chatbot")
    st.subheader("Your AI-powered legal assistant")

    if "messages" not in st.session_state:
        st.session_state.messages = []
        st.session_state.user_inputs = {}
        st.session_state.current_question = "start"

    for i, msg in enumerate(st.session_state.messages):
        message(msg["content"], is_user=msg["role"] == "user", key=str(i))

    if not st.session_state.messages:
        initial_message = "Welcome to the Justice League Chatbot! I'm here to help you find the right lawyer or provide general legal information. How can I assist you today?"
        st.session_state.messages.append(
            {"role": "assistant", "content": initial_message})

    show_options()


def show_options():
    options = get_options()

    if st.session_state.current_question == "additional_info":
        get_additional_info()
    else:
        col1, col2 = st.columns([3, 1])
        with col1:
            st.text_input(
                "Type your response or choose an option:",
                st.session_state["user_input"],
                key="user_input",
                on_change=ask_llm
            )
        with col2:
            if st.session_state["is_category_selected"] == False:
                st.write("Quick options:")
                for option in options:
                    if st.button(option, key=f"button_{option}"):
                        handle_user_input(option)


def get_options():
    options_dict = {
        "start": ["Find a lawyer", "Get general legal advice"],
        "category": ["Criminal", "Family", "Corporate", "Immigration"],
        "cost_range": ["Low", "Medium", "High", "Very High"],
        "experience": ["0-5 years", "6-10 years", "11-20 years", "20+ years"],
        "location": ["Jabalpur", "Bhopal", "Indore", "Gwalior"]
    }
    return options_dict.get(st.session_state.current_question, ["Find a lawyer", "Get general legal advice"])


def get_additional_info():
    user_input = st.text_area(
        "Please provide any additional information about your case or specific needs:", key="additional_info")
    if st.button("Submit"):
        handle_user_input(user_input)


def handle_user_input(user_input):
    st.session_state.messages.append({"role": "user", "content": user_input})

    if user_input == "Find a lawyer":
        ask_category()
    elif user_input == "Get general legal advice":
        provide_legal_advice()
    elif st.session_state.current_question == "category":
        st.session_state.user_inputs['category'] = user_input
        ask_cost_range()
    elif st.session_state.current_question == "cost_range":
        st.session_state.user_inputs['cost_range'] = user_input
        ask_experience()
    elif st.session_state.current_question == "experience":
        st.session_state.user_inputs['experience'] = user_input
        ask_location()
    elif st.session_state.current_question == "location":
        st.session_state.user_inputs['location'] = user_input
        ask_additional_info()
    elif st.session_state.current_question == "additional_info":
        st.session_state.user_inputs['additional_info'] = user_input
        show_results(user_input)

    st.experimental_rerun()


def ask_llm():
    user_input = st.session_state["user_input"]
    st.session_state.messages.append({"role": "user", "content": user_input})
    last_5_entries = st.session_state.messages[-5:]
    inputs = {
        "chat_history":last_5_entries,
        "input": user_input,
    }
    output = chains_obj.legal_adviser_bot_chain.run(inputs)
    st.session_state.messages.append(
        {"role": "assistant", "content": output})
    st.session_state["user_input"] = ""


def ask_category():
    response = "What type of lawyer are you looking for?"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "category"


def ask_cost_range():
    response = "What's your budget range?"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "cost_range"


def ask_experience():
    response = "How many years of experience should the lawyer have?"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "experience"


def ask_location():
    response = "Where are you looking for a lawyer?"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "location"


def ask_additional_info():
    response = "Please provide any additional information about your case or specific needs that might help us find the best lawyer for you:"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "additional_info"


def suggest_options():
    response = "I'm not sure how to help with that. Would you like to:"
    st.session_state.messages.append(
        {"role": "assistant", "content": response})
    st.session_state.current_question = "start"


def provide_legal_advice():
    response = "Hello, I'm LegalAssist, an AI chatbot specializing in legal information. I can answer general questions about law and legal procedures, but I can't provide personalized legal advice. How can I assist you with legal information today?"
    st.session_state["is_category_selected"] = True
    st.session_state.messages.append(
        {"role": "assistant", "content": response})


def show_results(additional_info):

    category = st.session_state.user_inputs['category']
    cost_range = st.session_state.user_inputs['cost_range']
    experience = st.session_state.user_inputs['experience']
    location = st.session_state.user_inputs['location']

    user_inputs = {"category":category,"cost_range":cost_range, "experience":experience, "location":location}

    matching_lawyers = search_lawyers(
        category, cost_range, experience, location)

    output = chains_obj.lawyer_recommendations_chain(
        {"user_inputs":user_inputs, "matching_lawyers":matching_lawyers, "additional_info":additional_info})
    
    st.session_state.messages.append(
        {"role": "assistant", "content": output['text']})

    st.session_state.current_question = "start"


if __name__ == "__main__":
    main()