import gradio as gr
import time
import threading
from datetime import datetime
from typing import Dict, Any, List, Optional
import asyncio
import json
import os 
import sys 

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from main.chatbot import VietnameseLegalRAG
from utils.data_loader import LegalDataLoader
from config import Config

class VietnameseLegalChatbot:
    """
    Vietnamese Legal Chatbot implementing the architecture described in the documentation.
    
    Core Components:
    1. Natural Language Understanding (NLU) / NLP Module - handled by RAG system
    2. Dialogue Manager - manages conversation flow and context
    3. Knowledge Base / Data Store - vector store + BM25 + external search
    4. Response Generation Module - LLM with contextual information
    """
    
    def __init__(self):
        # Core chatbot backend components
        self.rag_system: Optional[VietnameseLegalRAG] = None
        self.initialization_status = {
            "status": "initializing", 
            "message": "🚀 Đang khởi tạo hệ thống trợ lý pháp lý...",
            "progress": 0,
            "details": "Chuẩn bị khởi động..."
        }
        
        # Dialogue manager - conversation state management
        self.conversation_sessions = {}
        self.current_session_id = "default"
        
        # Performance metrics
        self.metrics = {
            "total_queries": 0,
            "successful_queries": 0,
            "fallback_queries": 0,
            "average_response_time": 0,
            "start_time": time.time()
        }
        
        # Initialize system in background
        self._initialize_system_async()
    
    def _initialize_system_async(self):
        """Initialize the RAG system asynchronously with detailed progress"""
        def initialize():
            try:
                # Phase 1: Data Loading
                self.initialization_status.update({
                    "status": "loading_data", 
                    "message": "📚 Đang tải dữ liệu pháp luật Việt Nam...",
                    "progress": 10,
                    "details": "Đọc tệp dữ liệu từ thư mục..."
                })
                time.sleep(1)  # Visual feedback
                
                # Initialize data loader
                data_loader = LegalDataLoader()
                
                self.initialization_status.update({
                    "progress": 20,
                    "details": "Xử lý và chuẩn bị tài liệu..."
                })
                
                documents = data_loader.prepare_documents_for_indexing()
                
                if not documents:
                    self.initialization_status.update({
                        "status": "error", 
                        "message": "❌ Không thể tải dữ liệu pháp luật",
                        "progress": 0,
                        "details": "Kiểm tra thư mục data và tệp dữ liệu"
                    })
                    return
                
                # Phase 2: RAG System Initialization
                self.initialization_status.update({
                    "status": "initializing_rag", 
                    "message": "🤖 Đang khởi tạo hệ thống RAG...",
                    "progress": 30,
                    "details": f"Đã tải {len(documents):,} tài liệu pháp luật"
                })
                time.sleep(1)
                
                # Initialize RAG system
                self.rag_system = VietnameseLegalRAG()
                
                self.initialization_status.update({
                    "progress": 50,
                    "details": "Kết nối với LLM và vector store..."
                })
                
                # Phase 3: Index Building
                self.initialization_status.update({
                    "status": "building_indices", 
                    "message": "🔍 Đang xây dựng chỉ mục tìm kiếm...",
                    "progress": 60,
                    "details": "Kiểm tra chỉ mục hiện có..."
                })
                
                # Setup indices with progress updates
                try:
                    collection_info = self.rag_system.vector_store.get_collection_info()
                    bm25_loaded = self.rag_system.bm25_retriever.load_index()
                    
                    if not collection_info or not bm25_loaded:
                        self.initialization_status.update({
                            "message": "🏗️ Đang xây dựng chỉ mục lần đầu...",
                            "progress": 70,
                            "details": "Quá trình này có thể mất vài phút..."
                        })
                        self.rag_system.setup_indices(documents, force_rebuild=False)
                        
                        self.initialization_status.update({
                            "progress": 90,
                            "details": "Hoàn thiện thiết lập..."
                        })
                    else:
                        self.initialization_status.update({
                            "progress": 90,
                            "details": "Sử dụng chỉ mục có sẵn..."
                        })
                        
                except Exception as e:
                    self.initialization_status.update({
                        "message": "🏗️ Đang xây dựng chỉ mục mới...",
                        "progress": 70,
                        "details": f"Xây dựng lại do lỗi: {str(e)[:50]}..."
                    })
                    self.rag_system.setup_indices(documents, force_rebuild=False)
                
                # Phase 4: Final Validation
                self.initialization_status.update({
                    "progress": 95,
                    "details": "Kiểm tra tính toàn vẹn hệ thống..."
                })
                time.sleep(1)
                
                # Validate system
                system_status = self.rag_system.get_system_status()
                if not all([
                    system_status.get('llm_available'),
                    system_status.get('vector_store_available'),
                    system_status.get('bm25_available')
                ]):
                    raise Exception("Một số thành phần hệ thống không khả dụng")
                
                # Success
                self.initialization_status.update({
                    "status": "ready", 
                    "message": "✅ Trợ lý pháp lý đã sẵn sàng phục vụ!",
                    "progress": 100,
                    "details": f"Hệ thống hoạt động với {len(documents):,} tài liệu"
                })
                
            except Exception as e:
                self.initialization_status.update({
                    "status": "error", 
                    "message": f"❌ Lỗi khởi tạo: {str(e)}",
                    "progress": 0,
                    "details": "Kiểm tra cấu hình và thử lại"
                })
        
        # Start initialization in background thread
        init_thread = threading.Thread(target=initialize)
        init_thread.daemon = True
        init_thread.start()
    
    def get_system_status(self):
        """Get comprehensive system status"""
        base_status = {
            "initialization": self.initialization_status.copy(),
            "metrics": self.metrics.copy(),
            "uptime": time.time() - self.metrics["start_time"]
        }
        
        if self.initialization_status["status"] == "ready" and self.rag_system:
            # Get detailed system status
            rag_status = self.rag_system.get_system_status()
            base_status.update(rag_status)
            
            # Calculate success rate
            total = self.metrics["total_queries"]
            if total > 0:
                base_status["success_rate"] = (self.metrics["successful_queries"] / total) * 100
                base_status["fallback_rate"] = (self.metrics["fallback_queries"] / total) * 100
            else:
                base_status["success_rate"] = 0
                base_status["fallback_rate"] = 0
        
        return base_status
    
    def get_formatted_status(self):
        """Get formatted status for display"""
        status = self.get_system_status()
        init_status = status["initialization"]
        
        if init_status["status"] == "ready" and self.rag_system:
            # System ready - show comprehensive status
            indicators = []
            if status.get('llm_available'):
                indicators.append("🤖 LLM")
            if status.get('vector_store_available'):
                indicators.append("🔍 Vector Store")
            if status.get('bm25_available'):
                indicators.append("📊 BM25")
            if status.get('reranking_enabled') and status.get('reranker_available'):
                indicators.append("🎯 Reranker")
            
            uptime_hours = status["uptime"] / 3600
            
            return f"""
            <div class="status-ready">
                <h3>✅ Hệ thống hoạt động</h3>
                <p><strong>Thành phần:</strong> {', '.join(indicators)}</p>
                <p><strong>Thời gian hoạt động:</strong> {uptime_hours:.1f} giờ</p>
                <p><strong>Truy vấn:</strong> {status['metrics']['total_queries']} 
                   (Thành công: {status.get('success_rate', 0):.1f}%)</p>
            </div>
            """
        elif init_status["status"] == "error":
            return f"""
            <div class="status-error">
                <h3>❌ Lỗi hệ thống</h3>
                <p>{init_status['message']}</p>
                <p class="status-details">{init_status['details']}</p>
            </div>
            """
        else:
            # Loading state with progress
            return f"""
            <div class="status-loading">
                <h3>{init_status['message']}</h3>
                <div class="progress-container">
                    <div class="progress-bar" style="width: {init_status['progress']}%"></div>
                </div>
                <p class="progress-text">{init_status['progress']}% - {init_status['details']}</p>
            </div>
            """
    
    def create_new_session(self):
        """Create a new conversation session"""
        session_id = f"chat_{int(time.time() * 1000)}"
        self.conversation_sessions[session_id] = {
            "title": "Cuộc trò chuyện mới",
            "messages": [],
            "created_at": datetime.now(),
            "context": {},
            "metrics": {"queries": 0, "avg_response_time": 0}
        }
        self.current_session_id = session_id
        return session_id
    
    def update_session_title(self, session_id: str, first_message: str):
        """Update session title based on first user message"""
        if session_id in self.conversation_sessions:
            title = first_message[:50] + "..." if len(first_message) > 50 else first_message
            self.conversation_sessions[session_id]["title"] = title
    
    def process_message(self, message: str, history: List, session_id: str = None):
        """
        Core message processing with enhanced error handling and metrics
        """
        start_time = time.time()
        processing_status = "🤔 Đang suy nghĩ..."
        
        # Check if system is ready
        if self.initialization_status["status"] != "ready" or not self.rag_system:
            error_response = f"⚠️ {self.initialization_status['message']}"
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": error_response})
            return history, "", "❌ Hệ thống chưa sẵn sàng"
        
        if not message.strip():
            return history, "", "💬 Nhập câu hỏi để bắt đầu"
        
        try:
            # Update metrics
            self.metrics["total_queries"] += 1
            
            # Create session if it doesn't exist
            if session_id is None:
                session_id = self.create_new_session()
            elif session_id not in self.conversation_sessions:
                self.conversation_sessions[session_id] = {
                    "title": "Cuộc trò chuyện mới",
                    "messages": [],
                    "created_at": datetime.now(),
                    "context": {},
                    "metrics": {"queries": 0, "avg_response_time": 0}
                }
            
            # Update session title if this is the first message
            session = self.conversation_sessions[session_id]
            if len(session["messages"]) == 0:
                self.update_session_title(session_id, message.strip())
            
            # Store user message in session context
            session["messages"].append({"role": "user", "content": message, "timestamp": datetime.now()})
            session["metrics"]["queries"] += 1
            
            # Show processing status
            processing_status = "🔍 Đang tìm kiếm tài liệu..."
            
            # Process query through RAG system
            result = self.rag_system.answer_question(message.strip())
            
            processing_status = "✍️ Đang tạo câu trả lời..."
            
            # Format response with enhanced context
            response = result['answer']
            
            # Add metadata about information sources, search trigger, and question refinement (if enabled)
            source_info = ""
            refinement_info = ""
            
            # Add question refinement info if available and enabled
            if (Config.SHOW_REFINEMENT_INFO and 
                result.get('question_refinement') and 
                result.get('refined_question') != result.get('original_question')):
                refinement = result['question_refinement']
                if refinement.get('refinement_steps'):
                    refinement_summary = self.rag_system.question_refiner.get_refinement_summary(refinement)
                    refinement_info = f"\n\n*🔧 Câu hỏi đã được tối ưu: {refinement_summary}*"
            
            # Add search and source information (if enabled)
            if Config.SHOW_SEARCH_TRIGGER_INFO or Config.SHOW_SOURCE_INFO:
                if result.get('search_triggered'):
                    # Search was triggered due to insufficient information
                    if result.get('fallback_used') and result.get('search_results'):
                        if Config.SHOW_SEARCH_TRIGGER_INFO:
                            source_info = "\n\n*🔍➡️🌐 Không tìm thấy đủ thông tin trong tài liệu tham khảo nên đã tự động tìm kiếm trên web.*"
                        self.metrics["fallback_queries"] += 1
                    else:
                        if Config.SHOW_SEARCH_TRIGGER_INFO:
                            source_info = "\n\n*🔍 Đã kích hoạt tìm kiếm tự động.*"
                        self.metrics["fallback_queries"] += 1
                elif result.get('fallback_used'):
                    self.metrics["fallback_queries"] += 1
                    if result.get('search_results'):
                        if Config.SHOW_SOURCE_INFO:
                            source_info = "\n\n*🌐 Thông tin này được tìm kiếm từ web do không tìm thấy đủ thông tin trong cơ sở dữ liệu pháp luật nội bộ.*"
                else:
                    # Check if this was a rejected non-legal question
                    if result.get('rejected_non_legal'):
                        # Don't count as successful or fallback - it's a rejection
                        pass
                    else:
                        self.metrics["successful_queries"] += 1
                        if Config.SHOW_SOURCE_INFO:
                            source_info = f"\n\n*📚 Dựa trên {len(result.get('retrieved_documents', []))} tài liệu pháp luật.*"
            else:
                # Update metrics without showing info
                if result.get('rejected_non_legal'):
                    # Don't count rejected questions in success/fallback metrics
                    pass
                elif result.get('search_triggered') or result.get('fallback_used'):
                    self.metrics["fallback_queries"] += 1
                else:
                    self.metrics["successful_queries"] += 1
            
            response += refinement_info + source_info
            
            # Store assistant response in session context
            session["messages"].append({"role": "assistant", "content": response, "timestamp": datetime.now()})
            
            # Update conversation history for display
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": response})
            
            # Update metrics
            response_time = time.time() - start_time
            self.metrics["average_response_time"] = (
                (self.metrics["average_response_time"] * (self.metrics["total_queries"] - 1) + response_time) 
                / self.metrics["total_queries"]
            )
            session["metrics"]["avg_response_time"] = response_time
            
            # Format retrieved documents for display
            docs_info = self._format_retrieved_documents(result.get('retrieved_documents', []))
            
            # If this was a rejected non-legal question, show legal guidance instead
            if result.get('rejected_non_legal'):
                docs_info = self._format_legal_guidance()
            
            processing_status = f"✅ Hoàn thành ({response_time:.1f}s)"
            
            return history, docs_info, processing_status
            
        except Exception as e:
            error_response = f"❌ Lỗi xử lý: {str(e)}"
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": error_response})
            processing_status = f"❌ Lỗi: {str(e)[:50]}..."
            return history, "", processing_status
    
    def _format_retrieved_documents(self, documents):
        """Format retrieved documents with simplified styling"""
        if not documents:
            return "📄 **Không tìm thấy tài liệu tham khảo**"
        
        try:
            docs_html = f"## 📄 Tài liệu tham khảo ({len(documents)} tài liệu)\n\n"
            
            for i, doc in enumerate(documents, 1):
                # Safe access to document properties
                title = str(doc.get('title', 'Không có tiêu đề'))
                content = str(doc.get('content', ''))
                law_id = str(doc['metadata'].get('law_id', ''))
                
                # Truncate content for display
                display_content = content[:200] + "..." if len(content) > 200 else content
                
                docs_html += f"### 📄 {i}. ({law_id}) {title}\n"
                docs_html += f"**📝 Nội dung:** {display_content}\n\n"
                docs_html += "---\n\n"
            
            return docs_html
        except Exception as e:
            return f"📄 **Lỗi hiển thị tài liệu: {str(e)}**"
    
    def _format_legal_guidance(self):
        """Format legal guidance for rejected non-legal questions"""
        return """## 📚 Hướng dẫn sử dụng trợ lý pháp lý

Câu hỏi của bạn không thuộc lĩnh vực pháp luật mà tôi có thể hỗ trợ.

### ⚖️ Tôi có thể giúp bạn với:
- **Doanh nghiệp**: Thành lập, giải thể, vốn điều lệ, giấy phép kinh doanh
- **Lao động**: Hợp đồng lao động, lương, nghỉ phép, sa thải, bảo hiểm
- **Thuế**: Kê khai thuế, miễn thuế, thuế thu nhập cá nhân/doanh nghiệp
- **Bất động sản**: Mua bán nhà đất, chuyển nhượng, sổ đỏ, quyền sử dụng đất
- **Gia đình**: Hôn nhân, ly hôn, thừa kế, nuôi con, quyền con cái
- **Dân sự**: Hợp đồng, tranh chấp, bồi thường, quyền sở hữu
- **Hành chính**: Thủ tục pháp lý, giấy tờ, cơ quan nhà nước

### 💡 Gợi ý:
Hãy đặt câu hỏi cụ thể về các vấn đề pháp lý trên để nhận được hỗ trợ tốt nhất!"""
    
    def get_sample_questions(self):
        """Get categorized sample questions"""
        return {
            "🏢 Doanh nghiệp": [
                "Thủ tục thành lập doanh nghiệp như thế nào?",
                "Quy định về vốn điều lệ tối thiểu?",
                "Thủ tục giải thể doanh nghiệp?"
            ],
            "⚖️ Lao động": [
                "Quyền lợi của người lao động khi bị sa thải?",
                "Quy định về thời gian làm việc?",
                "Chế độ nghỉ phép hàng năm?"
            ],
            "💰 Thuế": [
                "Điều kiện miễn thuế thu nhập cá nhân?",
                "Cách tính thuế giá trị gia tăng?",
                "Thủ tục kê khai thuế?"
            ],
            "🏠 Bất động sản": [
                "Hợp đồng mua bán nhà đất cần giấy tờ gì?",
                "Quy trình chuyển nhượng quyền sử dụng đất?",
                "Quy trình chuyển nhượng quyền sử dụng đất khi chỉ một bên (chồng hay vợ) được quyền thừa kế?",
                "Thủ tục cấp sổ đỏ?"
            ],
            "👨‍👩‍👧‍👦 Gia đình": [
                "Thủ tục ly hôn thuận tình?",
                "Quyền thừa kế của con cái?",
                "Quy định về nuôi con nuôi?"
            ]
        }

def load_css():
    """Load CSS from external file"""
    try:
        with open('css/style.css', 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        try:
            with open('css/app/style.css', 'r', encoding='utf-8') as f:
                return f.read()
        except FileNotFoundError:
            print("⚠️ Warning: CSS file not found. Using default styles.")
            return ""

def create_chatbot_interface():
    """Create a simplified Gradio interface for Hugging Face Spaces"""
    
    # Initialize chatbot
    chatbot = VietnameseLegalChatbot()
    
    # Load CSS from external file
    css = load_css()
    
    with gr.Blocks(
        css=css, 
        title="Trợ lý Pháp lý Việt Nam", 
        theme=gr.themes.Default(),
        analytics_enabled=False
    ) as interface:
        
        # Enhanced header with simple styling
        gr.HTML("""
        <div class="main-header">
            <h1>⚖️ Trợ lý Pháp lý Việt Nam</h1>
            <p>Hệ thống tư vấn pháp luật thông minh</p>
        </div>
        """)
        
        with gr.Row(elem_classes="main-container"):
            # Left sidebar - Sample questions with dropdowns
            with gr.Column(scale=2, min_width=280):
                gr.HTML('<div class="sidebar-header">💡 Câu hỏi mẫu</div>')
                
                # Sample questions as simple buttons instead of dropdowns
                sample_categories = chatbot.get_sample_questions()
                sample_buttons = []
                
                for category, questions in sample_categories.items():
                    gr.HTML(f'<div style="margin: 10px 0; font-weight: bold; color: #4285f4;">{category}</div>')
                    for question in questions[:2]:  # Limit to 2 questions per category
                        btn = gr.Button(
                            question[:40] + "..." if len(question) > 40 else question,
                            size="sm",
                            variant="secondary",
                            elem_classes="sample-question-btn"
                        )
                        sample_buttons.append((btn, question))
            
            # Center - Main chat interface (expanded)
            with gr.Column(scale=5, min_width=500):
                # Simplified chat interface
                chatbot_component = gr.Chatbot(
                    label="💬 Trợ lý Pháp lý",
                    elem_classes="chat-container-main",
                    height=500,
                    show_copy_button=True
                )
                
                # Enhanced input area
                with gr.Row():
                    message_input = gr.Textbox(
                        placeholder="Hỏi tôi về pháp luật Việt Nam...",
                        container=False,
                        scale=5,
                        lines=1,
                        elem_classes="main-input"
                    )
                    send_button = gr.Button("📤 Gửi", variant="primary", scale=1, elem_classes="send-button")
                
                # Control buttons
                with gr.Row():
                    clear_chat_btn = gr.Button("🗑️ Xóa cuộc trò chuyện", size="sm", variant="secondary")
            
            # Right sidebar - Reference documents (expanded)
            with gr.Column(scale=3, min_width=350):
                # Enhanced documents display
                docs_display = gr.Markdown(
                    value="📄 **Tài liệu tham khảo sẽ hiển thị ở đây**",
                    label="📚 Cơ sở pháp lý",
                    elem_classes="docs-display"
                )
        
        # Simplified event handlers without State
        def handle_message(message, history):
            """Handle user message with simplified approach"""
            if not message.strip():
                return history, "", ""
            
            # Simple session ID
            session_id = "default"
            
            # Convert Gradio history to internal format
            internal_history = []
            for item in history:
                if isinstance(item, (list, tuple)) and len(item) == 2:
                    internal_history.append({"role": "user", "content": item[0]})
                    if item[1]:
                        internal_history.append({"role": "assistant", "content": item[1]})
            
            # Process message
            try:
                new_history, docs_info, process_status = chatbot.process_message(message, internal_history, session_id)
                
                # Convert back to Gradio format
                gradio_history = []
                i = 0
                while i < len(new_history):
                    if (i + 1 < len(new_history) and 
                        new_history[i].get("role") == "user" and 
                        new_history[i + 1].get("role") == "assistant"):
                        gradio_history.append([
                            new_history[i]["content"], 
                            new_history[i + 1]["content"]
                        ])
                        i += 2
                    elif new_history[i].get("role") == "user":
                        gradio_history.append([new_history[i]["content"], ""])
                        i += 1
                    else:
                        i += 1
                
                return gradio_history, "", docs_info
                
            except Exception as e:
                error_msg = f"❌ Lỗi xử lý: {str(e)}"
                history.append([message, error_msg])
                return history, "", ""
        
        def handle_clear_chat():
            """Clear chat"""
            return [], ""
        

        
        # Wire up sample question buttons
        for btn, question in sample_buttons:
            btn.click(
                lambda q=question: q,
                outputs=[message_input]
            )
        
        # Wire up main events
        send_button.click(
            handle_message,
            inputs=[message_input, chatbot_component],
            outputs=[chatbot_component, message_input, docs_display]
        )
        
        message_input.submit(
            handle_message,
            inputs=[message_input, chatbot_component],
            outputs=[chatbot_component, message_input, docs_display]
        )
        
        clear_chat_btn.click(
            handle_clear_chat,
            outputs=[chatbot_component, docs_display]
        )
    
    return interface

def main():
    """Enhanced main application entry point"""
    print("🚀 Khởi động Trợ lý Pháp lý Việt Nam với Gradio...")
    print("🎨 Sử dụng theme màu cờ Việt Nam và thiết kế hiện đại")
    print("⚡ Hệ thống động với cập nhật real-time")
    
    # Create and launch the enhanced interface
    interface = create_chatbot_interface()
    
    # Launch with minimal configuration to avoid schema issues
    try:
        interface.launch(
            share=True,
            show_api=False,
            debug=True
        )
    except Exception as e:
        print(f"Error launching interface: {e}")
        # Fallback launch with absolute minimal config
        interface.launch()

if __name__ == "__main__":
    main()