import streamlit as st
import json
from retrieval import retrieve_documents
from generation import generate_response
import time

# Page configuration
st.set_page_config(
    page_title="VN Law AI Assistant",
    page_icon="🤖",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# Custom CSS
st.markdown("""
<style>
    .main > div {
        padding-top: 2rem;
    }
    
    .stChatMessage {
        border-radius: 10px;
        padding: 1rem;
        margin: 0.5rem 0;
    }
    
    .stChatMessage[data-testid="chat-message-user"] {
        background-color: #007bff;
        color: white;
    }
    
    .stChatMessage[data-testid="chat-message-assistant"] {
        background-color: #f8f9fa;
        color: #333;
        border: 1px solid #dee2e6;
    }
    
    .chat-input {
        position: fixed;
        bottom: 0;
        left: 0;
        right: 0;
        background: white;
        padding: 1rem;
        border-top: 1px solid #dee2e6;
    }
    
    .status-container {
        background-color: #e7f3ff;
        padding: 0.5rem;
        border-radius: 5px;
        border-left: 4px solid #007bff;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
def init_session_state():
    """Khởi tạo session state"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "conversation_history" not in st.session_state:
        st.session_state.conversation_history = []
    
    if "session_id" not in st.session_state:
        st.session_state.session_id = f"session_{int(time.time())}"

def add_message_to_history(message):
    """Thêm message vào lịch sử hội thoại"""
    st.session_state.conversation_history.append(message)

def reset_conversation():
    """Reset cuộc hội thoại"""
    st.session_state.messages = []
    st.session_state.conversation_history = []
    st.session_state.session_id = f"session_{int(time.time())}"

def process_user_message(user_message):
    """Xử lý tin nhắn từ user - tương đương Flask endpoints"""
    
    # Hiển thị trạng thái đang xử lý
    with st.status("🔍 Đang xử lý câu hỏi...", expanded=True) as status:
        try:
            # Bước 1: Retrieve documents (thay thế /retrieve endpoint)
            status.write("📚 Đang tìm kiếm tài liệu liên quan...")
            documents = retrieve_documents(user_message)
            status.write(f"✅ Tìm thấy {len(documents) if isinstance(documents, list) else 1} tài liệu")
            
            # Bước 2: Generate response (thay thế /generate endpoint)
            status.write("🤖 Đang tạo phản hồi...")
            
            # Thêm message vào lịch sử trước khi generate
            add_message_to_history(user_message)
            
            response_data = generate_response(
                context=user_message,
                retrieved_data=documents,
                conversation_history=st.session_state.conversation_history
            )
            
            # Xử lý response data
            if isinstance(response_data, dict):
                response_text = response_data.get('response', str(response_data))
                # Hiển thị thêm thông tin nếu có
                if 'sources' in response_data:
                    status.write(f"📖 Nguồn tham khảo: {len(response_data['sources'])} tài liệu")
            else:
                response_text = str(response_data)
            
            status.write("✅ Hoàn thành!")
            status.update(label="✅ Đã xử lý xong", state="complete")
            
            return response_text, documents
            
        except Exception as e:
            error_msg = f"❌ Lỗi: {str(e)}"
            status.update(label="❌ Có lỗi xảy ra", state="error")
            return error_msg, []

def display_retrieved_documents(documents):
    """Hiển thị tài liệu được retrieve"""
    if documents and len(documents) > 0:
        with st.expander(f"📚 Tài liệu tham khảo ({len(documents)} tài liệu)", expanded=False):
            for i, doc in enumerate(documents[:3]):  # Chỉ hiển thị 3 tài liệu đầu
                if isinstance(doc, dict):
                    st.write(f"**Tài liệu {i+1}:**")
                    if 'title' in doc:
                        st.write(f"📄 **Tiêu đề:** {doc['title']}")
                    if 'content' in doc:
                        st.write(f"📝 **Nội dung:** {doc['content'][:200]}...")
                    if 'source' in doc:
                        st.write(f"🔗 **Nguồn:** {doc['source']}")
                else:
                    st.write(f"**Tài liệu {i+1}:** {str(doc)[:200]}...")
                st.divider()

def main():
    """Hàm chính của ứng dụng"""
    
    # Khởi tạo session state
    init_session_state()
    
    # Header
    st.title("🤖 VN Law AI Assistant")
    st.markdown("### Trợ lý AI tư vấn pháp luật Việt Nam")
    
    # Sidebar với thông tin và controls
    with st.sidebar:
        st.header("⚙️ Điều khiển")
        
        # Reset button
        if st.button("🗑️ Reset cuộc hội thoại", use_container_width=True):
            reset_conversation()
            st.rerun()
        
        # Session info
        st.markdown("---")
        st.markdown("📊 **Thông tin phiên:**")
        st.write(f"🆔 Session: {st.session_state.session_id}")
        st.write(f"💬 Số tin nhắn: {len(st.session_state.messages)}")
        st.write(f"📝 Lịch sử: {len(st.session_state.conversation_history)} mục")
        
        # Examples
        st.markdown("---")
        st.markdown("💡 **Câu hỏi mẫu:**")
        example_questions = [
            "Luật lao động quy định gì về thời gian làm việc?",
            "Thủ tục ly hôn theo pháp luật Việt Nam như thế nào?",
            "Quyền lợi của người tiêu dùng được bảo vệ ra sao?",
            "Điều kiện để thành lập doanh nghiệp là gì?"
        ]
        
        for question in example_questions:
            if st.button(question, key=f"example_{hash(question)}", use_container_width=True):
                st.session_state.example_question = question
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            
            # Hiển thị tài liệu tham khảo nếu có
            if message["role"] == "assistant" and "documents" in message:
                display_retrieved_documents(message["documents"])
    
    # Handle example question
    if hasattr(st.session_state, 'example_question'):
        user_message = st.session_state.example_question
        del st.session_state.example_question
        
        # Add user message to chat
        st.session_state.messages.append({"role": "user", "content": user_message})
        with st.chat_message("user"):
            st.markdown(user_message)
        
        # Process and add assistant response
        with st.chat_message("assistant"):
            response_text, documents = process_user_message(user_message)
            st.markdown(response_text)
            display_retrieved_documents(documents)
        
        # Add to messages
        st.session_state.messages.append({
            "role": "assistant", 
            "content": response_text,
            "documents": documents
        })
        
        st.rerun()
    
    # Chat input
    if prompt := st.chat_input("Nhập câu hỏi về pháp luật Việt Nam..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Generate and display assistant response
        with st.chat_message("assistant"):
            response_text, documents = process_user_message(prompt)
            st.markdown(response_text)
            display_retrieved_documents(documents)
        
        # Add assistant response to chat history
        st.session_state.messages.append({
            "role": "assistant", 
            "content": response_text,
            "documents": documents
        })
    
    # Welcome message
    if len(st.session_state.messages) == 0:
        with st.chat_message("assistant"):
            welcome_msg = """
            👋 Xin chào! Tôi là **VN Law AI Assistant**.
            
            🎯 Tôi có thể giúp bạn:
            - Tư vấn về pháp luật Việt Nam
            - Giải đáp các thắc mắc pháp lý
            - Tìm kiếm quy định liên quan
            
            💡 Hãy đặt câu hỏi hoặc chọn câu hỏi mẫu bên sidebar!
            """
            st.markdown(welcome_msg)
        
        st.session_state.messages.append({
            "role": "assistant", 
            "content": welcome_msg,
            "documents": []
        })

if __name__ == "__main__":
    main()