import streamlit as st
from agents import create_judge_agent, create_advocate_agent, create_consultant_agent
from crewai import Task, Crew
from utils import is_arabic, format_legal_response
from config import LEGAL_CATEGORIES, DEFAULT_LANGUAGE

st.set_page_config(page_title="المساعد القانوني الإماراتي", layout="wide")

# Load custom CSS
with open('style.css') as f:
    st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)

st.title("المساعد القانوني الإماراتي")
st.write("احصل على المساعدة القانونية من خبراء قانونيين إماراتيين مدعومين بالذكاء الاصطناعي")

# Language selector
language = st.sidebar.selectbox(
    "اختر اللغة / Select Language",
    ["العربية", "English"],
    index=0
)

# Legal category selector
selected_category = st.sidebar.selectbox(
    "اختر الفئة القانونية / Select Legal Category",
    list(LEGAL_CATEGORIES.values()),
    index=0
)

# Initialize session state for chat history
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# Create tabs for different agents
tab1, tab2, tab3 = st.tabs(["القاضي", "المحامي", "المستشار"])

def get_agent_response(agent, query, category):
    try:
        # Prepare the task with context
        task_description = f"""
        تحليل والرد على الاستفسار التالي في مجال {category}:
        {query}
        
        يجب أن يكون الرد:
        1. مستنداً إلى القانون الإماراتي
        2. مدعوماً بالمراجع القانونية
        3. واضحاً ومفهوماً
        4. متوافقاً مع أحدث التشريعات
        """
        
        task = Task(
            description=task_description,
            agent=agent,
            expected_output="تحليل قانوني ورد بناءً على القانون الإماراتي"
        )
        
        crew = Crew(
            agents=[agent],
            tasks=[task],
            verbose=2  # Enable detailed logging
        )
        
        result = crew.kickoff()
        return format_legal_response(result, 'ar' if is_arabic(query) else 'en')
        
    except Exception as e:
        st.error(f"""
        حدث خطأ أثناء معالجة طلبك. يرجى التحقق من التالي:
        1. التأكد من صحة مفتاح API الخاص بـ Groq
        2. التأكد من اتصال الإنترنت
        3. المحاولة مرة أخرى بعد بضع دقائق
        
        التفاصيل التقنية: {str(e)}
        """)
        return None

# Judge Tab
with tab1:
    st.header("استشارة القاضي الإماراتي")
    judge_query = st.text_area("اكتب سؤالك القانوني للقاضي:", key="judge_input", placeholder="أدخل النص هنا...")
    st.markdown(
        """
        <style>
        .element-container textarea {
            direction: rtl;
            text-align: right;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
    if st.button("الحصول على رأي القاضي", key="judge_button"):
        if judge_query:
            with st.spinner("القاضي يحلل قضيتك..."):
                judge_agent = create_judge_agent()
                response = get_agent_response(judge_agent, judge_query, selected_category)
                if response:
                    st.session_state.chat_history.append(("القاضي", judge_query, response))
                    st.write("رد القاضي:")
                    st.markdown(response, unsafe_allow_html=True)

# Advocate Tab
with tab2:
    st.header("استشارة المحامي الإماراتي")
    advocate_query = st.text_area("اكتب سؤالك القانوني للمحامي:", key="advocate_input", placeholder="أدخل النص هنا...")
    st.markdown(
        """
        <style>
        .element-container textarea {
            direction: rtl;
            text-align: right;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
    if st.button("الحصول على رأي المحامي", key="advocate_button"):
        if advocate_query:
            with st.spinner("المحامي يحلل قضيتك..."):
                advocate_agent = create_advocate_agent()
                response = get_agent_response(advocate_agent, advocate_query, selected_category)
                if response:
                    st.session_state.chat_history.append(("المحامي", advocate_query, response))
                    st.write("رد المحامي:")
                    st.markdown(response, unsafe_allow_html=True)

# Consultant Tab
with tab3:
    st.header("استشارة المستشار القضائي الإماراتي")
    consultant_query = st.text_area("اكتب سؤالك القانوني للمستشار:", key="consultant_input", placeholder="أدخل النص هنا...")
    st.markdown(
        """
        <style>
        .element-container textarea {
            direction: rtl;
            text-align: right;
        }
        </style>
        """,
        unsafe_allow_html=True
        )
    # In app.py, under the consultant tab (tab3)
    if st.button("الحصول على رأي المستشار", key="consultant_button"):
        if consultant_query:
            with st.spinner("المستشار يحلل قضيتك..."):
                consultant_agent = create_consultant_agent()
                response = get_agent_response(consultant_agent, consultant_query, selected_category)
                if response:
                    st.session_state.chat_history.append(("المستشار", consultant_query, response))
                    st.write("رد المستشار:")
                    st.markdown(response, unsafe_allow_html=True)
    
    # In app.py, under the consultant tab (tab3)